166 lines
5.2 KiB
Python
166 lines
5.2 KiB
Python
from time import time
|
|
from numpy.typing import ArrayLike
|
|
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 = np.logspace(-8, -1, num=15) # in s
|
|
|
|
# acquisition parameter
|
|
acq_length = 4096
|
|
dt = 1e-6 # 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
|
|
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 = 1
|
|
|
|
|
|
def omega_q(delta_: float, eta_: float, theta_: ArrayLike, phi_: ArrayLike) -> 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 rotate(x_in: float, y_in: float, z_in: float, a: float) -> tuple[float, float]:
|
|
# rotation by tetrahedral angle is a given, only second angle is free parameter
|
|
beta = 109.45 * np.pi / 180.
|
|
cos_beta = np.cos(beta)
|
|
sin_beta = np.sin(beta)
|
|
|
|
cos_alpha = np.cos(a)
|
|
sin_alpha = np.sin(a)
|
|
|
|
scale = np.sqrt(1 - z_in * z_in) + 1e-12
|
|
|
|
x = x_in * cos_beta + sin_beta / scale * (x_in * z_in * cos_alpha - y_in * sin_alpha)
|
|
y = y_in * cos_beta + sin_beta / scale * (y_in * z_in * cos_alpha + x_in * sin_alpha)
|
|
z = z_in * cos_beta - scale * sin_beta*cos_alpha
|
|
z = max(-1, min(1, z))
|
|
|
|
return np.arccos(z), np.arctan2(y, x)
|
|
|
|
|
|
def new_tau(size=1) -> np.ndarray:
|
|
return rng.exponential(tau_i, size=size)
|
|
|
|
|
|
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)))
|
|
|
|
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):
|
|
# draw orientation
|
|
z_theta, z_phi, z_alpha = rng.random(3)
|
|
theta = np.arccos(1 - 2 * z_theta)
|
|
phi = 2 * np.pi * z_phi
|
|
alpha = 2 * np.pi * z_alpha
|
|
|
|
# orientation in cartesian coordinates
|
|
x_start = np.sin(theta) * np.cos(phi)
|
|
y_start = np.sin(theta) * np.sin(phi)
|
|
z_start = np.cos(theta)
|
|
|
|
# calculate orientation of tetrahedral edges and their frequencies
|
|
orientation = np.zeros(4)
|
|
orientation[0] = omega_q(delta, eta, theta, phi)
|
|
orientation[1] = omega_q(delta, eta, *rotate(x_start, y_start, z_start, alpha))
|
|
orientation[2] = omega_q(delta, eta, *rotate(x_start, y_start, z_start, alpha + 2 * np.pi / 3))
|
|
orientation[3] = omega_q(delta, eta, *rotate(x_start, y_start, z_start, alpha + 4 * np.pi / 3))
|
|
|
|
t_passed = 0
|
|
t = [0]
|
|
phase = [0]
|
|
accumulated_phase = 0
|
|
start_position = rng.choice([0, 1, 2, 3])
|
|
|
|
while t_passed < t_max:
|
|
# orientation until the next jump
|
|
jumps = rng.choice([1, 2, 3], size=chunks) + start_position
|
|
jumps = np.cumsum(jumps)
|
|
jumps %= 4
|
|
|
|
current_omega = orientation[jumps]
|
|
# current_omega = rng.choice(orientation, 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
|
|
reduction_factor[n, j] += np.cos(phase_interpol(2*t_echo_j) + start_amp)
|
|
|
|
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'tau = {tau_i}s')
|
|
ax.legend(lines, t_echo_strings)
|
|
# plt.savefig(f'{tau_i}.png')
|
|
|
|
# save time signals and spectra
|
|
# 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()
|