From b2164c944a58207e9a7ca45e463edd5f0951b34b Mon Sep 17 00:00:00 2001 From: Dominik Demuth Date: Wed, 19 Jun 2024 19:10:49 +0200 Subject: [PATCH] more modularity --- LICENSE | 2 +- pyproject.toml | 33 +++ python/rj_spectrum.py | 114 ---------- python/rj_spectrum_chunk.py | 123 ---------- src/config.json | 39 ++++ src/config_ste.json | 30 +++ src/rwsims/__init__.py | 0 src/rwsims/distributions.py | 16 ++ src/rwsims/motions.py | 33 +++ src/rwsims/parameter.py | 108 +++++++++ src/rwsims/parser.py | 130 +++++++++++ src/rwsims/sims.py | 212 ++++++++++++++++++ .../rwsims}/tetrahedral_spectrum.py | 0 .../rwsims}/tetrahedral_spectrum_chunk.py | 17 +- 14 files changed, 614 insertions(+), 243 deletions(-) create mode 100644 pyproject.toml delete mode 100644 python/rj_spectrum.py delete mode 100644 python/rj_spectrum_chunk.py create mode 100644 src/config.json create mode 100644 src/config_ste.json create mode 100644 src/rwsims/__init__.py create mode 100644 src/rwsims/distributions.py create mode 100644 src/rwsims/motions.py create mode 100644 src/rwsims/parameter.py create mode 100644 src/rwsims/parser.py create mode 100644 src/rwsims/sims.py rename {python => src/rwsims}/tetrahedral_spectrum.py (100%) rename {python => src/rwsims}/tetrahedral_spectrum_chunk.py (92%) diff --git a/LICENSE b/LICENSE index bbf9022..5caf13f 100644 --- a/LICENSE +++ b/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: diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..b6fa3fb --- /dev/null +++ b/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" \ No newline at end of file diff --git a/python/rj_spectrum.py b/python/rj_spectrum.py deleted file mode 100644 index 68145bf..0000000 --- a/python/rj_spectrum.py +++ /dev/null @@ -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() diff --git a/python/rj_spectrum_chunk.py b/python/rj_spectrum_chunk.py deleted file mode 100644 index 53077b7..0000000 --- a/python/rj_spectrum_chunk.py +++ /dev/null @@ -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() diff --git a/src/config.json b/src/config.json new file mode 100644 index 0000000..24ab8f3 --- /dev/null +++ b/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 + } + } +} \ No newline at end of file diff --git a/src/config_ste.json b/src/config_ste.json new file mode 100644 index 0000000..befdd8b --- /dev/null +++ b/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 + } + } +} \ No newline at end of file diff --git a/src/rwsims/__init__.py b/src/rwsims/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/rwsims/distributions.py b/src/rwsims/distributions.py new file mode 100644 index 0000000..a330944 --- /dev/null +++ b/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) diff --git a/src/rwsims/motions.py b/src/rwsims/motions.py new file mode 100644 index 0000000..f6f380d --- /dev/null +++ b/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) diff --git a/src/rwsims/parameter.py b/src/rwsims/parameter.py new file mode 100644 index 0000000..f55d5bf --- /dev/null +++ b/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 diff --git a/src/rwsims/parser.py b/src/rwsims/parser.py new file mode 100644 index 0000000..3f6cb95 --- /dev/null +++ b/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 diff --git a/src/rwsims/sims.py b/src/rwsims/sims.py new file mode 100644 index 0000000..fc1d8f5 --- /dev/null +++ b/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') diff --git a/python/tetrahedral_spectrum.py b/src/rwsims/tetrahedral_spectrum.py similarity index 100% rename from python/tetrahedral_spectrum.py rename to src/rwsims/tetrahedral_spectrum.py diff --git a/python/tetrahedral_spectrum_chunk.py b/src/rwsims/tetrahedral_spectrum_chunk.py similarity index 92% rename from python/tetrahedral_spectrum_chunk.py rename to src/rwsims/tetrahedral_spectrum_chunk.py index 9b9fafb..555676a 100644 --- a/python/tetrahedral_spectrum_chunk.py +++ b/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()