From 34d17a915a1bc45230ff24bc2f75ef3d593b5526 Mon Sep 17 00:00:00 2001 From: Dominik Demuth Date: Sat, 3 Aug 2024 19:04:13 +0200 Subject: [PATCH] proper saving of STE results; --- config.json | 16 ++--- main.py | 16 +---- rwsims/distributions.py | 12 +++- rwsims/functions.py | 10 +--- rwsims/motions.py | 57 +++++++++++------- rwsims/parameter.py | 46 +++++++++++---- rwsims/parsing.py | 4 +- rwsims/sims.py | 109 +++++++--------------------------- rwsims/spectrum.py | 39 +++++++++++++ rwsims/ste.py | 126 ++++++++++++++++++++++++++++++++++++++++ 10 files changed, 282 insertions(+), 153 deletions(-) diff --git a/config.json b/config.json index 925d7d5..717e702 100644 --- a/config.json +++ b/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,15 +35,15 @@ "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 } -} \ No newline at end of file +} diff --git a/main.py b/main.py index 8874982..9f8aa60 100644 --- a/main.py +++ b/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() diff --git a/rwsims/distributions.py b/rwsims/distributions.py index 87a83af..39ca23f 100644 --- a/rwsims/distributions.py +++ b/rwsims/distributions.py @@ -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) diff --git a/rwsims/functions.py b/rwsims/functions.py index 64d1982..c2e7403 100644 --- a/rwsims/functions.py +++ b/rwsims/functions.py @@ -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 diff --git a/rwsims/motions.py b/rwsims/motions.py index c18633a..6d6da80 100644 --- a/rwsims/motions.py +++ b/rwsims/motions.py @@ -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: diff --git a/rwsims/parameter.py b/rwsims/parameter.py index d4888fe..283ce5f 100644 --- a/rwsims/parameter.py +++ b/rwsims/parameter.py @@ -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 diff --git a/rwsims/parsing.py b/rwsims/parsing.py index 3121af4..f902947 100644 --- a/rwsims/parsing.py +++ b/rwsims/parsing.py @@ -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'} ) diff --git a/rwsims/sims.py b/rwsims/sims.py index 5e4e3a4..de4f7f5 100644 --- a/rwsims/sims.py +++ b/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() diff --git a/rwsims/spectrum.py b/rwsims/spectrum.py index e69de29..93c0691 100644 --- a/rwsims/spectrum.py +++ b/rwsims/spectrum.py @@ -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') diff --git a/rwsims/ste.py b/rwsims/ste.py index e69de29..14babf4 100644 --- a/rwsims/ste.py +++ b/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')