more modularity
This commit is contained in:
Dominik Demuth
parent
b88bd0dbd6
commit
b2164c944a
2
LICENSE
2
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:
|
||||
|
||||
|
||||
33
pyproject.toml
Normal file
33
pyproject.toml
Normal file
@ -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"
|
||||
@ -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()
|
||||
@ -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()
|
||||
39
src/config.json
Normal file
39
src/config.json
Normal file
@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
30
src/config_ste.json
Normal file
30
src/config_ste.json
Normal file
@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
0
src/rwsims/__init__.py
Normal file
0
src/rwsims/__init__.py
Normal file
16
src/rwsims/distributions.py
Normal file
16
src/rwsims/distributions.py
Normal file
@ -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)
|
||||
33
src/rwsims/motions.py
Normal file
33
src/rwsims/motions.py
Normal file
@ -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)
|
||||
108
src/rwsims/parameter.py
Normal file
108
src/rwsims/parameter.py
Normal file
@ -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
|
||||
130
src/rwsims/parser.py
Normal file
130
src/rwsims/parser.py
Normal file
@ -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
|
||||
212
src/rwsims/sims.py
Normal file
212
src/rwsims/sims.py
Normal file
@ -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')
|
||||
@ -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))
|
||||
|
||||
fig2, ax2 = plt.subplots()
|
||||
ax2.semilogx(tau, reduction_factor / num_traj, 'o--')
|
||||
|
||||
plt.show()
|
||||
Loading…
Reference in New Issue
Block a user