285 lines
9.5 KiB
Python
285 lines
9.5 KiB
Python
from __future__ import annotations
|
|
|
|
from time import perf_counter
|
|
|
|
import numpy as np
|
|
from numpy.random import Generator
|
|
from datetime import datetime
|
|
from scipy.interpolate import interp1d
|
|
import matplotlib.pyplot as plt
|
|
from scipy.optimize import curve_fit
|
|
|
|
from .functions import ste
|
|
from .parameter import Parameter
|
|
from .distributions import BaseDistribution
|
|
from .motions import BaseMotion
|
|
from .parsing import parse
|
|
|
|
|
|
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)
|
|
|
|
# second loop over parameter of motion 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}, {motion}')
|
|
|
|
start = last_print = perf_counter()
|
|
|
|
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(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)
|
|
# print(t_evo_k, t0 + t_evo_k + 2*t_echo, t0)
|
|
cc[:, k] += np.cos(dephased)*np.cos(rephased)
|
|
ss[:, k] += np.sin(dephased)*np.sin(rephased)
|
|
|
|
last_print = print_step(n, num_traj, start, last_print)
|
|
|
|
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, dist_values.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)
|
|
|
|
num_echos = len(p.spec.t_echo)
|
|
reduction_factor = np.zeros((num_variables, num_echos))
|
|
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)
|
|
# second loop over parameter of motion 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}, {motion}')
|
|
|
|
timesignal = np.zeros((p.spec.num_points, num_echos))
|
|
|
|
start = perf_counter()
|
|
last_print = start
|
|
|
|
# inner loop to create trajectories
|
|
for n in range(num_traj):
|
|
phase = make_trajectory(dist, motion, t_max)
|
|
|
|
for (k, t_echo_k) in enumerate(t_echo):
|
|
# effect of de-phasing and re-phasing
|
|
start_amp = -2 * phase(t_echo_k)
|
|
|
|
# start of actual acquisition
|
|
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(n+1, num_traj, start, last_print)
|
|
last_print = print_step(n+1, num_traj, start, last_print)
|
|
|
|
# apply line broadening
|
|
timesignal *= p.spec.dampening[:, None]
|
|
timesignal /= num_traj
|
|
timesignal[0, :] /= 2
|
|
|
|
# FT to spectrum
|
|
spec = np.fft.fftshift(np.fft.fft(timesignal, axis=0), axes=0).real
|
|
spec -= spec[0]
|
|
spec *= p.spec.pulse_attn[:, None]
|
|
|
|
save_spectrum_data(timesignal, spec, p, dist, motion, t_echo_strings)
|
|
|
|
fig2, ax2 = plt.subplots()
|
|
lines = ax2.semilogx(p.dist.variables['tau'], reduction_factor / num_traj, 'o--')
|
|
ax2.legend(lines, t_echo_strings)
|
|
plt.savefig(f'{dist.name}_{motion.name}_reduction.png')
|
|
|
|
plt.show()
|
|
|
|
|
|
def make_trajectory(
|
|
dist: BaseDistribution,
|
|
motion: BaseMotion,
|
|
t_max: float,
|
|
t_passed: float = 0.,
|
|
init_phase: float = 0.
|
|
):
|
|
|
|
# set initial orientations and correlation times
|
|
motion.start()
|
|
dist.start()
|
|
|
|
# number of jumps that are simulated at once
|
|
chunks = min(int(0.51 * t_max / dist.tau_jump), 100_000) + 1
|
|
# print(chunks)
|
|
|
|
t = [np.array([t_passed])]
|
|
phase = [np.array([init_phase])]
|
|
# omega = [np.array([0])]
|
|
while t_passed < t_max:
|
|
# frequencies between jumps
|
|
current_omega = motion.jump(size=chunks)
|
|
# times at a particular position
|
|
t_wait = dist.wait(size=chunks)
|
|
|
|
accumulated_phase = np.cumsum(t_wait * current_omega) + phase[-1][-1]
|
|
phase.append(accumulated_phase)
|
|
# omega.append(current_omega)
|
|
|
|
t_wait = np.cumsum(t_wait) + t_passed
|
|
t_passed = t_wait[-1]
|
|
t.append(t_wait)
|
|
|
|
t = np.concatenate(t)
|
|
phase = np.concatenate(phase)
|
|
# omega = np.concatenate(omega)
|
|
|
|
# fig_test, ax_test = plt.subplots()
|
|
# ax_test.plot(t, phase, 'x-')
|
|
|
|
# np.savetxt('trajectory.dat', np.c_[t, phase, omega])
|
|
|
|
# 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]:
|
|
# collect variables that are common to spectra and stimulated echo
|
|
|
|
# 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
|
|
|
|
|
|
def print_step(n: int, num_traj: int, start: float, last_print: float) -> float:
|
|
step_time = perf_counter()
|
|
dt = step_time - last_print
|
|
if dt > 10 or n == num_traj:
|
|
date = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
|
|
print(f'{date} - step {n} of {num_traj}', end=' - ')
|
|
|
|
elapsed = step_time - start
|
|
total = num_traj * elapsed / n
|
|
print(f'expected total: {total:.2f}s - elapsed: {elapsed:.2f}s - remaining: {total - elapsed:.2f}s')
|
|
if dt > 10:
|
|
last_print = step_time
|
|
|
|
return last_print
|
|
|
|
|
|
def make_filename(dist: BaseDistribution, motion: BaseMotion) -> str:
|
|
filename = f'{dist}_{motion}'
|
|
filename = filename.replace(' ', '_')
|
|
filename = filename.replace('.', 'p')
|
|
|
|
return filename
|
|
|
|
|
|
def save_spectrum_data(
|
|
timesignal: np.ndarray,
|
|
spectrum: np.ndarray,
|
|
param: Parameter,
|
|
dist: BaseDistribution,
|
|
motion: BaseMotion,
|
|
echo_strings: list[str]
|
|
):
|
|
filename = make_filename(dist, motion)
|
|
|
|
header = param.totext(sim=True, spec=True)
|
|
header += '\nx\t' + '\t'.join(echo_strings)
|
|
|
|
np.savetxt(filename + '_timesignal.dat', np.c_[param.spec.t_acq, timesignal], header=header)
|
|
np.savetxt(filename + '_spectrum.dat', np.c_[param.spec.freq, spectrum], header=header)
|
|
|
|
fig, ax = plt.subplots()
|
|
lines = ax.plot(param.spec.freq, spectrum)
|
|
ax.set_title(f'{dist}, {motion}')
|
|
ax.legend(lines, echo_strings)
|
|
plt.savefig(filename + '_spectrum.png')
|
|
|
|
fig1, ax1 = plt.subplots()
|
|
lines = ax1.plot(param.spec.t_acq, timesignal)
|
|
ax1.set_title(f'{dist}, {motion}')
|
|
ax1.legend(lines, echo_strings)
|
|
plt.savefig(filename + '_timesignal.png')
|
|
|
|
plt.show()
|