Example.py

import numpy as np
import matplotlib.pyplot as plt
import Simulation_ABP
import Simulation_AOUP
import Simulation_RTP
import Exiting_Statistics_toolkit as exit
import Fitting_NMSD_toolkit as fitting
import pandas as pd
import ABP_RT_hydro_fast
import ABP_RT_hydro_expanded_fast
import ABP_RT_Order3_fast
import Compute_NMSD_from_Continued_Fractions as CNMSD
from scipy.special import erf

import countoscope #Import the main counting code of the countoscope library which can be found at https://github.com/Countoscope/countoscope_PRX

""" 
Example of the use of the codes in this repository.

It consists in one simulation stage, one counting stage, one fitting stage, one exit probability computation stage, and one fitting stage.
"""

saving_path = "Results/Example/" # path to save the results of the example

def f_function_AOUPs(t, L, Dt, Dv, tau):
    MSD = 4 * (Dt + Dv) * t + 4 * Dv * tau * (np.exp(-t / tau) - 1)
    return np.sqrt(MSD / (np.pi*L**2) )*(np.exp(-L**2/MSD) - 1) + erf(L/(np.sqrt(MSD)))

# Simulation stage
# The provided codes generate simulations of ABPs, RTPs, or AOUPs

model = "RTP" # choose the model to simulate: "ABP", "RTP" or "AOUP"

COLORMAPS = {"ABP": "PuBu", "RTP": "YlOrRd", "AOUP": "RdPu"}
def colors(cmap, n):
    return plt.get_cmap(cmap)(np.linspace(0.35, 0.95, n))

# Simulation parameters
##Physical parameters
v = 5 #particules speed in μm/s
Dt = 0.1 #translational diffusion coefficient in μm**2/s**2
Dr, alpha = (0, 1) if model == "RTP" else (1, 0) # Rotational diffusion coefficient in μm**2/s**2 and tumbling rate in 1/s. 
#For ABPs, Dr is the rotational diffusion coefficient, for RTPs, alpha is the tumbling rate, and for AOUPs, Dr = 1/tau where tau is the persistence time.

if model == "AOUP":
    Dv = v**2 / (2*Dr) #active diffusion coefficient in μm**2/s
    tau = 1/Dr #persistence time in seconds

Dimension = 2 # number of dimensions of the system, this study is for now limited to 2D

LBox_x = 250 # size of one side of the simulation box in μm
LBox_y = 250 # size of one side of the simulation box in μm
LBox_z = 250 # size of one side of the simulation box in μm

N0 = 1500 #number of particles
dt= 0.01 #in seconds
T = 1000 #total simulation time in seconds
Nsave = 5 # interval for saving positions, we save the positions every Nsave time steps

if model == "ABP" :
    Result_values = Simulation_ABP.simulation_ABP(v, Dt, Dr, Dimension, LBox_x, LBox_y, LBox_z, N0, dt, T, Nsave)
elif model == "RTP" :
    Result_values = Simulation_RTP.simulation_RTP(v, Dt, alpha, Dimension, LBox_x, LBox_y, LBox_z, N0, dt, T, Nsave)
elif model == "AOUP" :
    Result_values = Simulation_AOUP.simulation_AOUP(Dt, Dv, tau, Dimension, LBox_x, LBox_y, LBox_z, N0, dt, T, Nsave)



#optional: save the results of the simulation
savename = f'Simulation_{model}_{Dimension}D_Dr-{Dr}_Dt-{Dt}_v-{v}_Time-{T}_Interval-{round(dt*Nsave, 3)}'
np.save(saving_path + f"{savename}.npy", Result_values)

# Counting stage
# We use the Countoscope library to count the particles in the simulation box

Box_Ls = np.logspace(0.5, 1.7 ,10)
print(Box_Ls)
a = 1.0
sep = [2*a, 2*a, 2*a, -a,-a, -2*a, -2*a, -4*a, -7*a, -15*a]
data = Result_values
data.reshape(-1, 4) # Reshape the data to have the correct number of columns x, y, frame number, particle number (the counting is for now limited to 2D)
results = countoscope.calculate_nmsd(data=data[:,0:3],
                                                                    window_size_x=LBox_x, window_size_y=LBox_y, 
                                                                    box_sizes=Box_Ls, sep_sizes=sep)
