moving around

pyproject.toml

@@ -25,8 +25,8 @@ classifiers = [
 ]
 
 [project.scripts]
-rw_spectra = "scripts.sim_spectra"
+rw_spectra = "sims.run_spectrum_sim"
-rw_ste = "scripts.sim_ste"
+rw_ste = "sims.run_ste_sim"
 
 [project.urls]
 Homepage = "https://github.com/pypa/sampleproject"

src/config.json

@@ -1,6 +1,6 @@
 {
   "simulation": {
-    "num_walker": 2500,
+    "num_walker": 5000,
     "seed": null
   },
   "molecule": {
@@ -8,32 +8,40 @@
     "eta": 0.0
   },
   "correlation_times": {
-    "distribution": "DeltaDistribution",
+    "distribution": "LogGaussian",
     "tau": {
       "start": 1e-4,
-      "stop": 1e-2,
+      "stop": 1e-8,
-      "steps": 6,
+      "steps": 9,
       "is_log": true
+    },
+    "sigma": {
+      "list": [0.5, 1, 2]
     }
   },
   "motion": {
-    "model": "RandomJump"
+    "model": "TetrahedralJump"
   },
   "spectrum": {
     "dwell_time": 1e-6,
     "num_points": 4096,
     "t_echo": {
-      "list": [0, 5e-6, 10e-6, 20e-6, 50e-6, 100e-6, 200e-6]
+      "list": [5e-6, 10e-6, 20e-6, 40e-6, 60e-6, 100e-6]
     },
     "line_broadening": 2e3
   },
   "stimulated_echo": {
-    "t_evo": 10e-6,
+    "t_evo": {
+      "start": 2e-6,
+      "stop": 100e-6,
+      "steps": 98
+    },
     "t_mix": {
-      "start": 1e-5,
+      "start": 1e-7,
-      "stop": 1e-2,
+      "stop": 1e-1,
-      "steps": 10,
+      "steps": 31,
       "is_log": true
-    }
+    },
+    "t_echo": 15e-6
   }
 }

src/main.py

@@ -0,0 +1,5 @@
+from rwsims.sims import run_ste_sim, run_spectrum_sim
+
+# run_ste_sim('config.json')
+
+run_spectrum_sim('config.json')

src/rwsims/distributions.py

@@ -1,16 +1,61 @@
 from __future__ import annotations
 
+from abc import ABC, abstractmethod
+
+import numpy as np
 from numpy.typing import ArrayLike
 from numpy.random import Generator
 
 
-class DeltaDistribution:
+class BaseDistribution(ABC):
     def __init__(self, tau: float, rng: Generator | None = None):
         self._tau = tau
         self._rng = rng
 
+        self._tau_jump = tau
+
+    @abstractmethod
     def __repr__(self):
-        return f'DeltaDistribution (tau={self._tau})'
+        pass
+
+    @abstractmethod
+    def start(self):
+        pass
+
+    @property
+    @abstractmethod
+    def mean_tau(self):
+        pass
 
     def wait(self, size: int = 1) -> ArrayLike:
-        return self._rng.exponential(self._tau, size=size)
+        return self._rng.exponential(self._tau_jump, size=size)
+
+
+class DeltaDistribution(BaseDistribution):
+
+    def __repr__(self):
+        return f'No distribution(tau={self._tau})'
+
+    def start(self):
+        self._tau_jump = self._tau
+
+    @property
+    def mean_tau(self):
+        return self._tau
+
+
+class LogGaussianDistribution(BaseDistribution):
+    def __init__(self, tau: float, sigma: float, rng: Generator):
+        super().__init__(tau=tau, rng=rng)
+
+        self._sigma = sigma
+
+    def __repr__(self):
+        return f'Log-Gaussian(tau={self._tau}, sigma={self._sigma})'
+
+    def start(self):
+        self._tau_jump = self._rng.lognormal(np.log(self._tau), self._sigma)
+
+    @property
+    def mean_tau(self):
+        return self._tau * np.exp(self._sigma**2 / 2)

