103 lines
3.3 KiB
Python
103 lines
3.3 KiB
Python
|
import pathlib
|
||
|
|
|
||
|
|
import numpy
|
||
|
|
import numpy as np
|
||
|
|
from matplotlib import pyplot as plt
|
||
|
|
from scipy.optimize import curve_fit
|
||
|
|
|
||
|
|
from python.helpers import read_parameter_file
|
||
|
|
|
||
|
|
|
||
|
|
def ste_decay(x: np.ndarray, m0: float, t: float, beta: float, finfty: float) -> np.ndarray:
|
||
|
|
"""
|
||
|
|
Calculate stimulated-echo decay.
|
||
|
|
|
||
|
|
:param x: Mixing times in seconds.
|
||
|
|
:param m0: Amplitude
|
||
|
|
:param t: Correlation time in seconds
|
||
|
|
:param beta: Stretching parameter
|
||
|
|
:param finfty: Final plateau
|
||
|
|
:return: Stimulated-echo decay
|
||
|
|
"""
|
||
|
|
return m0 * ((1-finfty) * np.exp(-(x/t)**beta) + finfty)
|
||
|
|
|
||
|
|
|
||
|
|
def fit_decay(x: np.ndarray, y: np.ndarray, tevo: np.ndarray, verbose: bool = True) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
||
|
|
num_evo = y.shape[1]
|
||
|
|
|
||
|
|
tau_fit = np.empty((num_evo, 2))
|
||
|
|
tau_fit[:, 0] = tevo
|
||
|
|
|
||
|
|
beta_fit = np.empty((num_evo, 2))
|
||
|
|
beta_fit[:, 0] = tevo
|
||
|
|
|
||
|
|
finfty_fit = np.empty((num_evo, 2))
|
||
|
|
finfty_fit[:, 0] = tevo
|
||
|
|
|
||
|
|
scaled_y = (y-y[-1, :]) / (y[0, :]-y[-1, :])
|
||
|
|
|
||
|
|
for j in range(num_evo):
|
||
|
|
p0 = [scaled_y[0, 1], x[np.argmin(np.abs(scaled_y[:, j]-np.exp(-1)))], 0.5, 0.1]
|
||
|
|
|
||
|
|
try:
|
||
|
|
res = curve_fit(ste_decay, x, y[:, j], p0, bounds=[(0, 0, 0., 0), (np.inf, np.inf, 1, 1)])
|
||
|
|
except RuntimeError as e:
|
||
|
|
print(f'Fit {j+1} of {num_evo} failed with {e.args}')
|
||
|
|
continue
|
||
|
|
|
||
|
|
m0, tauc, beta, finfty = res[0]
|
||
|
|
|
||
|
|
if verbose:
|
||
|
|
print(f'Fit {j+1} of {num_evo}: tau_c = {tauc:.6e}, beta={beta:.4e}, amplitude = {m0: .4e}, f_infty={finfty:.4e}')
|
||
|
|
|
||
|
|
tau_fit[j, 1] = tauc
|
||
|
|
beta_fit[j, 1] = beta
|
||
|
|
finfty_fit[j, 1] = finfty
|
||
|
|
|
||
|
|
return tau_fit, beta_fit, finfty_fit
|
||
|
|
|
||
|
|
|
||
|
|
def fit_and_save_ste(parameter_file: pathlib.Path, plot_decays: bool = True, verbose: bool = True):
|
||
|
|
# read simulation parameters
|
||
|
|
parameter = read_parameter_file(parameter_file)
|
||
|
|
|
||
|
|
# files have form ste_arg=0.000000e+01_parameter.txt, first remove ste part then parameter.txt to get variables
|
||
|
|
varied_string = str(parameter_file).partition('_')[-1].rpartition('_')[0]
|
||
|
|
|
||
|
|
# make evolution times
|
||
|
|
tevo = np.linspace(parameter['tevo_start'], parameter['tevo_stop'], num=int(parameter['tevo_steps']))
|
||
|
|
|
||
|
|
raw_data_cc = np.loadtxt(f'coscos_{varied_string}.dat')
|
||
|
|
raw_data_ss = np.loadtxt(f'sinsin_{varied_string}.dat')
|
||
|
|
|
||
|
|
t_mix = raw_data_cc[:, 0]
|
||
|
|
cc_decay = raw_data_cc[:, 1:]
|
||
|
|
ss_decay = raw_data_ss[:, 1:]
|
||
|
|
|
||
|
|
if plot_decays:
|
||
|
|
fig_cc, ax_cc = plt.subplots()
|
||
|
|
ax_cc.set_title('Cos-Cos')
|
||
|
|
ax_cc.semilogx(t_mix, cc_decay, '.')
|
||
|
|
|
||
|
|
fig_ss, ax_ss = plt.subplots()
|
||
|
|
ax_ss.set_title('Sin-Sin')
|
||
|
|
ax_ss.semilogx(t_mix, ss_decay, '.')
|
||
|
|
|
||
|
|
plt.show()
|
||
|
|
|
||
|
|
print('Fit Cos-Cos')
|
||
|
|
tau_cc, beta_cc, finfty_cc = fit_decay(t_mix, cc_decay, tevo, verbose=verbose)
|
||
|
|
|
||
|
|
np.savetxt(f'tau_cc_{varied_string}.dat', tau_cc)
|
||
|
|
np.savetxt(f'beta_cc_{varied_string}.dat', beta_cc)
|
||
|
|
np.savetxt(f'finfty_cc_{varied_string}.dat', finfty_cc)
|
||
|
|
|
||
|
|
print('Fit Sin-Sin')
|
||
|
|
tau_ss, beta_ss, finfty_ss = fit_decay(t_mix, ss_decay, tevo, verbose=verbose)
|
||
|
|
|
||
|
|
np.savetxt(f'tau_ss_{varied_string}.dat', tau_ss)
|
||
|
|
np.savetxt(f'beta_ss_{varied_string}.dat', beta_ss)
|
||
|
|
np.savetxt(f'finfty_ss_{varied_string}.dat', finfty_ss)
|
||
|
|
|
||
|
|
return varied_string, tau_cc, beta_cc, finfty_cc, tau_ss, beta_ss, finfty_ss
|