moving around
This commit is contained in:
Dominik Demuth
parent
b2164c944a
commit
82acade7d5
@ -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"
|
||||
|
||||
@ -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
|
||||
}
|
||||
}
|
||||
5
src/main.py
Normal file
5
src/main.py
Normal file
@ -0,0 +1,5 @@
|
||||
from rwsims.sims import run_ste_sim, run_spectrum_sim
|
||||
|
||||
# run_ste_sim('config.json')
|
||||
|
||||
run_spectrum_sim('config.json')
|
||||
@ -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)
|
||||
65
src/rwsims/helper.py
Normal file
65
src/rwsims/helper.py
Normal file
@ -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
|
||||
@ -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]
|
||||
|
||||
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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(
|
||||
@ -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)
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user