more modularity

LICENSE

@@ -1,4 +1,4 @@
-Copyright (c) 2024 dominik. 
+Copyright (c) 2024 Dominik Demuth
 
 Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
 

pyproject.toml

@@ -0,0 +1,33 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "rwsims"
+version = "0.0.1"
+authors = [
+  { name="Dominik Demuth", email="dominik.demuth@pkm.tu-darmstadt.de" },
+]
+maintainers = [
+    { name="Dominik Demuth", email="dominik.demuth@pkm.tu-darmstadt.de"}
+]
+description = "A small example package"
+readme = "README.md"
+requires-python = ">=3.8"
+dependencies = [
+    "numpy",
+    "matplotlib"
+]
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: BSD-3 Clause",
+    "Operating System :: OS Independent",
+]
+
+[project.scripts]
+rw_spectra = "scripts.sim_spectra"
+rw_ste = "scripts.sim_ste"
+
+[project.urls]
+Homepage = "https://github.com/pypa/sampleproject"
+Issues = "https://github.com/pypa/sampleproject/issues"

python/rj_spectrum.py

@@ -1,114 +0,0 @@
-from time import time
-
-import numpy as np
-from scipy.interpolate import interp1d
-import matplotlib.pyplot as plt
-
-
-# spectral parameter
-delta = 161e3  # in Hz
-eta = 0
-lb = 2e3  # in Hz
-
-# correlation time
-tau = [1e-7]  # in s
-
-# acquisition parameter
-acq_length = 4096
-dt = 1e-6  # in s
-t_echo = [0, 5e-6, 10e-6, 20e-6, 50e-6, 100e-6, 200e-6]  # all in s
-
-# derived parameter
-t_acq = np.arange(acq_length) * dt
-t_max = acq_length*dt + 2*max(t_echo)
-dampening = np.exp(-lb * t_acq)
-
-# random number generator
-seed = 1234
-rng = np.random.default_rng(seed)
-
-# number of random walkers
-num_traj = 5000
-
-
-def omega_q(delta_: float, eta_: float, theta_: float, phi_: float) -> float:
-    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 new_orientation(delta_: float, eta_: float) -> float:
-    z_theta, z_phi = rng.random(2)
-    theta = np.arccos(1 - 2 * z_theta)
-    phi = 2 * np.pi * z_phi
-
-    return omega_q(delta_, eta_, theta, phi)
-
-
-for tau_i in tau:
-    print(f'\nStart for tau={tau_i}')
-
-    timesignal = np.zeros((acq_length, len(t_echo)))
-
-    start = time()
-    expected_jumps = round(t_max/tau_i)
-    if expected_jumps > 1e7:
-        print(f'Too many jumps to process, Skip {tau_i}s')
-        continue
-
-    for i in range(num_traj):
-        t_passed = 0
-        t = [0]
-        phase = [0]
-        accumulated_phase = 0
-
-        while t_passed < t_max:
-            # orientation until the next jump
-            current_omega = new_orientation(delta, eta)
-
-            # time to next jump
-            t_wait = rng.exponential(tau_i)
-
-            t_passed += t_wait
-            accumulated_phase += t_wait * current_omega
-
-            t.append(t_passed)
-            phase.append(accumulated_phase)
-
-        # convenient interpolation to get phase at arbitrary times
-        phase_interpol = interp1d(t, phase)
-
-        for j, t_echo_j in enumerate(t_echo):
-            # effect of de-phasing and re-phasing
-            start_amp = -2 * phase_interpol(t_echo_j)
-
-            # start of actual acquisition
-            timesignal[:, j] += np.cos(start_amp + phase_interpol(t_acq + 2*t_echo_j)) * dampening
-
-        if (i+1) % 200 == 0:
-            elapsed = time()-start
-            print(f'Step {i+1} of {num_traj}', end=' - ')
-            total = num_traj * elapsed / (i+1)
-            print(f'elapsed: {elapsed:.2f}s - total: {total:.2f}s - remaining: {total-elapsed:.2f}s')
-
-    timesignal /= num_traj
-
-    # FT to spectrum
-    freq = np.fft.fftshift(np.fft.fftfreq(acq_length, dt))
-    spec = np.fft.fftshift(np.fft.fft(timesignal, axis=0), axes=0).real
-    spec -= spec[0]
-
-    t_echo_strings = list(map(str, t_echo))
-
-    # plot spectra
-    fig, ax = plt.subplots()
-    lines = ax.plot(freq, spec)
-    ax.set_title(f'RJ (tau = {tau_i}s)')
-    ax.legend(lines, t_echo_strings)
-    # plt.savefig(f'RJ_{tau_i}.png')
-
-    # # save time signals and spectra
-    # np.savetxt(f'rj_spectrum_{tau_i}.dat', np.c_[freq, spec], header='f\t' + '\t'.join(t_echo_strings))
-    # np.savetxt(f'rj_timesignal_{tau_i}.dat', np.c_[t_acq, timesignal], header='t\t' + '\t'.join(t_echo_strings))
-
-    plt.show()

