cpp/python/ste.py
2024-12-01 19:24:19 +01:00

93 lines
2.7 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]
if num_evo != tevo.size:
tevo = np.arange(num_evo)
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 = [y[0, 0], x[np.argmin(np.abs(scaled_y[:, j]-np.exp(-1)))], 0.8, 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_ste(
parameter_file: pathlib.Path,
prefix: str,
plot_decays: bool = True,
verbose: bool = True,
) -> tuple[str, np.ndarray, np.ndarray, np.ndarray]:
# 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 = np.loadtxt(parameter_file.with_name(f'{prefix}_{varied_string}.dat'))
t_mix = raw_data[:, 0]
decay = raw_data[:, 1:]
if plot_decays:
fig, ax = plt.subplots()
ax.set_title(prefix)
ax.semilogx(t_mix, decay, '.')
fig.show()
print(f'Fit {prefix}')
tau, beta, finfty = fit_decay(t_mix, decay, tevo, verbose=verbose)
return varied_string, tau, beta, finfty