src/rwsims/helper.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+import numpy as np
+from numpy.typing import ArrayLike
+from numpy.random import Generator
+
+
+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])
+    theta = np.arccos(z_in)
+    phi = np.arctan2(y_in, x_in)
+
+    return r, theta, phi
+
+
+def spherical_to_xyz(r: float, theta: float, phi: float) -> tuple[float, float, float]:
+    sin_theta = np.sin(theta)
+    return r*np.cos(phi)*sin_theta, r*np.sin(phi)*sin_theta, r*np.cos(theta)
+
+
+def get_rotation_matrix(vec_in: np.ndarray, vec_out: np.ndarray):
+    rotation = np.eye(3)
+
+    # rotation by angle around given axis
+    cos_angle = np.dot(vec_in, vec_out)
+
+    # check for parallel / anti-parallel vectors
+    if cos_angle == 1:
+        return rotation
+    elif cos_angle == -1:
+        return -rotation
+
+    else:
+        axis = np.cross(vec_in, vec_out)
+        scale = np.linalg.norm(axis)
+        axis /= scale
+        sin_angle = scale / np.linalg.norm(vec_in) / np.linalg.norm(vec_out)
+
+        v_cross = np.array([
+            [0,       -axis[2], axis[1]],
+            [axis[2], 0,        -axis[0]],
+            [-axis[1], axis[0], 0],
+        ])
+
+        rotation += sin_angle * v_cross
+        rotation += (1-cos_angle) * v_cross @ v_cross
+
+        return rotation
+
+
+def omega_q(delta: float, eta: float, theta: ArrayLike, phi: ArrayLike) -> ArrayLike:
+    cos_theta = np.cos(theta)
+    sin_theta = np.sin(theta)
+    return np.pi * delta * (3 * cos_theta**2 - 1 + eta * sin_theta**2 * np.cos(2*phi))
+
+
+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:
+        z_theta, z_phi = rng.random(2)
+    theta = np.arccos(1 - 2 * z_theta)
+    phi = 2 * np.pi * z_phi
+
+    return theta, phi

src/rwsims/motions.py

@@ -1,33 +1,92 @@
 from __future__ import annotations
 
+from abc import ABC, abstractmethod
+
 import numpy as np
 from numpy.random import Generator
 from numpy.typing import ArrayLike
 
-
+from .helper import xyz_to_spherical, spherical_to_xyz, omega_q, draw_orientation, get_rotation_matrix
-def omega_q(delta: float, eta: float, theta: float, phi: float) -> ArrayLike:
-    cos_theta = np.cos(theta)
-    sin_theta = np.sin(theta)
-    return 2 * np.pi * delta * (3 * cos_theta * cos_theta - 1 + eta * sin_theta*sin_theta * np.cos(2*phi))
 
 
-def draw_orientation(delta: float, eta: float, rng: Generator, size: int = 1) -> ArrayLike:
+class BaseMotion(ABC):
-    z_theta, z_phi = rng.random((2, size))
+    def __init__(self, delta: float, eta: float, rng: Generator):
-    theta = np.arccos(1 - 2 * z_theta)
-    phi = 2 * np.pi * z_phi
-
-    return omega_q(delta, eta, theta, phi)
-
-
-class RandomJump:
-    def __init__(self, delta: float, eta: float, rng: Generator | None = None):
         self._delta = delta
         self._eta = eta
 
         self._rng = rng
 
+    @abstractmethod
+    def __repr__(self):
+        pass
+
+    def start(self):
+        pass
+
+    @abstractmethod
+    def jump(self, size: int = 1) -> ArrayLike:
+        pass
+
+
+class RandomJump(BaseMotion):
+
     def __repr__(self):
         return 'Random Jump'
 
     def jump(self, size: int = 1) -> ArrayLike:
