python/rwsims/tetrahedral_spectrum.py
2024-08-01 18:46:28 +02:00

146 lines
4.5 KiB
Python

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 = 10e3 # 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 = 10000
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 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)
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)
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
current_orientation = rng.choice([0, 1, 2, 3])
while t_passed < t_max:
# orientation until the next jump
# always jump to a different position i -> i + {1, 2, 3}
current_orientation += rng.choice([1, 2, 3])
current_orientation %= 4
current_omega = orientation[current_orientation]
# 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'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))
plt.show()