python/rj_spectrum_chunk.py

@@ -1,123 +0,0 @@
-from time import time
-
-import numpy as np
-from scipy.interpolate import interp1d
-import matplotlib.pyplot as plt
-
-
-# spectral parameter
-delta = 161e3  # in Hz
-eta = 0
-lb = 5e3  # in Hz
-
-# correlation time
-tau = [1e-5]  # in s
-
-# acquisition parameter
-acq_length = 4096
-dt = 1e-6  # in s
-t_echo = [0, 5e-6, 10e-6, 20e-6, 50e-6, 100e-6, 200e-6]  # all in s
-
-# derived parameter
-t_acq = np.arange(acq_length) * dt
-t_max = acq_length*dt + 2*max(t_echo)
-dampening = np.exp(-lb * t_acq)
-
-# random number generator
-seed = None
-rng = np.random.default_rng(seed)
-
-# number of random walkers
-num_traj = 50000
-
-
-def omega_q(delta_: float, eta_: float, theta_: float, phi_: float) -> np.ndarray:
-    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 new_orientation(delta_: float, eta_: float, size=1) -> np.ndarray:
-    z_theta, z_phi = rng.random((2, size))
-    theta = np.arccos(1 - 2 * z_theta)
-    phi = 2 * np.pi * z_phi
-
-    return omega_q(delta_, eta_, theta, phi)
-
-
-def new_tau(size=1) -> np.ndarray:
-    return rng.exponential(tau_i, size=size)
-
-
-for tau_i in tau:
-    print(f'\nStart for tau={tau_i}')
-
-    timesignal = np.zeros((acq_length, len(t_echo)))
-
-    start = time()
-    expected_jumps = round(t_max/tau_i)
-    if expected_jumps > 1e7:
-        print(f'Too many jumps to process, Skip {tau_i}s')
-        continue
-
-    chunks = int(0.6 * t_max / tau_i) + 1
-    print(f'Chunk size for trajectories: {chunks}')
-
-    for i in range(num_traj):
-        t_passed = 0
-        t = [0]
-        phase = [0]
-        accumulated_phase = 0
-
-        while t_passed < t_max:
-            # orientation until the next jump
-            current_omega = new_orientation(delta, eta, size=chunks)
-
-            # time to next jump
-            t_wait = new_tau(size=chunks)
-            accumulated_phase = np.cumsum(t_wait*current_omega) + phase[-1]
-
-            t_wait = np.cumsum(t_wait) + t_passed
-            t_passed = t_wait[-1]
-            t.extend(t_wait.tolist())
-
-            phase.extend(accumulated_phase.tolist())
-
-        # convenient interpolation to get phase at arbitrary times
-        phase_interpol = interp1d(t, phase)
-
-        for j, t_echo_j in enumerate(t_echo):
-            # effect of de-phasing and re-phasing
-            start_amp = -2 * phase_interpol(t_echo_j)
-
-            # start of actual acquisition
-            timesignal[:, j] += np.cos(start_amp + phase_interpol(t_acq + 2*t_echo_j)) * dampening
-
-        if (i+1) % 200 == 0:
-            elapsed = time()-start
-            print(f'Step {i+1} of {num_traj}', end=' - ')
-            total = num_traj * elapsed / (i+1)
-            print(f'elapsed: {elapsed:.2f}s - total: {total:.2f}s - remaining: {total-elapsed:.2f}s')
-
-    timesignal /= num_traj
-
-    # FT to spectrum
-    freq = np.fft.fftshift(np.fft.fftfreq(acq_length, dt))
-    spec = np.fft.fftshift(np.fft.fft(timesignal, axis=0), axes=0).real
-    spec -= spec[0]
-    # spec /= np.max(spec, axis=0)
-
-    t_echo_strings = list(map(str, t_echo))
-
-    # plot spectra
-    fig, ax = plt.subplots()
-    lines = ax.plot(freq, spec)
-    ax.set_title(f'RJ (tau = {tau_i}s)')
-    ax.legend(lines, t_echo_strings)
-    # plt.savefig(f'RJ_{tau_i}.png')
-
-    # # save time signals and spectra
-    # np.savetxt(f'rj_spectrum_{tau_i}_chunky.dat', np.c_[freq, spec], header='f\t' + '\t'.join(t_echo_strings))
-    # np.savetxt(f'rj_timesignal_{tau_i}_chunky.dat', np.c_[t_acq, timesignal], header='t\t' + '\t'.join(t_echo_strings))
-
-    plt.show()