-        return draw_orientation(self._delta, self._eta, self._rng, size=size)
+        return omega_q(self._delta, self._eta, *draw_orientation(self._rng, size=size))
+
+
+class TetrahedralJump(BaseMotion):
+    def __init__(self, delta: float, eta: float, rng: Generator):
+        super().__init__(delta, eta, rng)
+
+        self._orientation = None
+        self._start = None
+
+    def __repr__(self):
+        return 'Tetrahedral Jump'
+
+    def start(self):
+        self._orientation = self._make_tetrahedron()
+        self._start = self._rng.choice([0, 1, 2, 3])
+
+    def _make_tetrahedron(self) -> np.ndarray:
+        beta = np.arccos(-1/3)  # tetrahedral angle 109.5 degrees
+        sin_beta = np.sin(beta)
+        cos_beta = np.cos(beta)
+
+        # corners of a tetrahedron
+        alpha = 2 * np.pi * self._rng.random()
+        corners = np.array([
+            [0, 0, 1],
+            [sin_beta * np.cos(alpha),           sin_beta * np.sin(alpha),           cos_beta],
+            [sin_beta * np.cos(alpha+2*np.pi/3), sin_beta * np.sin(alpha+2*np.pi/3), cos_beta],
+            [sin_beta * np.cos(alpha+4*np.pi/3), sin_beta * np.sin(alpha+4*np.pi/3), cos_beta]
+        ])
+
+        # orientation in lab system
+        theta0, phi0 = draw_orientation(self._rng)
+
+        rot = get_rotation_matrix(
+            corners[0],
+            np.array(spherical_to_xyz(1., theta0, phi0)),
+        )
+        orientations = np.zeros(4)
+        for i in range(4):
+            corner_lab = np.dot(rot, corners[i])
+            _, theta_i, phi_i = xyz_to_spherical(*corner_lab)
+            orientations[i] = omega_q(self._delta, self._eta, theta_i, phi_i)
+
+        return orientations
+
+    def jump(self, size: int = 1) -> ArrayLike:
+        jumps = self._rng.choice([1, 2, 3], size=size)
+        jumps = np.cumsum(jumps) + self._start
+        jumps %= 4
+
+        self._start = jumps[-1]
+
+        return self._orientation[jumps]
+
+

src/rwsims/parameter.py

@@ -2,15 +2,23 @@ from __future__ import annotations
 
 from dataclasses import dataclass, field
 from itertools import product
+from math import prod
 from typing import Any
 
 import numpy as np
 
-from src.rwsims.distributions import DeltaDistribution
+from .distributions import BaseDistribution
-from src.rwsims.motions import RandomJump
+from .motions import BaseMotion
 
-
+__all__ = [
-__all__ = ['SimParameter', 'MoleculeParameter', 'StimEchoParameter', 'SpectrumParameter', 'DistParameter', 'MotionParameter', 'Parameter']
+    'SimParameter',
+    'MoleculeParameter',
+    'StimEchoParameter',
+    'SpectrumParameter',
+    'DistParameter',
+    'MotionParameter',
+    'Parameter',
+]
 
 
 @dataclass
@@ -30,10 +38,11 @@ class MoleculeParameter:
 class StimEchoParameter:
     t_evo: np.ndarray
     t_mix: np.ndarray
+    t_echo: float
     t_max: float = field(init=False)
 
     def __post_init__(self):
-        self.t_max = np.max(self.t_mix) + 2 * np.max(self.t_evo)
+        self.t_max = np.max(self.t_mix) + 2 * np.max(self.t_evo) + 2*self.t_echo
 
 
 @dataclass
@@ -54,13 +63,13 @@ class SpectrumParameter:
 
 @dataclass
 class DistParameter:
-    dist_type: DeltaDistribution
+    dist_type: BaseDistribution
     variables: field(default_factory=dict)
     num_variables: int = 0
     iter: field(init=False) = None
 
     def __post_init__(self):
-        self.num_variables = sum(map(len, self.variables.values()))
+        self.num_variables = prod(map(len, self.variables.values()))
 
     def __iter__(self):
         return self
@@ -77,13 +86,13 @@ class DistParameter:
 
 @dataclass
 class MotionParameter:
-    model: RandomJump
+    model: BaseMotion
     variables: field(default_factory=dict)
     num_variables: int = 0
     iter: field(init=False) = None
 
     def __post_init__(self):
