proper saving of STE results;

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,
+      "stop": 40e-6,
-      "steps": 99
+      "steps": 39
     },
     "t_mix": {
       "start": 1e-6,
       "stop": 1e0,
-      "steps": 19,
+      "steps": 31,
       "is_log": true
     },
     "t_echo": 0e-6

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()

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)
 

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

rwsims/motions.py

@@ -17,15 +17,32 @@ class BaseMotion(ABC):
     def __repr__(self):
         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
 
     @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]:
+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 += 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:

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:
+    def header(self) -> str:
-        return f'num_traj={self.num_walker}\nseed={self.seed}'
+        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:
+    def header(self) -> str:
-        return (f'dwell_time{self.dwell_time}\n'
+        return (
-                f'num_points={self.num_points}\n'
+            f'dwell_time = {self.dwell_time}\n'
-                f't_echo={self.t_echo}\n'
+            f'num_points = {self.num_points}\n'
-                f'lb={self.lb}\n'
+            f't_echo = {self.t_echo}\n'
-                f't_pulse={self.t_pulse}')
+            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

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'}
     )

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)
+    fits_cc = []
-    fig2, ax2 = plt.subplots(2)
+    fits_ss = []
-    fig3, ax3 = plt.subplots(2)
 
     # 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()
+            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()
 
 
 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()
+    # 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()
 
@@ -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()

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')

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')