From 82acade7d53da959394202e5d9b1f278532e3c81 Mon Sep 17 00:00:00 2001 From: Dominik Demuth Date: Thu, 20 Jun 2024 19:19:55 +0200 Subject: [PATCH] moving around --- pyproject.toml | 4 +- src/config.json | 30 ++-- src/main.py | 5 + src/rwsims/distributions.py | 51 ++++++- src/rwsims/helper.py | 65 +++++++++ src/rwsims/motions.py | 91 ++++++++++-- src/rwsims/parameter.py | 27 ++-- src/rwsims/{parser.py => parsing.py} | 13 +- src/rwsims/sims.py | 210 ++++++++++++++------------- 9 files changed, 349 insertions(+), 147 deletions(-) create mode 100644 src/main.py create mode 100644 src/rwsims/helper.py rename src/rwsims/{parser.py => parsing.py} (90%) diff --git a/pyproject.toml b/pyproject.toml index b6fa3fb..11a86b3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,8 +25,8 @@ classifiers = [ ] [project.scripts] -rw_spectra = "scripts.sim_spectra" -rw_ste = "scripts.sim_ste" +rw_spectra = "sims.run_spectrum_sim" +rw_ste = "sims.run_ste_sim" [project.urls] Homepage = "https://github.com/pypa/sampleproject" diff --git a/src/config.json b/src/config.json index 24ab8f3..19ebc43 100644 --- a/src/config.json +++ b/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, - "steps": 6, + "stop": 1e-8, + "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, - "stop": 1e-2, - "steps": 10, + "start": 1e-7, + "stop": 1e-1, + "steps": 31, "is_log": true - } + }, + "t_echo": 15e-6 } } \ No newline at end of file diff --git a/src/main.py b/src/main.py new file mode 100644 index 0000000..d75e40e --- /dev/null +++ b/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') diff --git a/src/rwsims/distributions.py b/src/rwsims/distributions.py index a330944..1490039 100644 --- a/src/rwsims/distributions.py +++ b/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) \ No newline at end of file diff --git a/src/rwsims/helper.py b/src/rwsims/helper.py new file mode 100644 index 0000000..8f7fb1f --- /dev/null +++ b/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 diff --git a/src/rwsims/motions.py b/src/rwsims/motions.py index f6f380d..6e55591 100644 --- a/src/rwsims/motions.py +++ b/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 - -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)) +from .helper import xyz_to_spherical, spherical_to_xyz, omega_q, draw_orientation, get_rotation_matrix -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): +class BaseMotion(ABC): + def __init__(self, delta: float, eta: float, rng: Generator): 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] + + diff --git a/src/rwsims/parameter.py b/src/rwsims/parameter.py index f55d5bf..637029e 100644 --- a/src/rwsims/parameter.py +++ b/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 src.rwsims.motions import RandomJump +from .distributions import BaseDistribution +from .motions import BaseMotion - -__all__ = ['SimParameter', 'MoleculeParameter', 'StimEchoParameter', 'SpectrumParameter', 'DistParameter', 'MotionParameter', 'Parameter'] +__all__ = [ + '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 diff --git a/src/rwsims/parser.py b/src/rwsims/parsing.py similarity index 90% rename from src/rwsims/parser.py rename to src/rwsims/parsing.py index 3f6cb95..a2cc9bf 100644 --- a/src/rwsims/parser.py +++ b/src/rwsims/parsing.py @@ -5,11 +5,9 @@ from typing import Any import numpy as np -from distributions import ( - DeltaDistribution -) -from motions import RandomJump -from parameter import * +from .distributions import DeltaDistribution, LogGaussianDistribution +from .motions import RandomJump, TetrahedralJump +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( diff --git a/src/rwsims/sims.py b/src/rwsims/sims.py index fc1d8f5..3bd19a6 100644 --- a/src/rwsims/sims.py +++ b/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 parser import parse +from .parameter import Parameter +from .distributions import BaseDistribution +from .motions import BaseMotion +from .parsing import parse -def ste(x, a, f_infty, tau, beta): - return a*((1-f_infty) * np.exp(-(x/tau)**beta) + f_infty) +def run_ste_sim(config_file: str): + 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 - reduction_factor[max(p.motion.num_variables, 1)*i + j, k] += np.cos(phase_interpol(2 * t_echo_k) + start_amp) + 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(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]) + ax2.semilogx(p.dist.variables['tau'], reduction_factor / num_traj, 'o--') 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=' - ') - total = num_traj * elapsed / (n + 1) + 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): +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__': - run_ste_sim('../config.json') - run_spectrum_sim('../config.json') +def ste(x, a, f_infty, tau, beta): + return a*((1-f_infty) * np.exp(-(x/tau)**beta) + f_infty) +