-        self.num_variables = sum(map(len, self.variables.values()))
+        self.num_variables = prod(map(len, self.variables.values()))
 
     def __iter__(self):
         return self

src/rwsims/parsing.py

@@ -5,11 +5,9 @@ from typing import Any
 
 import numpy as np
 
-from distributions import (
+from .distributions import DeltaDistribution, LogGaussianDistribution
-    DeltaDistribution
+from .motions import RandomJump, TetrahedralJump
-)
+from .parameter import *
-from motions import RandomJump
-from parameter import *
 
 
 def parse(config_file: str) -> Parameter:
@@ -55,6 +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']
     )
     return ste
 
@@ -75,7 +74,8 @@ def _parse_spectrum(params: dict[str, Any] | None) -> SpectrumParameter | None:
 
 def _parse_dist(params: dict[str, Any]) -> DistParameter:
     mapping: dict = {
-        'DeltaDistribution': DeltaDistribution
+        'DeltaDistribution': DeltaDistribution,
+        'LogGaussian': LogGaussianDistribution
     }
     p = DistParameter(
         dist_type=mapping[params['distribution']],
@@ -88,6 +88,7 @@ def _parse_dist(params: dict[str, Any]) -> DistParameter:
 def _parse_motion(params: dict[str, Any]) -> MotionParameter:
     mapping = {
         'RandomJump':  RandomJump,
+        'TetrahedralJump': TetrahedralJump,
     }
 
     p = MotionParameter(

src/rwsims/sims.py

@@ -5,21 +5,103 @@ from time import time
 import numpy as np
 from numpy.random import Generator
 from scipy.interpolate import interp1d
-from scipy.optimize import curve_fit
 import matplotlib.pyplot as plt
+from scipy.optimize import curve_fit
 
-from parameter import Parameter
+from .parameter import Parameter
-from parser import parse
+from .distributions import BaseDistribution
+from .motions import BaseMotion
+from .parsing import parse
 
 
-def ste(x, a, f_infty, tau, beta):
+def run_ste_sim(config_file: str):
-    return a*((1-f_infty) * np.exp(-(x/tau)**beta) + f_infty)
+    p = parse(config_file)
+
+    rng, num_traj, t_max, delta, eta, num_variables = _prepare_sim(p)
+
+    t_mix = p.ste.t_mix
+    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)
+
+    # outer loop over variables of distribution of correlation times
+    for (i, dist_values) in enumerate(p.dist):
+        # noinspection PyCallingNonCallable
+        dist = p.dist.dist_type(**dist_values, rng=rng)
+        chunks = int(0.6 * t_max / dist_values.get('tau', 1)) + 1
+
+        # second loop over parameter of motional 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} and {motion}')
+
+            start = time()
+
+            cc = np.zeros((len(t_mix), len(t_evo)))
+            ss = np.zeros((len(t_mix), len(t_evo)))
+
+            # inner loop to create trajectories
+            for n in range(num_traj):
+                phase = make_trajectory(chunks, dist, motion, t_max)
+
+                for (k, t_evo_k) in enumerate(t_evo):
+                    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)
+                    cc[:, k] += np.cos(dephased)*np.cos(rephased)
+                    ss[:, k] += np.sin(dephased)*np.sin(rephased)
+
+                print_step(n, num_traj, start)
+
+            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, :], '.-')
+
+            for k in range(num_variables):
+                p0 = [0.5, 0, p.dist.variables.get('tau', 1), 1]
+
+                p_final = []
+                p_final1 = []
+                for k, t_evo_k in enumerate(p.ste.t_evo):
+                    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], '.-')
+
+    plt.show()
 
 
 def run_spectrum_sim(config_file: str):
     p = parse(config_file)
 
     rng, num_traj, t_max, delta, eta, num_variables = _prepare_sim(p)
+    print(num_traj)
 
     num_echos = len(p.spec.t_echo)
     reduction_factor = np.zeros((num_variables, num_echos))