src/config.json

@@ -0,0 +1,39 @@
+{
+  "simulation": {
+    "num_walker": 2500,
+    "seed": null
+  },
+  "molecule": {
+    "delta": 161e3,
+    "eta": 0.0
+  },
+  "correlation_times": {
+    "distribution": "DeltaDistribution",
+    "tau": {
+      "start": 1e-4,
+      "stop": 1e-2,
+      "steps": 6,
+      "is_log": true
+    }
+  },
+  "motion": {
+    "model": "RandomJump"
+  },
+  "spectrum": {
+    "dwell_time": 1e-6,
+    "num_points": 4096,
+    "t_echo": {
+      "list": [0, 5e-6, 10e-6, 20e-6, 50e-6, 100e-6, 200e-6]
+    },
+    "line_broadening": 2e3
+  },
+  "stimulated_echo": {
+    "t_evo": 10e-6,
+    "t_mix": {
+      "start": 1e-5,
+      "stop": 1e-2,
+      "steps": 10,
+      "is_log": true
+    }
+  }
+}

src/config_ste.json

@@ -0,0 +1,30 @@
+{
+  "simulation": {
+    "num_walker": 20000,
+    "seed": null
+  },
+  "molecule": {
+    "delta": 161e3,
+    "eta": 0.0
+  },
+  "correlation_times": {
+    "distribution": "DeltaDistribution",
+    "tau": 1e-2
+  },
+  "motion": {
+    "model": "RandomJump"
+  },
+  "stimulated_echo": {
+    "t_evo": {
+      "start": 1e-6,
+      "stop": 40e-6,
+      "steps": 80
+    },
+    "t_mix": {
+      "start": 1e-5,
+      "stop": 1e0,
+      "steps": 21,
+      "is_log": true
+    }
+  }
+}

src/rwsims/distributions.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+from numpy.typing import ArrayLike
+from numpy.random import Generator
+
+
+class DeltaDistribution:
+    def __init__(self, tau: float, rng: Generator | None = None):
+        self._tau = tau
+        self._rng = rng
+
+    def __repr__(self):
+        return f'DeltaDistribution (tau={self._tau})'
+
+    def wait(self, size: int = 1) -> ArrayLike:
+        return self._rng.exponential(self._tau, size=size)

src/rwsims/motions.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+import numpy as np
+from numpy.random import Generator
+from numpy.typing import ArrayLike
+
+
+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:
+    z_theta, z_phi = rng.random((2, size))
+    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
+
+    def __repr__(self):
+        return 'Random Jump'
+
+    def jump(self, size: int = 1) -> ArrayLike:
+        return draw_orientation(self._delta, self._eta, self._rng, size=size)

