proper saving of STE results;

This commit is contained in:
Dominik Demuth 2024-08-03 19:04:13 +02:00
parent 66e8925241
commit 34d17a915a
10 changed files with 282 additions and 153 deletions

View File

@ -1,6 +1,6 @@
{ {
"simulation": { "simulation": {
"num_walker": 30000, "num_walker": 10000,
"seed": null "seed": null
}, },
"molecule": { "molecule": {
@ -9,10 +9,12 @@
}, },
"correlation_times": { "correlation_times": {
"distribution": "DeltaDistribution", "distribution": "DeltaDistribution",
"tau": 1e-3 "tau": {
"list": [1e-4, 1e-3]
}
}, },
"motion": { "motion": {
"model": "TetrahedralJump" "model": "RandomJump"
}, },
"spectrum": { "spectrum": {
"dwell_time": 1e-6, "dwell_time": 1e-6,
@ -33,13 +35,13 @@
"stimulated_echo": { "stimulated_echo": {
"t_evo": { "t_evo": {
"start": 1e-6, "start": 1e-6,
"stop": 100e-6, "stop": 40e-6,
"steps": 99 "steps": 39
}, },
"t_mix": { "t_mix": {
"start": 1e-6, "start": 1e-6,
"stop": 1e0, "stop": 1e0,
"steps": 19, "steps": 31,
"is_log": true "is_log": true
}, },
"t_echo": 0e-6 "t_echo": 0e-6

16
main.py
View File

@ -1,22 +1,8 @@
from numpy.random import default_rng
import matplotlib.pyplot as plt import matplotlib.pyplot as plt
from rwsims.sims import run_ste_sim, run_spectrum_sim from rwsims.sims import run_ste_sim, run_spectrum_sim
from rwsims.motions import TetrahedralJump
# run_ste_sim('config.json') run_ste_sim('config.json')
# run_spectrum_sim('config.json') # run_spectrum_sim('config.json')
rng = default_rng()
tetra = TetrahedralJump(1, 0, rng)
for _ in range(100):
tetra.start()
omegas = tetra.jump(100)
plt.plot(omegas, '.')
break
plt.show()

View File

@ -3,7 +3,6 @@ from __future__ import annotations
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
import numpy as np import numpy as np
# from numpy.typing import ArrayLike
from numpy.random import Generator from numpy.random import Generator
@ -22,6 +21,9 @@ class BaseDistribution(ABC):
def __repr__(self): def __repr__(self):
pass pass
def header(self):
return f'tau = {self._tau}'
@abstractmethod @abstractmethod
def start(self): def start(self):
pass pass
@ -31,7 +33,7 @@ class BaseDistribution(ABC):
def mean_tau(self): def mean_tau(self):
pass pass
def wait(self, size: int = 1) -> ArrayLike: def wait(self, size: int = 1) -> 'ArrayLike':
return self._rng.exponential(self.tau_jump, size=size) return self._rng.exponential(self.tau_jump, size=size)
@ -57,6 +59,12 @@ class LogGaussianDistribution(BaseDistribution):
def __repr__(self): def __repr__(self):
return f'Log-Gaussian(tau={self._tau}, sigma={self._sigma})' return f'Log-Gaussian(tau={self._tau}, sigma={self._sigma})'
def header(self) -> str:
return (
f'tau = {self._tau}\n'
f'sigma = {self._sigma}'
)
def start(self): def start(self):
self.tau_jump = self._rng.lognormal(np.log(self._tau), self._sigma) self.tau_jump = self._rng.lognormal(np.log(self._tau), self._sigma)

View File

@ -1,17 +1,13 @@
from __future__ import annotations from __future__ import annotations
import numpy as np import numpy as np
# from numpy.typing import ArrayLike
from .distributions import BaseDistribution
from .motions import BaseMotion
def ste(x, a, f_infty, tau, beta): def ste(x: np.ndarray, a: float, f_infty: float, tau: float, beta: float) -> np.ndarray:
return a*((1-f_infty) * np.exp(-(x/tau)**beta) + f_infty) return a*((1-f_infty) * np.exp(-(x/tau)**beta) + f_infty)
def pulse_attn(freq, t_pulse: float): def pulse_attn(freq: np.ndarray, t_pulse: float):
# cf. Schmitt-Rohr/Spieß eq. 2.126; omega_1 * t_p = pi/2 # cf. Schmitt-Rohr/Spieß eq. 2.126; omega_1 * t_p = pi/2
pi_half_squared = np.pi**2 / 4 pi_half_squared = np.pi**2 / 4
omega = 2 * np.pi * freq omega = 2 * np.pi * freq

View File

@ -17,15 +17,32 @@ class BaseMotion(ABC):
def __repr__(self): def __repr__(self):
pass pass
@property @classmethod
def name(self) -> str: def name(cls) -> str:
return self.__class__.__name__ """
:return: Name of the actual class
"""
return cls.__class__.__name__
def start(self): def header(self) -> str:
return ''
def start(self) -> None:
"""
Function that should be called at the beginning of a trajectory.
"""
pass pass
@abstractmethod @abstractmethod
def jump(self, size: int = 1) -> 'ArrayLike': def jump(self, size: int = 1) -> 'ArrayLike':
"""
Array of omega_q for trajectory of length `size`.
Implementation is done in subclasses.
:param size: number of jumps that are processed in one go
:return: Array of omega_q of length `size`
"""
pass pass
@ -75,23 +92,15 @@ class TetrahedralJump(BaseMotion):
) )
orientations = np.zeros(4) orientations = np.zeros(4)
c = []
for i in range(4): for i in range(4):
corner_lab = rot @ corners[i] corner_lab = np.dot(rot, corners[i])
_, theta_i, phi_i = xyz_to_spherical(*corner_lab) _, theta_i, phi_i = xyz_to_spherical(*corner_lab)
c.append(corner_lab)
orientations[i] = omega_q(self._delta, self._eta, theta_i, phi_i) orientations[i] = omega_q(self._delta, self._eta, theta_i, phi_i)
# print(orientations)
#
# theta0 = np.arccos(cos_theta0)
# v0 = np.array([np.sin(theta0) * np.cos(phi0), np.sin(theta0)*np.sin(theta0), cos_theta0])
# norm = np.linalg.norm(v0)
# print(norm)
#
#
# corners = np.zeros((4, 3))
# corners[0] = v0
return orientations return orientations
def jump(self, size: int = 1) -> 'ArrayLike': def jump(self, size: int = 1) -> 'ArrayLike':
@ -106,10 +115,16 @@ class TetrahedralJump(BaseMotion):
# Helper functions # Helper functions
def xyz_to_spherical(x_in: float, y_in: float, z_in: float) -> tuple[np.floating, float, float]: def xyz_to_spherical(x_in: float, y_in: float, z_in: float) -> tuple[float, float, float]:
r = np.linalg.norm([x_in, y_in, z_in]) r: float = np.linalg.norm([x_in, y_in, z_in])
theta: float = np.arccos(z_in) theta: float = np.arccos(z_in)
theta += 2*np.pi
theta %= 2*np.pi
phi: float = np.arctan2(y_in, x_in) phi: float = np.arctan2(y_in, x_in)
phi += 2*np.pi
phi %= 2*np.pi
return r, theta, phi return r, theta, phi
@ -149,14 +164,14 @@ def get_rotation_matrix(vec_in: np.ndarray, vec_out: np.ndarray):
return rotation return rotation
def omega_q(delta: float, eta: float, cos_theta: ArrayLike, phi: ArrayLike) -> ArrayLike: def omega_q(delta: float, eta: float, cos_theta: 'ArrayLike', phi: 'ArrayLike') -> 'ArrayLike':
# sin_theta = np.sin(cos_theta) # sin_theta = np.sin(cos_theta)
# cos_theta = np.cos(cos_theta) # cos_theta = np.cos(cos_theta)
sin_theta_sq = 1 - cos_theta**2 sin_theta_sq = 1 - cos_theta**2
return np.pi * delta * (3 * cos_theta**2 - 1 + eta * sin_theta_sq * np.cos(2*phi)) return np.pi * delta * (3 * cos_theta**2 - 1 + eta * sin_theta_sq * np.cos(2*phi))
def draw_orientation(rng: Generator, size: int | None = None) -> tuple[ArrayLike, ArrayLike]: def draw_orientation(rng: Generator, size: int | None = None) -> tuple['ArrayLike', 'ArrayLike']:
if size is not None: if size is not None:
z_theta, z_phi = rng.random((2, size)) z_theta, z_phi = rng.random((2, size))
else: else:

View File

@ -6,12 +6,12 @@ from math import prod
from typing import Any from typing import Any
import numpy as np import numpy as np
# from numpy.typing import ArrayLike
from .functions import pulse_attn from .functions import pulse_attn
from .distributions import BaseDistribution from .distributions import BaseDistribution
from .motions import BaseMotion from .motions import BaseMotion
__all__ = [ __all__ = [
'SimParameter', 'SimParameter',
'MoleculeParameter', 'MoleculeParameter',
@ -20,6 +20,7 @@ __all__ = [
'DistParameter', 'DistParameter',
'MotionParameter', 'MotionParameter',
'Parameter', 'Parameter',
'make_filename'
] ]
@ -29,7 +30,7 @@ class SimParameter:
num_walker: int num_walker: int
t_max: float t_max: float
def totext(self) -> str: def header(self) -> str:
return f'num_traj = {self.num_walker}\nseed = {self.seed}' return f'num_traj = {self.num_walker}\nseed = {self.seed}'
@ -49,6 +50,13 @@ class StimEchoParameter:
def __post_init__(self): def __post_init__(self):
self.t_max = np.max(self.t_mix) + 2 * np.max(self.t_evo) + 2*self.t_echo self.t_max = np.max(self.t_mix) + 2 * np.max(self.t_evo) + 2*self.t_echo
def header(self) -> str:
return (
f't_evo = {self.t_evo}\n'
f't_mix = {self.t_mix}\n'
f't_echo={self.t_echo}\n'
)
@dataclass @dataclass
class SpectrumParameter: class SpectrumParameter:
@ -70,16 +78,19 @@ class SpectrumParameter:
self.freq = np.fft.fftshift(np.fft.fftfreq(self.num_points, self.dwell_time)) self.freq = np.fft.fftshift(np.fft.fftfreq(self.num_points, self.dwell_time))
self.pulse_attn = pulse_attn(self.freq, self.t_pulse) self.pulse_attn = pulse_attn(self.freq, self.t_pulse)
def totext(self) -> str: def header(self) -> str:
return (f'dwell_time{self.dwell_time}\n' return (
f'dwell_time = {self.dwell_time}\n'
f'num_points = {self.num_points}\n' f'num_points = {self.num_points}\n'
f't_echo = {self.t_echo}\n' f't_echo = {self.t_echo}\n'
f'lb = {self.lb}\n' f'lb = {self.lb}\n'
f't_pulse={self.t_pulse}') f't_pulse = {self.t_pulse}'
)
@dataclass @dataclass
class DistParameter: class DistParameter:
name: str
dist_type: BaseDistribution dist_type: BaseDistribution
variables: field(default_factory=dict) variables: field(default_factory=dict)
num_variables: int = 0 num_variables: int = 0
@ -103,6 +114,7 @@ class DistParameter:
@dataclass @dataclass
class MotionParameter: class MotionParameter:
name: str
model: BaseMotion model: BaseMotion
variables: field(default_factory=dict) variables: field(default_factory=dict)
num_variables: int = 0 num_variables: int = 0
@ -133,13 +145,23 @@ class Parameter:
motion: MotionParameter motion: MotionParameter
molecule: MoleculeParameter molecule: MoleculeParameter
def totext(self, sim: bool = True, spec: bool = True) -> str: def header(self, sim: bool = True, spec: bool = False, ste: bool = False) -> str:
text = [] text = []
if sim: if sim:
text.append(self.sim.totext()) text.append(self.sim.header())
if spec: if spec:
text.append(self.spec.totext()) text.append(self.spec.header())
if ste:
text.append(self.ste.header())
return '\n'.join(text) return '\n'.join(text)
def make_filename(dist: BaseDistribution, motion: BaseMotion) -> str:
filename = f'{dist}_{motion}'
filename = filename.replace(' ', '_')
filename = filename.replace('.', 'p')
return filename

View File

@ -53,7 +53,7 @@ def _parse_ste(params: dict[str, Any] | None) -> StimEchoParameter | None:
ste = StimEchoParameter( ste = StimEchoParameter(
t_mix=_make_times(params['t_mix']), t_mix=_make_times(params['t_mix']),
t_evo=_make_times(params['t_evo']), t_evo=_make_times(params['t_evo']),
t_echo=params['t_echo'] t_echo=params.get('t_echo', 0)
) )
return ste return ste
@ -79,6 +79,7 @@ def _parse_dist(params: dict[str, Any]) -> DistParameter:
'LogGaussian': LogGaussianDistribution 'LogGaussian': LogGaussianDistribution
} }
p = DistParameter( p = DistParameter(
name=params['distribution'],
dist_type=mapping[params['distribution']], dist_type=mapping[params['distribution']],
variables={k: _make_times(v) for k, v in params.items() if k != 'distribution'}, variables={k: _make_times(v) for k, v in params.items() if k != 'distribution'},
) )
@ -93,6 +94,7 @@ def _parse_motion(params: dict[str, Any]) -> MotionParameter:
} }
p = MotionParameter( p = MotionParameter(
name=params['model'],
model=mapping[params['model']], model=mapping[params['model']],
variables={k: _make_times(v) for k, v in params.items() if k != 'model'} variables={k: _make_times(v) for k, v in params.items() if k != 'model'}
) )

View File

@ -7,9 +7,9 @@ from numpy.random import Generator
from datetime import datetime from datetime import datetime
from scipy.interpolate import interp1d from scipy.interpolate import interp1d
import matplotlib.pyplot as plt import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
from .functions import ste from .spectrum import save_spectrum_data
from .ste import save_ste_data, fit_ste, save_ste_fit, plot_ste_fits
from .parameter import Parameter from .parameter import Parameter
from .distributions import BaseDistribution from .distributions import BaseDistribution
from .motions import BaseMotion from .motions import BaseMotion
@ -25,9 +25,8 @@ def run_ste_sim(config_file: str):
t_evo = p.ste.t_evo t_evo = p.ste.t_evo
t_echo = p.ste.t_echo t_echo = p.ste.t_echo
fig, ax = plt.subplots(2) fits_cc = []
fig2, ax2 = plt.subplots(2) fits_ss = []
fig3, ax3 = plt.subplots(2)
# outer loop over variables of distribution of correlation times # outer loop over variables of distribution of correlation times
for (i, dist_values) in enumerate(p.dist): for (i, dist_values) in enumerate(p.dist):
@ -54,7 +53,6 @@ def run_ste_sim(config_file: str):
dephased = phase(t_evo_k) dephased = phase(t_evo_k)
t0 = t_mix + t_evo_k t0 = t_mix + t_evo_k
rephased = phase(t0 + t_evo_k + 2*t_echo) + phase(t0) - 2 * phase(t0+t_echo) rephased = phase(t0 + t_evo_k + 2*t_echo) + phase(t0) - 2 * phase(t0+t_echo)
# print(t_evo_k, t0 + t_evo_k + 2*t_echo, t0)
cc[:, k] += np.cos(dephased)*np.cos(rephased) cc[:, k] += np.cos(dephased)*np.cos(rephased)
ss[:, k] += np.sin(dephased)*np.sin(rephased) ss[:, k] += np.sin(dephased)*np.sin(rephased)
@ -63,46 +61,21 @@ def run_ste_sim(config_file: str):
cc[:, 1:] /= num_traj cc[:, 1:] /= num_traj
ss[:, 1:] /= num_traj ss[:, 1:] /= num_traj
fig4, ax4 = plt.subplots() save_ste_data(cc, ss, p, dist, motion)
ax4.semilogx(t_mix, cc/cc[0, :], '.-')
fig5, ax5 = plt.subplots()
ax5.semilogx(t_mix, ss/ss[0, :], '.-')
for k in range(num_variables): p_fit_cc, p_fit_ss = fit_ste(cc, ss, t_evo, t_mix, dist_values, num_variables)
p0 = [0.5, 0, dist_values.get('tau', 1), 1] fits_cc.append(p_fit_cc)
fits_ss.append(p_fit_ss)
p_final = [] save_ste_fit(p_fit_cc, p_fit_ss, p, dist, motion)
p_final1 = []
for k, t_evo_k in enumerate(p.ste.t_evo):
try: plot_ste_fits(fits_cc, fits_ss, p.dist, p.motion)
res = curve_fit(ste, t_mix, cc[:, k], p0=p0, bounds=([0, 0, 0, 0], [np.inf, 1, np.inf, 1]))
p_final.append(res[0].tolist())
except RuntimeError:
p_final.append([np.nan, np.nan, np.nan, np.nan])
try:
res2 = curve_fit(ste, t_mix, ss[:, k], p0=p0, bounds=([0, 0, 0, 0], [np.inf, 1, np.inf, 1]))
p_final1.append(res2[0].tolist())
except RuntimeError:
p_final1.append([np.nan, np.nan, np.nan, np.nan])
p_final = np.array(p_final)
p_final1 = np.array(p_final1)
# ax[0].semilogy(p.ste.t_evo, p_final[:, 0], '.--')
# ax[1].semilogy(t_evo, p_final1[:, 0], '.--')
ax[0].plot(t_evo, p_final[:, 1], '.-')
ax[1].plot(t_evo, p_final1[:, 1], '.-')
ax2[0].semilogy(t_evo, p_final[:, 2], '.-')
ax2[1].semilogy(t_evo, p_final1[:, 2], '.-')
ax3[0].plot(t_evo, p_final[:, 3], '.-')
ax3[1].plot(t_evo, p_final1[:, 3], '.-')
plt.show() plt.show()
def run_spectrum_sim(config_file: str): def run_spectrum_sim(config_file: str):
p = parse(config_file) p: Parameter = parse(config_file)
rng, num_traj, t_max, delta, eta, num_variables = _prepare_sim(p) rng, num_traj, t_max, delta, eta, num_variables = _prepare_sim(p)
@ -115,8 +88,10 @@ def run_spectrum_sim(config_file: str):
for (i, dist_values) in enumerate(p.dist): for (i, dist_values) in enumerate(p.dist):
# noinspection PyCallingNonCallable # noinspection PyCallingNonCallable
dist = p.dist.dist_type(**dist_values, rng=rng) dist = p.dist.dist_type(**dist_values, rng=rng)
# second loop over parameter of motion model # second loop over parameter of motion model
for (j, motion_values) in enumerate(p.motion): for (j, motion_values) in enumerate(p.motion):
# noinspection PyCallingNonCallable # noinspection PyCallingNonCallable
motion = p.motion.model(delta, eta, **motion_values, rng=rng) motion = p.motion.model(delta, eta, **motion_values, rng=rng)
print(f'\nStart of {dist}, {motion}') print(f'\nStart of {dist}, {motion}')
@ -151,12 +126,16 @@ def run_spectrum_sim(config_file: str):
spec -= spec[0] spec -= spec[0]
spec *= p.spec.pulse_attn[:, None] spec *= p.spec.pulse_attn[:, None]
# save timesignals and spectra, also plots them
save_spectrum_data(timesignal, spec, p, dist, motion, t_echo_strings) save_spectrum_data(timesignal, spec, p, dist, motion, t_echo_strings)
fig2, ax2 = plt.subplots() # plot and save reduction factor
lines = ax2.semilogx(p.dist.variables['tau'], reduction_factor / num_traj, 'o--') reduction_factor /= num_traj
ax2.legend(lines, t_echo_strings) fig, ax = plt.subplots()
plt.savefig(f'{dist.name}_{motion.name}_reduction.png') lines = ax.semilogx(p.dist.variables['tau'], reduction_factor, 'o--')
ax.legend(lines, t_echo_strings)
plt.savefig(f'{dist.name()}_{motion.name()}_reduction.png')
plt.show() plt.show()
@ -175,11 +154,9 @@ def make_trajectory(
# number of jumps that are simulated at once # number of jumps that are simulated at once
chunks = min(int(0.51 * t_max / dist.tau_jump), 100_000) + 1 chunks = min(int(0.51 * t_max / dist.tau_jump), 100_000) + 1
# print(chunks)
t = [np.array([t_passed])] t = [np.array([t_passed])]
phase = [np.array([init_phase])] phase = [np.array([init_phase])]
# omega = [np.array([0])]
while t_passed < t_max: while t_passed < t_max:
# frequencies between jumps # frequencies between jumps
current_omega = motion.jump(size=chunks) current_omega = motion.jump(size=chunks)
@ -188,7 +165,6 @@ def make_trajectory(
accumulated_phase = np.cumsum(t_wait * current_omega) + phase[-1][-1] accumulated_phase = np.cumsum(t_wait * current_omega) + phase[-1][-1]
phase.append(accumulated_phase) phase.append(accumulated_phase)
# omega.append(current_omega)
t_wait = np.cumsum(t_wait) + t_passed t_wait = np.cumsum(t_wait) + t_passed
t_passed = t_wait[-1] t_passed = t_wait[-1]
@ -196,12 +172,6 @@ def make_trajectory(
t = np.concatenate(t) t = np.concatenate(t)
phase = np.concatenate(phase) phase = np.concatenate(phase)
# omega = np.concatenate(omega)
# fig_test, ax_test = plt.subplots()
# ax_test.plot(t, phase, 'x-')
# np.savetxt('trajectory.dat', np.c_[t, phase, omega])
# convenient interpolation to get phase at arbitrary times # convenient interpolation to get phase at arbitrary times
phase_interpol = interp1d(t, phase) phase_interpol = interp1d(t, phase)
@ -245,40 +215,3 @@ def print_step(n: int, num_traj: int, start: float, last_print: float) -> float:
return last_print return last_print
def make_filename(dist: BaseDistribution, motion: BaseMotion) -> str:
filename = f'{dist}_{motion}'
filename = filename.replace(' ', '_')
filename = filename.replace('.', 'p')
return filename
def save_spectrum_data(
timesignal: np.ndarray,
spectrum: np.ndarray,
param: Parameter,
dist: BaseDistribution,
motion: BaseMotion,
echo_strings: list[str]
):
filename = make_filename(dist, motion)
header = param.totext(sim=True, spec=True)
header += '\nx\t' + '\t'.join(echo_strings)
np.savetxt(filename + '_timesignal.dat', np.c_[param.spec.t_acq, timesignal], header=header)
np.savetxt(filename + '_spectrum.dat', np.c_[param.spec.freq, spectrum], header=header)
fig, ax = plt.subplots()
lines = ax.plot(param.spec.freq, spectrum)
ax.set_title(f'{dist}, {motion}')
ax.legend(lines, echo_strings)
plt.savefig(filename + '_spectrum.png')
fig1, ax1 = plt.subplots()
lines = ax1.plot(param.spec.t_acq, timesignal)
ax1.set_title(f'{dist}, {motion}')
ax1.legend(lines, echo_strings)
plt.savefig(filename + '_timesignal.png')
plt.show()

View File

@ -0,0 +1,39 @@
from __future__ import annotations
import numpy as np
from matplotlib import pyplot as plt
from .distributions import BaseDistribution
from .motions import BaseMotion
from .parameter import Parameter, make_filename
def save_spectrum_data(
timesignal: np.ndarray,
spectrum: np.ndarray,
param: Parameter,
dist: BaseDistribution,
motion: BaseMotion,
echo_strings: list[str]
) -> None:
filename = make_filename(dist, motion)
header = param.header(sim=True, spec=True)
header += '\n' + dist.header()
header += '\n' + motion.header()
header += '\nx\t' + '\t'.join(echo_strings)
np.savetxt(filename + '_timesignal.dat', np.c_[param.spec.t_acq, timesignal], header=header)
np.savetxt(filename + '_spectrum.dat', np.c_[param.spec.freq, spectrum], header=header)
fig, ax = plt.subplots()
lines = ax.plot(param.spec.freq, spectrum)
ax.set_title(f'{dist}, {motion}')
ax.legend(lines, echo_strings)
plt.savefig(filename + '_spectrum.png')
fig1, ax1 = plt.subplots()
lines = ax1.plot(param.spec.t_acq, timesignal)
ax1.set_title(f'{dist}, {motion}')
ax1.legend(lines, echo_strings)
plt.savefig(filename + '_timesignal.png')

View File

@ -0,0 +1,126 @@
from __future__ import annotations
import numpy as np
from matplotlib import pyplot as plt
from scipy.optimize import curve_fit
from .distributions import BaseDistribution
from .functions import ste
from .motions import BaseMotion
from .parameter import Parameter, make_filename
def save_ste_data(
cc: np.ndarray,
ss: np.ndarray,
param: Parameter,
dist: BaseDistribution,
motion: BaseMotion,
) -> None:
filename = make_filename(dist, motion)
header = param.header(sim=True, ste=True)
header += '\n' + dist.header()
header += '\n' + motion.header()
t_evo_string = list(map(lambda x: f'{x:.3e}', param.ste.t_evo))
header += '\nx\t' + '\t'.join(t_evo_string)
for ste_data, ste_label in ((cc, 'cc'), (ss, 'ss')):
np.savetxt(filename + f'_{ste_label}.dat', np.c_[param.ste.t_mix, ste_data], header=header)
fig, ax = plt.subplots()
lines = ax.semilogx(param.ste.t_mix, ste_data/ste_data[0, :])
ax.set_title(f'{dist}, {motion}')
ax.set_xlabel('t_mix / s')
ax.set_ylabel(f'F_{ste_label}(t) / F_{ste_label}(0)')
ax.legend(lines, t_evo_string)
plt.savefig(filename + f'_{ste_label}.png')
def fit_ste(
cc: np.ndarray,
ss: np.ndarray,
t_evo: np.ndarray,
t_mix: np.ndarray,
dist_values: dict,
num_variables: int
) -> tuple[np.ndarray, np.ndarray]:
for k in range(num_variables):
p_cc = []
p_ss = []
# fit ste decay for every evolution time
for k, t_evo_k in enumerate(t_evo):
for ste_data, ste_fits in ((cc, p_cc), (ss, p_ss)):
# [amplitude, f_infty, tau, beta]
p0 = [ste_data[0, k], 0.1, dist_values.get('tau', 1), 1]
try:
res = curve_fit(ste, t_mix, ste_data[:, k], p0=p0, bounds=([0, 0, 0, 0], [np.inf, 1, np.inf, 1]))
ste_fits.append([t_evo_k] + res[0].tolist())
except RuntimeError:
ste_fits.append([t_evo_k, np.nan, np.nan, np.nan, np.nan])
p_cc = np.array(p_cc)
p_ss = np.array(p_ss)
return p_cc, p_ss
def save_ste_fit(cc: np.ndarray, ss: np.ndarray, param: Parameter, dist: BaseDistribution, motion: BaseMotion):
filename = make_filename(dist, motion)
header = param.header(sim=True, ste=True)
header += '\n' + dist.header()
header += '\n' + motion.header()
header += '\nt_echo\tamp\tf_infty\ttau\tbeta'
np.savetxt(filename + '_cc_fit.dat', cc, header=header)
np.savetxt(filename + '_ss_fit.dat', ss, header=header)
def plot_ste_fits(fits_cc, fits_ss, dist, motion):
fits_cc = np.array(fits_cc)
fits_ss = np.array(fits_ss)
fig, ax = plt.subplots(2)
fig2, ax2 = plt.subplots(2)
fig3, ax3 = plt.subplots(2)
fig4, ax4 = plt.subplots(2)
num_motion = motion.num_variables
filename = f'{dist.name}_{motion.name}'
for (i, dist_values) in enumerate(dist):
for (j, motion_values) in enumerate(motion):
row = i*num_motion + j
label = ([f'{key}={val}' for key, val in dist_values.items()] +
[f'{key}={val}' for key, val in motion_values.items()])
for k, ax_k in enumerate((ax, ax2, ax3, ax4)):
ax_k[0].plot(fits_cc[row, :, 0], fits_cc[row, :, k+1], 'o--', label=', '.join(label))
ax_k[1].plot(fits_ss[row, :, 0], fits_ss[row, :, k+1], 'o--')
ax[0].legend()
ax[0].set_title('Amplitude (top: CC, bottom: SS)')
ax[0].set_yscale('log')
ax[1].set_yscale('log')
plt.savefig(filename + '_amp.png')
ax2[0].legend()
ax2[0].set_title('F_infty (top: CC, bottom: SS)')
ax2[0].set_yscale('log')
ax2[1].set_yscale('log')
plt.savefig(filename + '_finfty.png')
ax3[0].legend()
ax3[0].set_title('tau (top: CC, bottom: SS)')
ax3[0].set_yscale('log')
ax3[1].set_yscale('log')
plt.savefig(filename + '_tau.png')
ax4[0].legend()
ax4[0].set_title('beta (top: CC, bottom: SS)')
plt.savefig(filename + '_beta.png')