proper saving of STE results;
This commit is contained in:
Dominik Demuth
parent
66e8925241
commit
34d17a915a
14
config.json
14
config.json
@ -1,6 +1,6 @@
|
||||
{
|
||||
"simulation": {
|
||||
"num_walker": 30000,
|
||||
"num_walker": 10000,
|
||||
"seed": null
|
||||
},
|
||||
"molecule": {
|
||||
@ -9,10 +9,12 @@
|
||||
},
|
||||
"correlation_times": {
|
||||
"distribution": "DeltaDistribution",
|
||||
"tau": 1e-3
|
||||
"tau": {
|
||||
"list": [1e-4, 1e-3]
|
||||
}
|
||||
},
|
||||
"motion": {
|
||||
"model": "TetrahedralJump"
|
||||
"model": "RandomJump"
|
||||
},
|
||||
"spectrum": {
|
||||
"dwell_time": 1e-6,
|
||||
@ -33,13 +35,13 @@
|
||||
"stimulated_echo": {
|
||||
"t_evo": {
|
||||
"start": 1e-6,
|
||||
"stop": 100e-6,
|
||||
"steps": 99
|
||||
"stop": 40e-6,
|
||||
"steps": 39
|
||||
},
|
||||
"t_mix": {
|
||||
"start": 1e-6,
|
||||
"stop": 1e0,
|
||||
"steps": 19,
|
||||
"steps": 31,
|
||||
"is_log": true
|
||||
},
|
||||
"t_echo": 0e-6
|
||||
|
||||
16
main.py
16
main.py
@ -1,22 +1,8 @@
|
||||
from numpy.random import default_rng
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
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')
|
||||
|
||||
rng = default_rng()
|
||||
|
||||
tetra = TetrahedralJump(1, 0, rng)
|
||||
|
||||
for _ in range(100):
|
||||
tetra.start()
|
||||
omegas = tetra.jump(100)
|
||||
|
||||
plt.plot(omegas, '.')
|
||||
break
|
||||
|
||||
plt.show()
|
||||
|
||||
@ -3,7 +3,6 @@ from __future__ import annotations
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
import numpy as np
|
||||
# from numpy.typing import ArrayLike
|
||||
from numpy.random import Generator
|
||||
|
||||
|
||||
@ -22,6 +21,9 @@ class BaseDistribution(ABC):
|
||||
def __repr__(self):
|
||||
pass
|
||||
|
||||
def header(self):
|
||||
return f'tau = {self._tau}'
|
||||
|
||||
@abstractmethod
|
||||
def start(self):
|
||||
pass
|
||||
@ -31,7 +33,7 @@ class BaseDistribution(ABC):
|
||||
def mean_tau(self):
|
||||
pass
|
||||
|
||||
def wait(self, size: int = 1) -> ArrayLike:
|
||||
def wait(self, size: int = 1) -> 'ArrayLike':
|
||||
return self._rng.exponential(self.tau_jump, size=size)
|
||||
|
||||
|
||||
@ -57,6 +59,12 @@ class LogGaussianDistribution(BaseDistribution):
|
||||
def __repr__(self):
|
||||
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):
|
||||
self.tau_jump = self._rng.lognormal(np.log(self._tau), self._sigma)
|
||||
|
||||
|
||||
@ -1,17 +1,13 @@
|
||||
from __future__ import annotations
|
||||
|
||||
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)
|
||||
|
||||
|
||||
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
|
||||
pi_half_squared = np.pi**2 / 4
|
||||
omega = 2 * np.pi * freq
|
||||
@ -19,4 +15,4 @@ def pulse_attn(freq, t_pulse: float):
|
||||
numerator = np.sin(np.sqrt(pi_half_squared + omega**2 * t_pulse**2 / 2))
|
||||
denominator = np.sqrt(pi_half_squared + omega**2 * t_pulse**2 / 4)
|
||||
|
||||
return np.pi * numerator/denominator / 2
|
||||
return np.pi * numerator / denominator / 2
|
||||
|
||||
@ -17,15 +17,32 @@ class BaseMotion(ABC):
|
||||
def __repr__(self):
|
||||
pass
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return self.__class__.__name__
|
||||
@classmethod
|
||||
def name(cls) -> str:
|
||||
"""
|
||||
: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
|
||||
|
||||
@abstractmethod
|
||||
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
|
||||
|
||||
|
||||
@ -75,23 +92,15 @@ class TetrahedralJump(BaseMotion):
|
||||
)
|
||||
|
||||
orientations = np.zeros(4)
|
||||
c = []
|
||||
|
||||
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)
|
||||
c.append(corner_lab)
|
||||
|
||||
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
|
||||
|
||||
def jump(self, size: int = 1) -> 'ArrayLike':
|
||||
@ -106,10 +115,16 @@ class TetrahedralJump(BaseMotion):
|
||||
|
||||
# Helper functions
|
||||
|
||||
def xyz_to_spherical(x_in: float, y_in: float, z_in: float) -> tuple[np.floating, float, float]:
|
||||
r = np.linalg.norm([x_in, y_in, z_in])
|
||||
def xyz_to_spherical(x_in: float, y_in: float, z_in: float) -> tuple[float, float, float]:
|
||||
r: float = np.linalg.norm([x_in, y_in, 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 += 2*np.pi
|
||||
phi %= 2*np.pi
|
||||
|
||||
return r, theta, phi
|
||||
|
||||
@ -149,14 +164,14 @@ def get_rotation_matrix(vec_in: np.ndarray, vec_out: np.ndarray):
|
||||
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)
|
||||
# cos_theta = np.cos(cos_theta)
|
||||
sin_theta_sq = 1 - cos_theta**2
|
||||
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:
|
||||
z_theta, z_phi = rng.random((2, size))
|
||||
else:
|
||||
|
||||
@ -6,12 +6,12 @@ from math import prod
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
# from numpy.typing import ArrayLike
|
||||
|
||||
from .functions import pulse_attn
|
||||
from .distributions import BaseDistribution
|
||||
from .motions import BaseMotion
|
||||
|
||||
|
||||
__all__ = [
|
||||
'SimParameter',
|
||||
'MoleculeParameter',
|
||||
@ -20,6 +20,7 @@ __all__ = [
|
||||
'DistParameter',
|
||||
'MotionParameter',
|
||||
'Parameter',
|
||||
'make_filename'
|
||||
]
|
||||
|
||||
|
||||
@ -29,8 +30,8 @@ class SimParameter:
|
||||
num_walker: int
|
||||
t_max: float
|
||||
|
||||
def totext(self) -> str:
|
||||
return f'num_traj={self.num_walker}\nseed={self.seed}'
|
||||
def header(self) -> str:
|
||||
return f'num_traj = {self.num_walker}\nseed = {self.seed}'
|
||||
|
||||
|
||||
@dataclass
|
||||
@ -49,6 +50,13 @@ class StimEchoParameter:
|
||||
def __post_init__(self):
|
||||
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
|
||||
class SpectrumParameter:
|
||||
@ -70,16 +78,19 @@ class SpectrumParameter:
|
||||
self.freq = np.fft.fftshift(np.fft.fftfreq(self.num_points, self.dwell_time))
|
||||
self.pulse_attn = pulse_attn(self.freq, self.t_pulse)
|
||||
|
||||
def totext(self) -> str:
|
||||
return (f'dwell_time{self.dwell_time}\n'
|
||||
f'num_points={self.num_points}\n'
|
||||
f't_echo={self.t_echo}\n'
|
||||
f'lb={self.lb}\n'
|
||||
f't_pulse={self.t_pulse}')
|
||||
def header(self) -> str:
|
||||
return (
|
||||
f'dwell_time = {self.dwell_time}\n'
|
||||
f'num_points = {self.num_points}\n'
|
||||
f't_echo = {self.t_echo}\n'
|
||||
f'lb = {self.lb}\n'
|
||||
f't_pulse = {self.t_pulse}'
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class DistParameter:
|
||||
name: str
|
||||
dist_type: BaseDistribution
|
||||
variables: field(default_factory=dict)
|
||||
num_variables: int = 0
|
||||
@ -103,6 +114,7 @@ class DistParameter:
|
||||
|
||||
@dataclass
|
||||
class MotionParameter:
|
||||
name: str
|
||||
model: BaseMotion
|
||||
variables: field(default_factory=dict)
|
||||
num_variables: int = 0
|
||||
@ -133,13 +145,23 @@ class Parameter:
|
||||
motion: MotionParameter
|
||||
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 = []
|
||||
if sim:
|
||||
text.append(self.sim.totext())
|
||||
text.append(self.sim.header())
|
||||
|
||||
if spec:
|
||||
text.append(self.spec.totext())
|
||||
text.append(self.spec.header())
|
||||
|
||||
if ste:
|
||||
text.append(self.ste.header())
|
||||
|
||||
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
|
||||
|
||||
@ -53,7 +53,7 @@ def _parse_ste(params: dict[str, Any] | None) -> StimEchoParameter | None:
|
||||
ste = StimEchoParameter(
|
||||
t_mix=_make_times(params['t_mix']),
|
||||
t_evo=_make_times(params['t_evo']),
|
||||
t_echo=params['t_echo']
|
||||
t_echo=params.get('t_echo', 0)
|
||||
)
|
||||
return ste
|
||||
|
||||
@ -79,6 +79,7 @@ def _parse_dist(params: dict[str, Any]) -> DistParameter:
|
||||
'LogGaussian': LogGaussianDistribution
|
||||
}
|
||||
p = DistParameter(
|
||||
name=params['distribution'],
|
||||
dist_type=mapping[params['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(
|
||||
name=params['model'],
|
||||
model=mapping[params['model']],
|
||||
variables={k: _make_times(v) for k, v in params.items() if k != 'model'}
|
||||
)
|
||||
|
||||
109
rwsims/sims.py
109
rwsims/sims.py
@ -7,9 +7,9 @@ from numpy.random import Generator
|
||||
from datetime import datetime
|
||||
from scipy.interpolate import interp1d
|
||||
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 .distributions import BaseDistribution
|
||||
from .motions import BaseMotion
|
||||
@ -25,9 +25,8 @@ def run_ste_sim(config_file: str):
|
||||
t_evo = p.ste.t_evo
|
||||
t_echo = p.ste.t_echo
|
||||
|
||||
fig, ax = plt.subplots(2)
|
||||
fig2, ax2 = plt.subplots(2)
|
||||
fig3, ax3 = plt.subplots(2)
|
||||
fits_cc = []
|
||||
fits_ss = []
|
||||
|
||||
# outer loop over variables of distribution of correlation times
|
||||
for (i, dist_values) in enumerate(p.dist):
|
||||
@ -54,7 +53,6 @@ def run_ste_sim(config_file: str):
|
||||
dephased = phase(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)
|
||||
# print(t_evo_k, t0 + t_evo_k + 2*t_echo, t0)
|
||||
cc[:, k] += np.cos(dephased)*np.cos(rephased)
|
||||
ss[:, k] += np.sin(dephased)*np.sin(rephased)
|
||||
|
||||
@ -63,46 +61,21 @@ def run_ste_sim(config_file: str):
|
||||
cc[:, 1:] /= num_traj
|
||||
ss[:, 1:] /= num_traj
|
||||
|
||||
fig4, ax4 = plt.subplots()
|
||||
ax4.semilogx(t_mix, cc/cc[0, :], '.-')
|
||||
fig5, ax5 = plt.subplots()
|
||||
ax5.semilogx(t_mix, ss/ss[0, :], '.-')
|
||||
save_ste_data(cc, ss, p, dist, motion)
|
||||
|
||||
for k in range(num_variables):
|
||||
p0 = [0.5, 0, dist_values.get('tau', 1), 1]
|
||||
p_fit_cc, p_fit_ss = fit_ste(cc, ss, t_evo, t_mix, dist_values, num_variables)
|
||||
fits_cc.append(p_fit_cc)
|
||||
fits_ss.append(p_fit_ss)
|
||||
|
||||
p_final = []
|
||||
p_final1 = []
|
||||
for k, t_evo_k in enumerate(p.ste.t_evo):
|
||||
save_ste_fit(p_fit_cc, p_fit_ss, p, dist, motion)
|
||||
|
||||
try:
|
||||
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], '.-')
|
||||
plot_ste_fits(fits_cc, fits_ss, p.dist, p.motion)
|
||||
|
||||
plt.show()
|
||||
|
||||
|
||||
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)
|
||||
|
||||
@ -115,8 +88,10 @@ def run_spectrum_sim(config_file: str):
|
||||
for (i, dist_values) in enumerate(p.dist):
|
||||
# noinspection PyCallingNonCallable
|
||||
dist = p.dist.dist_type(**dist_values, rng=rng)
|
||||
|
||||
# second loop over parameter of motion model
|
||||
for (j, motion_values) in enumerate(p.motion):
|
||||
|
||||
# noinspection PyCallingNonCallable
|
||||
motion = p.motion.model(delta, eta, **motion_values, rng=rng)
|
||||
print(f'\nStart of {dist}, {motion}')
|
||||
@ -151,12 +126,16 @@ def run_spectrum_sim(config_file: str):
|
||||
spec -= spec[0]
|
||||
spec *= p.spec.pulse_attn[:, None]
|
||||
|
||||
# save timesignals and spectra, also plots them
|
||||
save_spectrum_data(timesignal, spec, p, dist, motion, t_echo_strings)
|
||||
|
||||
fig2, ax2 = plt.subplots()
|
||||
lines = ax2.semilogx(p.dist.variables['tau'], reduction_factor / num_traj, 'o--')
|
||||
ax2.legend(lines, t_echo_strings)
|
||||
plt.savefig(f'{dist.name}_{motion.name}_reduction.png')
|
||||
# plot and save reduction factor
|
||||
reduction_factor /= num_traj
|
||||
fig, ax = plt.subplots()
|
||||
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()
|
||||
|
||||
@ -175,11 +154,9 @@ def make_trajectory(
|
||||
|
||||
# number of jumps that are simulated at once
|
||||
chunks = min(int(0.51 * t_max / dist.tau_jump), 100_000) + 1
|
||||
# print(chunks)
|
||||
|
||||
t = [np.array([t_passed])]
|
||||
phase = [np.array([init_phase])]
|
||||
# omega = [np.array([0])]
|
||||
while t_passed < t_max:
|
||||
# frequencies between jumps
|
||||
current_omega = motion.jump(size=chunks)
|
||||
@ -188,7 +165,6 @@ def make_trajectory(
|
||||
|
||||
accumulated_phase = np.cumsum(t_wait * current_omega) + phase[-1][-1]
|
||||
phase.append(accumulated_phase)
|
||||
# omega.append(current_omega)
|
||||
|
||||
t_wait = np.cumsum(t_wait) + t_passed
|
||||
t_passed = t_wait[-1]
|
||||
@ -196,12 +172,6 @@ def make_trajectory(
|
||||
|
||||
t = np.concatenate(t)
|
||||
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
|
||||
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
|
||||
|
||||
|
||||
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()
|
||||
|
||||
@ -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')
|
||||
126
rwsims/ste.py
126
rwsims/ste.py
@ -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')
|
||||
Loading…
Reference in New Issue
Block a user