src/rwsims/parameter.py

@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from itertools import product
+from typing import Any
+
+import numpy as np
+
+from src.rwsims.distributions import DeltaDistribution
+from src.rwsims.motions import RandomJump
+
+
+__all__ = ['SimParameter', 'MoleculeParameter', 'StimEchoParameter', 'SpectrumParameter', 'DistParameter', 'MotionParameter', 'Parameter']
+
+
+@dataclass
+class SimParameter:
+    seed: int | None
+    num_walker: int
+    t_max: float
+
+
+@dataclass
+class MoleculeParameter:
+    delta: float
+    eta: float
+
+
+@dataclass
+class StimEchoParameter:
+    t_evo: np.ndarray
+    t_mix: np.ndarray
+    t_max: float = field(init=False)
+
+    def __post_init__(self):
+        self.t_max = np.max(self.t_mix) + 2 * np.max(self.t_evo)
+
+
+@dataclass
+class SpectrumParameter:
+    dwell_time: float
+    num_points: int
+    t_echo: np.ndarray
+    t_acq: np.ndarray = field(init=False)
+    t_max: float = field(init=False)
+    lb: float
+    dampening: np.ndarray = field(init=False)
+
+    def __post_init__(self):
+        self.t_acq = np.arange(self.num_points) * self.dwell_time
+        self.dampening = np.exp(-self.lb * self.t_acq)
+        self.t_max = np.max(self.t_acq) + 2 * np.max(self.t_echo)
+
+
+@dataclass
+class DistParameter:
+    dist_type: DeltaDistribution
+    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()))
+
+    def __iter__(self):
+        return self
+
+    def __next__(self) -> dict[str, Any]:
+        if self.iter is None:
+            self.iter = product(*self.variables.values())
+        try:
+            return dict(zip(self.variables.keys(), next(self.iter)))
+        except StopIteration:
+            self.iter = None
+            raise StopIteration
+
+
+@dataclass
+class MotionParameter:
+    model: RandomJump
+    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()))
+
+    def __iter__(self):
+        return self
+
+    def __next__(self) -> dict[str, Any]:
+        if self.iter is None:
+            self.iter = product(*self.variables.values())
+        try:
+            return dict(zip(self.variables.keys(), next(self.iter)))
+        except StopIteration:
+            self.iter = None
+            raise StopIteration
+
+
+@dataclass
+class Parameter:
+    ste: StimEchoParameter | None
+    spec: SpectrumParameter | None
+    sim: SimParameter
+    dist: DistParameter
+    motion: MotionParameter
+    molecule: MoleculeParameter

src/rwsims/parser.py