N2_mean = results.nmsd
N2_std = results.nmsd_std
N_mean = results.N_mean
N_var = results.N_var_mod
Num_boxes = results.num_boxes

# Optional plot the NMSD alongside the theory for the model used in the simulation
fig, ax = plt.subplots()
frame_dt = dt * Nsave # time between frames in seconds
nt = int(T/frame_dt) # number of frames
t = [it*frame_dt for it in range(1, nt)] # time in seconds, the point at t=0 messes up the log-log plot, so we start at t=dt
# For RTP and ABP the theory can be computed at any Order in the angular moments expansion
# Orders above 3 require solving the Coninued Fractions and inverse Fourier transform it to obtain the ISF, which is time consuming
Order = 2 # Chosen Order for the theory, None returns the exact result for RTPs

# Inverse Fourier transform parameters, lowering N_max increases speed but decreases precision, especially at short times.
n_k = 800
N_min = 2**14
N_max = 2**20

cmap = colors(COLORMAPS[model], len(Box_Ls))
t_theory = np.logspace(np.log10(frame_dt), np.log10(T), num=500) # time in seconds for the theory, we can use a different time grid than the simulation

for i, L in enumerate(Box_Ls):
    print(f"Computing the theory for L = {L}, Order = {Order}, model = {model}")
    ax.plot(t, N2_mean[i,1:], color = cmap[i], linestyle = 'None', marker = 'o')
    Npart = N0/LBox_x/LBox_y*L**2 # mean number of particles in a box of size L, for the simulation parameters
    if model == "ABP" or model == "RTP" :
        if Order == 1 :
            Ntplot, Pval = ABP_RT_hydro_fast.sDFT_ABP_RT(L, t_theory, Dt, Dr, v, alpha)
        elif Order == 2 :
            Ntplot, Pval = ABP_RT_hydro_expanded_fast.sDFT_ABP_RT_expanded(L, t_theory, Dt, Dr, v, alpha)
        elif Order == 3 :
            Ntplot, Pval = ABP_RT_Order3_fast.Compute_NMSD(L, t_theory, Dt, Dr, v, alpha)
        else : 
            Ntplot, Pval = CNMSD.nmsd_from_Fkt(L, t_theory, Dt, alpha, Dr, v, model=model, order=Order, force_recompute=False, use_universal_cache=True, precision_factor=1, N_min=N_min, N_max=N_max, n_k=n_k)

    elif model == "AOUP" :
        Ntplot, Pval = f_function_AOUPs(t_theory, L, Dt, Dv, tau)**2 , 1

    ax.plot(t_theory, 2*Npart*(Pval-Ntplot), label = f'L = {np.round(L, 1)}' + ' \N{MICRO SIGN}m', color = cmap[i])
ax.set_xlabel("Lag time (s)")
ax.set_ylabel("NMSD")
ax.set_xscale("log")
ax.set_yscale("log")
ax.legend()
ax.grid()
plt.tight_layout()

plt.show()

# optional: save the results of the counting
np.savez(saving_path + f'{savename}_counted', N2_mean=N2_mean, N2_std=N2_std, N_mean=N_mean, N_var = N_var, box_sizes=Box_Ls, sep_sizes=sep, Num_boxes=Num_boxes)

# Computing the probability stage
# This section computes the probability of particles to leave a box in time, and to return to it after leaving it, as a function of lag time.
# Computing the probabilities is time consuming, so it is optional. The results can be saved and reloaded for plotting and fitting.
compute_exit_probabilities = False
if compute_exit_probabilities:
    t= pd.DataFrame({'x' : data[:,0], 'y' : data[:,1], 'frame' : data[:,2].astype(int), 'particle' : data[:,3].astype(int)})
    del data

    max_lag = 30000   # number of frames to evaluate (adjust)
    # lags = np.arange(0, max_lag+1)   
    # If your frames correspond to real time dt (e.g. 0.01s per frame), pass frame_dt=0.01
    LBox = 5 # box size to probe the probabilities for
    Num_Boxes = 30
    df_metrics_boxes = exit.compute_probabilities_boxes_fast(t, LBox = LBox, Num_Boxes = Num_Boxes, max_lag = max_lag, num_points = 3000, frames_dt= dt * Nsave)

    # optional: save the results of the probability computation to a pickle file, to reload it later for plotting and fitting
    df_metrics_boxes.to_pickle(saving_path + f'{savename}_probabilities.pkl')


