python/rwsims/tetrahedral_spectrum.py

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()