@@ -33,7 +115,7 @@ def run_spectrum_sim(config_file: str):
         dist = p.dist.dist_type(**dist_values, rng=rng)
         print(f'\nStart of {dist}')
 
-        chunks = int(0.6 * t_max / dist_values.get('tau', 1)) + 1
+        chunks = int(0.6 * t_max / dist.mean_tau) + 1
 
         # second loop over parameter of motional model
         for (j, motion_values) in enumerate(p.motion):
@@ -41,26 +123,25 @@ def run_spectrum_sim(config_file: str):
             motion = p.motion.model(delta, eta, **motion_values, rng=rng)
             print(f'Start of {motion}')
 
-            print(f'Simulate in chunks of {chunks}')
-
             timesignal = np.zeros((p.spec.num_points, num_echos))
 
             start = time()
 
             # inner loop to create trajectories
             for n in range(num_traj):
-                phase_interpol = make_trajectory(chunks, dist, motion, t_max)
+                phase = make_trajectory(chunks, dist, motion, t_max)
 
                 for (k, t_echo_k) in enumerate(t_echo):
                     # effect of de-phasing and re-phasing
-                    start_amp = -2 * phase_interpol(t_echo_k)
+                    start_amp = -2 * phase(t_echo_k)
 
                     # start of actual acquisition
-                    timesignal[:, k] += np.cos(start_amp + phase_interpol(p.spec.t_acq + 2 * t_echo_k)) * p.spec.dampening
+                    timesignal[:, k] += np.cos(start_amp + phase(p.spec.t_acq + 2*t_echo_k))
-                    reduction_factor[max(p.motion.num_variables, 1)*i + j, k] += np.cos(phase_interpol(2 * t_echo_k) + start_amp)
+                    reduction_factor[max(p.motion.num_variables, 1)*i+j, k] += np.cos(phase(2*t_echo_k) + start_amp)
 
                 print_step(n, num_traj, start)
 
+            timesignal *= p.spec.dampening[:, None]
             timesignal /= num_traj
 
             # FT to spectrum
@@ -73,106 +154,35 @@ def run_spectrum_sim(config_file: str):
             ax.set_title(f'{dist}, {motion}')
             ax.legend(lines, t_echo_strings)
 
+            plt.show()
+
     fig2, ax2 = plt.subplots()
     ax2.semilogx(p.dist.variables['tau'], reduction_factor / num_traj, 'o--')
 
     plt.show()
 
 