@@ -0,0 +1,130 @@
+from __future__ import annotations
+
+import json
+from typing import Any
+
+import numpy as np
+
+from distributions import (
+    DeltaDistribution
+)
+from motions import RandomJump
+from parameter import *
+
+
+def parse(config_file: str) -> Parameter:
+    with open(config_file, 'r') as f:
+        parameter: dict = json.load(f)
+
+    ste = _parse_ste(parameter.get('stimulated_echo'))
+    spec = _parse_spectrum(parameter.get('spectrum'))
+
+    if ste is None and spec is None:
+        raise ValueError("No parameter for STE or spectra given")
+
+    t_max = 0
+    if spec is not None:
+        t_max = max(spec.t_max, t_max)
+    if ste is not None:
+        t_max = max(ste.t_max, t_max)
+    parameter['simulation'].update({'t_max': t_max})
+
+    sim = _parse_sim(parameter['simulation'])
+    dist = _parse_dist(parameter['correlation_times'])
+    motion = _parse_motion(parameter['motion'])
+    mol = _parse_molecule(parameter['molecule'])
+
+    p = Parameter(sim=sim, ste=ste, spec=spec, dist=dist, motion=motion, molecule=mol)
+
+    return p
+
+
+def _parse_sim(params: dict[str, Any]) -> SimParameter:
+    sim = SimParameter(
+        num_walker=params['num_walker'],
+        seed=params['seed'],
+        t_max=params['t_max']
+    )
+    return sim
+
+
+def _parse_ste(params: dict[str, Any] | None) -> StimEchoParameter | None:
+    if params is None:
+        return
+
+    ste = StimEchoParameter(
+        t_mix=_make_times(params['t_mix']),
+        t_evo=_make_times(params['t_evo']),
+    )
+    return ste
+
+
+def _parse_spectrum(params: dict[str, Any] | None) -> SpectrumParameter | None:
+    if params is None:
+        return
+
+    spec = SpectrumParameter(
+        num_points=params['num_points'],
+        dwell_time=params['dwell_time'],
+        t_echo=_make_times(params['t_echo']),
+        lb=params['line_broadening']
+    )
+
+    return spec
+
+
+def _parse_dist(params: dict[str, Any]) -> DistParameter:
+    mapping: dict = {
+        'DeltaDistribution': DeltaDistribution
+    }
+    p = DistParameter(
+        dist_type=mapping[params['distribution']],
+        variables={k: _make_times(v) for k, v in params.items() if k != 'distribution'},
+    )
+
+    return p
+
+
+def _parse_motion(params: dict[str, Any]) -> MotionParameter:
+    mapping = {
+        'RandomJump':  RandomJump,
+    }
+
+    p = MotionParameter(
+        model=mapping[params['model']],
+        variables={k: _make_times(v) for k, v in params.items() if k != 'model'}
+    )
+    return p
+
+
+def _parse_molecule(params: dict[str, Any]) -> MoleculeParameter:
+    return MoleculeParameter(
+        delta=params['delta'],
+        eta=params['eta']
+    )
+
+
+def _make_times(params: float | int | dict[str, Any]) -> np.ndarray:
+    times = None
+
+    if isinstance(params, (int, float, complex)):
+        times = np.array([params])
+
+    else:
+        if all(k in params for k in ('start', 'stop', 'steps')):
+            space_func = np.linspace
+            if 'is_log' in params and params['is_log']:
+                space_func = np.geomspace
+
+            times = space_func(start=params['start'], stop=params['stop'], num=params['steps'])
+
+        if 'list' in params:
+            if times is not None:
+                raise ValueError('list and range is given')
+
+            times = np.array(params['list'])
+
+    if times is None:
+        raise ValueError('No times are given')
+
+    return times

src/rwsims/sims.py