# fitting stage
# This section fits the computed NMSD to the theoretical models developped in the papers
# it uses the limiting regimes, as well as the Order 3 theoretical model for the NMSD of ABPs and RTPs
# and the full analyical model for the NMSD of AOUPS

# Mean number of particles per box, from the total number of particles rather than from the counts,
# whose time average converges slowly (Mackay et al., PRX 14, 041016 (2024), Appendix C)
N_mean_fit = N0 * Box_Ls**2 / (LBox_x * LBox_y)

# The NMSD of each box size is resampled on log-spaced lag times, and each box size is then fitted on its own
dataset = fitting.build_dataset(Box_Ls, N2_mean, N2_std, Num_boxes, N_mean_fit, frames_dt = dt * Nsave)


# 1. Limiting regimes, box by box (ABPs and RTPs): v from the advective regime, Deff from the long-time
# diffusive regime (t > k_diff * t_diff), Dt from the short-time diffusive regime (only if the frame rate
# resolves it), and Dr = v^2 / (2 (Deff - Dt)). No prior knowledge of the parameters is needed.
# k_diff should be high enough to be in the diffusive regime, and low enough to keep data before the plateau.
# Add Dt_known = ... if Dt was measured on passive particles.
# AOUPs skip this step: their exact NMSD is simple enough to be fitted directly on the whole curve.
k_diff = 3
if model == "AOUP":
    Limits = None
    p0 = (0.1, 1.0, 5.0) # rough guess of (Dt, 1/tau, v) to start the full fit, the order of magnitude is enough
else:
    Limits = fitting.fit_limiting_regimes(dataset, model = model, k_diff = k_diff)
    print(Limits.to_string())
    p0 = Limits[["Dt", "Dr", "v"]].median() # the full fit starts from the median over box sizes

# 2. Full model fitted box by box. Dr is the tumbling rate for RTPs and 1/tau for AOUPs.
# The plateau 2<N> is measured rather than fitted, so only the rise of the NMSD is kept (plateau_frac).
# Add Dt_fixed = ... to hold Dt at a value measured on passive particles.
dataset_rise = fitting.build_dataset(Box_Ls, N2_mean, N2_std, Num_boxes, N_mean_fit, frames_dt = dt * Nsave, plateau_frac = 0.9)
Fits = fitting.fit_nmsd(dataset_rise, model = model, p0 = p0)
print(Fits.to_string())

# optional: check the fits visually, then look at the fitted parameters as a function of box size.
# truth: the simulation input, to compare with (drop it for experimental data).
# For AOUPs simulated with Simulation_AOUP in 2D: truth = dict(v = np.sqrt(2*Dv/tau), Dr = 1/tau, Dt = Dt)
if Limits is not None:
    ax = fitting.plot_limiting_regimes(dataset, Limits)
    ax.figure.savefig(saving_path + f'{savename}_limiting_regimes.png', dpi = 300)
    plt.show()
ax = fitting.plot_fit(dataset, Fits)
plt.show()
ax.figure.savefig(saving_path + f'{savename}_full_model_fits.png', dpi = 300)
axes = fitting.plot_parameters(Limits, Fits, truth = dict(v = v, Dr = alpha if model == "RTP" else Dr, Dt = Dt))
plt.show()
# axes[0].figure.savefig(saving_path + f'{savename}_parameters.png', dpi = 300)

# optional: save the results of the fitting, one row per box size
# if Limits is not None:
    # Limits.to_csv(saving_path + f'{savename}_limiting_regimes.csv', index = False)
# Fits.to_csv(saving_path + f'{savename}_full_model_fits.csv', index = False)