-def run_ste_sim(config_file: str):
-    p = parse(config_file)
-
-    rng, num_traj, t_max, delta, eta, num_variables = _prepare_sim(p)
-
-    cc = np.zeros((len(p.ste.t_mix), num_variables, len(p.ste.t_evo)))
-    ss = np.zeros((len(p.ste.t_mix), num_variables, len(p.ste.t_evo)))
-
-    # outer loop over variables of distribution of correlation times
-    for (i, dist_values) in enumerate(p.dist):
-        # noinspection PyCallingNonCallable
-        dist = p.dist.dist_type(**dist_values, rng=rng)
-        print(f'\nStart of {dist}')
-
-        chunks = int(0.6 * t_max / dist_values.get('tau', 1)) + 1
-
-        # second loop over parameter of motional model
-        for (j, motion_values) in enumerate(p.motion):
-            # noinspection PyCallingNonCallable
-            motion = p.motion.model(delta, eta, **motion_values, rng=rng)
-
-            print(f'Start of {motion}')
-            print(f'Simulate in chunks of {chunks}')
-
-            start = time()
-
-            # inner loop to create trajectories
-            for n in range(num_traj):
-                phase_interpol = make_trajectory(chunks, dist, motion, t_max)
-
-                for (k, t_evo_k) in enumerate(p.ste.t_evo):
-                    dephased = phase_interpol(t_evo_k)
-                    rephased = phase_interpol(p.ste.t_mix + 2*t_evo_k)-phase_interpol(p.ste.t_mix+t_evo_k)
-                    cc[:, max(p.motion.num_variables, 1)*i + j, k] += np.cos(dephased)*np.cos(rephased)
-                    ss[:, max(p.motion.num_variables, 1)*i + j, k] += np.sin(dephased)*np.sin(rephased)
-
-                print_step(n, num_traj, start)
-
-    cc /= num_traj
-    ss /= num_traj
-
-    fig, ax = plt.subplots()
-    fig2, ax2 = plt.subplots()
-    fig5, ax5 = plt.subplots()
-    fig3, ax3 = plt.subplots()
-    fig4, ax4 = plt.subplots()
-
-    for j in range(num_variables):
-        p0 = [0.5, 0, 1e-2, 1]
-
-        ax3.plot(p.ste.t_evo, cc[0, j, :])
-        ax3.plot(p.ste.t_evo, ss[0, j, :])
-        ax4.plot(p.ste.t_evo, cc[-1, j, :] / cc[0, j, :])
-        ax4.plot(p.ste.t_evo, ss[-1, j, :] / ss[0, j, :])
-        p_final = []
-        p_final1 = []
-        for k, t_evo_k in enumerate(p.ste.t_evo):
-            res = curve_fit(ste, p.ste.t_mix, cc[:, j, k], p0=p0)
-            res2 = curve_fit(ste, p.ste.t_mix, ss[:, j, k], p0=p0)
-            p_final.append(res[0].tolist())
-            p_final1.append(res2[0].tolist())
-
-        p_final = np.array(p_final)
-        p_final1 = np.array(p_final1)
-        ax.plot(p.ste.t_evo, p_final[:, 0])
-        ax.plot(p.ste.t_evo, p_final1[:, 0])
-        ax.plot(p.ste.t_evo, p_final[:, 1])
-        ax.plot(p.ste.t_evo, p_final1[:, 1])
-        ax5.semilogy(p.ste.t_evo, p_final[:, 2])
-        ax5.semilogy(p.ste.t_evo, p_final1[:, 2])
-        ax2.plot(p.ste.t_evo, p_final[:, 3])
-        ax2.plot(p.ste.t_evo, p_final1[:, 3])
-
-    plt.show()
-
-
 def print_step(n, num_traj, start):
-    if (n + 1) % 200 == 0:
+    n_1 = n+1
+    if n_1 % 200 == 0 or n_1 == num_traj:
         elapsed = time() - start
-        print(f'Step {n + 1} of {num_traj}', end=' - ')
+        print(f'Step {n_1} of {num_traj}', end=' - ')
-        total = num_traj * elapsed / (n + 1)
+        total = num_traj * elapsed / (n_1)
         print(f'total: {total:.2f}s - elapsed: {elapsed:.2f}s - remaining: {total - elapsed:.2f}s')
 
 
-def make_trajectory(chunks: int, dist, motion, t_max: float):
+def make_trajectory(chunks: int, dist: BaseDistribution, motion: BaseMotion, t_max: float):
+    motion.start()
+    dist.start()
+
     t_passed = 0
     t = [0]
     phase = [0]
-    accumulated_phase = 0
     while t_passed < t_max:
-        # orientation until the next jump
+        # frequencies between jumps
         current_omega = motion.jump(size=chunks)
 
-        # time to next jump
+        # times at a particular position
         t_wait = dist.wait(size=chunks)
 
         accumulated_phase = np.cumsum(t_wait * current_omega) + phase[-1]
@@ -202,11 +212,11 @@ def _prepare_sim(parameter: Parameter) -> tuple[Generator, int, float, float, fl
     # parameter for omega_q
     delta, eta = parameter.molecule.delta, parameter.molecule.eta
 
-    num_variables = parameter.dist.num_variables + parameter.motion.num_variables
+    num_variables = parameter.dist.num_variables * parameter.motion.num_variables
 
     return rng, num_traj, t_max, delta, eta, num_variables
 
 
-if __name__ == '__main__':
+def ste(x, a, f_infty, tau, beta):
-    run_ste_sim('../config.json')
+    return a*((1-f_infty) * np.exp(-(x/tau)**beta) + f_infty)
-    run_spectrum_sim('../config.json')
+