@@ -0,0 +1,212 @@
+from __future__ import annotations
+
+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 parameter import Parameter
+from parser import parse
+
+
+def ste(x, a, f_infty, tau, beta):
+    return a*((1-f_infty) * np.exp(-(x/tau)**beta) + f_infty)
+
+
+def run_spectrum_sim(config_file: str):
+    p = parse(config_file)
+
+    rng, num_traj, t_max, delta, eta, num_variables = _prepare_sim(p)
+
+    num_echos = len(p.spec.t_echo)
+    reduction_factor = np.zeros((num_variables, num_echos))
+    freq = np.fft.fftshift(np.fft.fftfreq(p.spec.num_points, p.spec.dwell_time))
+    t_echo = p.spec.t_echo
+    t_echo_strings = list(map(str, t_echo))
+
+    # 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}')
+
+            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)
+
+                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 of actual acquisition
+                    timesignal[:, k] += np.cos(start_amp + phase_interpol(p.spec.t_acq + 2 * t_echo_k)) * p.spec.dampening
+                    reduction_factor[max(p.motion.num_variables, 1)*i + j, k] += np.cos(phase_interpol(2 * t_echo_k) + start_amp)
+
+                print_step(n, num_traj, start)
+
+            timesignal /= num_traj
+
+            # FT to spectrum
+            spec = np.fft.fftshift(np.fft.fft(timesignal, axis=0), axes=0).real
+            spec -= spec[0]
+
+            # plot spectra
+            fig, ax = plt.subplots()
+            lines = ax.plot(freq, spec)
+            ax.set_title(f'{dist}, {motion}')
+            ax.legend(lines, t_echo_strings)
+
+    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:
+        elapsed = time() - start
+        print(f'Step {n + 1} of {num_traj}', end=' - ')
+        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):
+    t_passed = 0
+    t = [0]
+    phase = [0]
+    accumulated_phase = 0
+    while t_passed < t_max:
+        # orientation until the next jump
+        current_omega = motion.jump(size=chunks)
+
+        # time to next jump
+        t_wait = dist.wait(size=chunks)
+
+        accumulated_phase = np.cumsum(t_wait * current_omega) + phase[-1]
+
+        t_wait = np.cumsum(t_wait) + t_passed
+        t_passed = t_wait[-1]
+        t.extend(t_wait.tolist())
+
+        phase.extend(accumulated_phase.tolist())
+
+    # convenient interpolation to get phase at arbitrary times
+    phase_interpol = interp1d(t, phase)
+
+    return phase_interpol
+
+
+def _prepare_sim(parameter: Parameter) -> tuple[Generator, int, float, float, float, int]:
+    # random number generator
+    rng = np.random.default_rng(parameter.sim.seed)
+
+    # number of random walkers
+    num_traj = parameter.sim.num_walker
+
+    # length of trajectories
+    t_max = parameter.sim.t_max
+
+    # parameter for omega_q
+    delta, eta = parameter.molecule.delta, parameter.molecule.eta
+
+    num_variables = parameter.dist.num_variables + parameter.motion.num_variables
+
+    return rng, num_traj, t_max, delta, eta, num_variables
+
+
+if __name__ == '__main__':
+    run_ste_sim('../config.json')
+    run_spectrum_sim('../config.json')

src/rwsims/tetrahedral_spectrum_chunk.py

@@ -11,12 +11,12 @@ eta = 0
 lb = 5e3  # in Hz
 
 # correlation time
-tau = [1e-6]  # in s
+tau = np.logspace(-8, -1, num=15)  # in s
 
 # acquisition parameter
 acq_length = 4096
 dt = 1e-6  # in s
-t_echo = [0, 5e-6, 10e-6, 20e-6, 50e-6, 100e-6, 200e-6]  # all in s
+t_echo = [5e-6, 10e-6, 20e-6, 50e-6, 100e-6, 200e-6]  # all in s
 
 # derived parameter
 t_acq = np.arange(acq_length) * dt
@@ -28,7 +28,7 @@ seed = None
 rng = np.random.default_rng(seed)
 
 # number of random walkers
-num_traj = 50000
+num_traj = 1
 
 
 def omega_q(delta_: float, eta_: float, theta_: ArrayLike, phi_: ArrayLike) -> np.ndarray:
@@ -60,7 +60,10 @@ def new_tau(size=1) -> np.ndarray:
     return rng.exponential(tau_i, size=size)
 
 
-for tau_i in tau:
+reduction_factor = np.zeros((len(tau), len(t_echo)))
+
+
+for (n, tau_i) in enumerate(tau):
     print(f'\nStart for tau={tau_i}')
 
     timesignal = np.zeros((acq_length, len(t_echo)))
@@ -127,6 +130,7 @@ for tau_i in tau:
 
             # start of actual acquisition
             timesignal[:, j] += np.cos(start_amp + phase_interpol(t_acq + 2*t_echo_j)) * dampening
+            reduction_factor[n, j] += np.cos(phase_interpol(2*t_echo_j) + start_amp)
 
         if (i+1) % 200 == 0:
             elapsed = time()-start
@@ -155,4 +159,7 @@ for tau_i in tau:
     # np.savetxt(f'spectrum_{tau_i}.dat', np.c_[freq, spec], header='f\t' + '\t'.join(t_echo_strings))
     # np.savetxt(f'timesignal_{tau_i}.dat', np.c_[t_acq, timesignal], header='t\t' + '\t'.join(t_echo_strings))
 
-    plt.show()
+fig2, ax2 = plt.subplots()
+ax2.semilogx(tau, reduction_factor / num_traj, 'o--')
+
+plt.show()