1D Double Well estimation

import numpy as np
import matplotlib.pyplot as plt
import folie as fl
import csv
from mpl_toolkits.mplot3d import Axes3D
from copy import deepcopy

1D UNBIASED Double Well Potential

1) The model

Here we model the double well potential as a quartic function \(V(q)= \sum_{i=0}^4 c_iq^i\) and choose a constant diffusion coefficient \(D(q)=q\) :

The force parameter to pass to the simulator will then be : \(F = - \frac{dV(q)}{dq}\)

coeff=0.1*np.array([0,0,-4.5,0,0.1]) # coefficients of the free energy
free_energy = np.polynomial.Polynomial(coeff)
force_coeff=np.array([-coeff[1],-2*coeff[2],-3*coeff[3],-4*coeff[4]]) #coefficients of the free energy
force_function = fl.functions.Polynomial(deg=3,coefficients=force_coeff)
diff_function= fl.functions.Polynomial(deg=0,coefficients=np.asarray([0.5]))

# Plot of Free Energy and Force
x_values = np.linspace(-7, 7, 100)
fig, axs = plt.subplots(1, 2)
axs[0].plot(x_values,free_energy(x_values))
axs[1].plot(x_values,force_function(x_values.reshape(len(x_values),1)))
axs[0].set_title("Potential")
axs[0].set_xlabel("$x$")
axs[0].set_ylabel("$V(x)$")
axs[0].grid()
axs[1].set_title("Force")
axs[1].set_xlabel("$x$")
axs[1].set_ylabel("$F(x)$")
axs[1].grid()

2) Simulation

# Define model to simulate and type of simulator to use
dt=1e-3
model_simu = fl.models.overdamped.Overdamped(force_function,diffusion=diff_function)
simulator = fl.simulations.Simulator(fl.simulations.EulerStepper(model_simu), dt)

ntraj=30
q0= np.empty(ntraj)
for i in range(len(q0)):
    q0[i]=0
# Calculate Trajectory
time_steps=10000
data = simulator.run(time_steps, q0, 1)

# Plot the trajecories
fig, axs = plt.subplots(1,1)
for n, trj in enumerate(data):
    axs.plot(trj["x"])
    axs.set_title("Trajectory")
    axs.set_xlabel("$timestep$")

3) Model Training

3.1) Training using same functional form of true force and diffusion

# Parameters of the training

trainforce =fl.functions.Polynomial(deg=3,coefficients=np.asarray([1,1,1,1]))
traindiff = fl.functions.Polynomial(deg=0,coefficients=np.asarray([0.0]))
trainmodel=fl.models.Overdamped(force = deepcopy(trainforce),diffusion=deepcopy(traindiff), has_bias=None)

KM_estimator = fl.KramersMoyalEstimator(deepcopy(trainmodel), n_jobs=4)
Eul_estimator = fl.LikelihoodEstimator(fl.EulerDensity(deepcopy(trainmodel)),n_jobs=4)
Eln_estimator = fl.LikelihoodEstimator(fl.ElerianDensity(deepcopy(trainmodel)),n_jobs=4)
Ksl_estimator = fl.LikelihoodEstimator(fl.KesslerDensity(deepcopy(trainmodel)),n_jobs=4)
Drz_estimator = fl.LikelihoodEstimator(fl.DrozdovDensity(deepcopy(trainmodel)),n_jobs=4)

KM_res=KM_estimator.fit_fetch(data)
Eul_res=Eul_estimator.fit_fetch(data)
Eln_res=Eln_estimator.fit_fetch(data)
Ksl_res=Ksl_estimator.fit_fetch(data)
Drz_res=Drz_estimator.fit_fetch(data)

res_vec = [KM_res,Eul_res,Eln_res,Ksl_res,Drz_res] # made a list of all the trained estimators

fig, axs = plt.subplots(1, 2, figsize=(14, 8))
axs[0].set_title("Force")
axs[0].set_xlabel("$x$")
axs[0].set_ylabel("$F(x)$")
axs[0].grid()

axs[1].set_title("Diffusion Coefficient")
axs[1].set_xlabel("$x$")
axs[1].set_ylabel("$D(x)$")
axs[1].grid()
xfa = np.linspace(-7.0, 7.0, 75)

#Plot exact quantities

axs[0].plot(xfa, model_simu.pos_drift(xfa.reshape(-1, 1)), label="Exact")
axs[1].plot(xfa, model_simu.diffusion(xfa.reshape(-1, 1)), label="Exact")
# #Plot inferred quantities

names = ["KM","Euler", "Elerian", "Kessler", "Drozdov"]
markers = ["x", "1","2","3","|"]
for i in range(len(names)):
    res_vec[i].remove_bias()
    axs[0].plot(xfa, res_vec[i].pos_drift(xfa.reshape(-1, 1)), markers[i],label=names[i] )
    axs[1].plot(xfa, res_vec[i].diffusion(xfa.reshape(-1, 1)), markers[i],label=names[i])
    print(names[i],res_vec[i].coefficients)
axs[0].legend()
axs[1].legend()

3.2) Training using splines

# Parameters of the training
n_knots= 5
domain = fl.MeshedDomain.create_from_range(np.linspace(data.stats.min , data.stats.max , n_knots).ravel())
trainmodel = fl.models.Overdamped(force=fl.functions.BSplinesFunction(domain), diffusion = fl.functions.Polynomial(deg=0,coefficients=np.asarray([0.0])), has_bias=None)

KM_estimator = fl.KramersMoyalEstimator(deepcopy(trainmodel), n_jobs=4)
Eul_estimator = fl.LikelihoodEstimator(fl.EulerDensity(deepcopy(trainmodel)),n_jobs=4)
Eln_estimator = fl.LikelihoodEstimator(fl.ElerianDensity(deepcopy(trainmodel)),n_jobs=4)
Ksl_estimator = fl.LikelihoodEstimator(fl.KesslerDensity(deepcopy(trainmodel)),n_jobs=4)
Drz_estimator = fl.LikelihoodEstimator(fl.DrozdovDensity(deepcopy(trainmodel)),n_jobs=4)

KM_res=Eul_estimator.fit_fetch(data)
Eul_res=Eul_estimator.fit_fetch(data)
Eln_res=Eln_estimator.fit_fetch(data)
Ksl_res=Ksl_estimator.fit_fetch(data)
Drz_res=Drz_estimator.fit_fetch(data)

res_vec = [KM_res,Eul_res,Eln_res,Ksl_res,Drz_res] # made a list of all the trained estimators

# Plot of OVerdamped( spline, constant)
fig, axs = plt.subplots(1, 2, figsize=(10, 6))
axs[0].set_title("Force")
axs[0].set_xlabel("$x$")
axs[0].set_ylabel("$F(x)$")
axs[0].grid()

axs[1].set_title("Diffusion Coefficient")
axs[1].set_xlabel("$x$")
axs[1].set_ylabel("$D(x)$")
axs[1].grid()
fig.suptitle('B-spline Fitting with '+str(n_knots)+ ' knots')

xfa = np.linspace(-7.0, 7.0, 75)

#Plot exact quantities

axs[0].plot(xfa, model_simu.pos_drift(xfa.reshape(-1, 1)), label="Exact")
axs[1].plot(xfa, model_simu.diffusion(xfa.reshape(-1, 1)), label="Exact")
#Plot inferred quantities
names = ["KM","Euler", "Elerian", "Kessler", "Drozdov"]
markers = ["x", "1","2","3","|"]
for i in range(len(names)):
    res_vec[i].remove_bias()
    axs[0].plot(xfa, res_vec[i].pos_driftriftriftrift(xfa.reshape(-1, 1)), markers[i],label=names[i] )
    axs[1].plot(xfa, res_vec[i].diffusion(xfa.reshape(-1, 1)), markers[i],label=names[i])
    print(names[i],res_vec[i].coefficients)
axs[0].legend()
axs[1].legend()

Check if the methods are returning all the the same results

1D BIASED Double Well Potential

1) Model

Here we model the double well potential as a quartic function \(V(q)= \sum_{i=0}^4 c_iq^i\) and choose a constant diffusion coefficient \(D(q)=D\) : \(\newline\) The force parameter to pass to the simulator will then be : \(F = - \frac{dV(q)}{dq}\) \(\newline\) Adiabaic bias used : \(V_{bias}(q)=\frac{1}{2}k(q-q_0)^2 \longmapsto\) ABMD_Simulator \(\newline\) The center of the parabola, \(q_0\), is choosen as : \(max(q,q_0)\) at every iteration

coeff=0.1*np.array([0,0,-4.5,0,0.1]) # coefficients of the free energy
free_energy = np.polynomial.Polynomial(coeff)
force_coeff=np.array([-coeff[1],-2*coeff[2],-3*coeff[3],-4*coeff[4]]) #coefficients of the free energy
force_function = fl.functions.Polynomial(deg=3,coefficients=force_coeff)
diff_function= fl.functions.Polynomial(deg=0,coefficients=np.asarray([0.5]))

# Plot of Free Energy and Force
x_values = np.linspace(-7, 7, 100)
fig, axs = plt.subplots(1, 2)
axs[0].plot(x_values,free_energy(x_values))
axs[1].plot(x_values,force_function(x_values.reshape(len(x_values),1)))
axs[0].set_title("Potential")
axs[0].set_xlabel("$x$")
axs[0].set_ylabel("$V(x)$")
axs[0].grid()
axs[1].set_title("Force")
axs[1].set_xlabel("$x$")
axs[1].set_ylabel("$F(x)$")
axs[1].grid()

2) Simulation

# Define model to simulate and type of simulator to use
dt=1e-3
biased_model_simu = fl.models.overdamped.Overdamped(force_function,diffusion=diff_function)
biased_simulator = fl.simulations.ABMD_Simulator(fl.simulations.EulerStepper(biased_model_simu), dt, k=10.0, xstop=6.0)
ntraj=30
q0= np.empty(ntraj)
for i in range(len(q0)):
    q0[i]=-6.0
# Calculate Trajectory
time_steps=35000
biased_data = biased_simulator.run(time_steps, q0, 1)
xmax = np.concatenate(biased_simulator.xmax_hist, axis=1).T    # if you rerun simulator.run without reinializing the simulator object it will probably append the results making xmax twice as long

# Plot the trajecories
fig, axs = plt.subplots(1,2,figsize=(10,6))
for n, trj in enumerate(biased_data):
    axs[0].plot(trj["x"])
    axs[0].set_title("Trajectory")
    axs[0].set_xlabel("$timestep$")
    axs[0].set_ylabel("q(t)")
    axs[1].plot(xmax)
    axs[1].set_title("q_0")
    axs[1].set_xlabel("$timesteps$")

3) Model Training

3.1) Training using same functional form of true force and diffusion

# Parameters of the training

trainforce =fl.functions.Polynomial(deg=3,coefficients=np.asarray([1,1,1,1]))
traindiff = fl.functions.Polynomial(deg=0,coefficients=np.asarray([0.0]))
trainmodel=fl.models.Overdamped(force = deepcopy(trainforce),diffusion=deepcopy(traindiff), has_bias=True)

KM_estimator = fl.KramersMoyalEstimator(deepcopy(trainmodel), n_jobs=4)
Eul_estimator = fl.LikelihoodEstimator(fl.EulerDensity(deepcopy(trainmodel)),n_jobs=4)
Eln_estimator = fl.LikelihoodEstimator(fl.ElerianDensity(deepcopy(trainmodel)),n_jobs=4)
Ksl_estimator = fl.LikelihoodEstimator(fl.KesslerDensity(deepcopy(trainmodel)),n_jobs=4)
Drz_estimator = fl.LikelihoodEstimator(fl.DrozdovDensity(deepcopy(trainmodel)),n_jobs=4)

KM_res=KM_estimator.fit_fetch(biased_data)
Eul_res=Eul_estimator.fit_fetch(biased_data)
Eln_res=Eln_estimator.fit_fetch(biased_data)
Ksl_res=Ksl_estimator.fit_fetch(biased_data)
Drz_res=Drz_estimator.fit_fetch(biased_data)

res_vec = [KM_res,Eul_res,Eln_res,Ksl_res,Drz_res] # made a list of all the trained estimators

fig, axs = plt.subplots(1, 2, figsize=(14, 8))
axs[0].set_title("Force")
axs[0].set_xlabel("$x$")
axs[0].set_ylabel("$F(x)$")
axs[0].grid()

axs[1].set_title("Diffusion Coefficient")
axs[1].set_xlabel("$x$")
axs[1].set_ylabel("$D(x)$")
axs[1].grid()
xfa = np.linspace(-7.0, 7.0, 75)

#Plot exact quantities
biased_model_simu.remove_bias()
axs[0].plot(xfa, biased_model_simu.pos_drift(xfa.reshape(-1, 1)), label="Exact")
axs[1].plot(xfa, biased_model_simu.diffusion(xfa.reshape(-1, 1)), label="Exact")
# #Plot inferred quantities

names = ["KM","Euler", "Elerian", "Kessler", "Drozdov"]
markers = ["x", "1","2","3","|"]
for i in range(len(names)):
    res_vec[i].remove_bias()
    axs[0].plot(xfa, res_vec[i].pos_driftriftriftrift(xfa.reshape(-1, 1)), markers[i],label=names[i] )
    axs[1].plot(xfa, res_vec[i].diffusion(xfa.reshape(-1, 1)), markers[i],label=names[i])
    print(names[i],res_vec[i].coefficients)
axs[0].legend()
axs[1].legend()

3.2) Training using splines

# Parameters of the training
n_knots= 5
domain = fl.MeshedDomain.create_from_range(np.linspace(biased_data.stats.min , biased_data.stats.max , n_knots).ravel())
trainmodel = fl.models.Overdamped(force=fl.functions.BSplinesFunction(domain), diffusion = fl.functions.Polynomial(deg=0,coefficients=np.asarray([0.0])), has_bias=True)

KM_estimator = fl.KramersMoyalEstimator(deepcopy(trainmodel), n_jobs=4)
Eul_estimator = fl.LikelihoodEstimator(fl.EulerDensity(deepcopy(trainmodel)),n_jobs=4)
Eln_estimator = fl.LikelihoodEstimator(fl.ElerianDensity(deepcopy(trainmodel)),n_jobs=4)
Ksl_estimator = fl.LikelihoodEstimator(fl.KesslerDensity(deepcopy(trainmodel)),n_jobs=4)
Drz_estimator = fl.LikelihoodEstimator(fl.DrozdovDensity(deepcopy(trainmodel)),n_jobs=4)

KM_res=Eul_estimator.fit_fetch(biased_data)
Eul_res=Eul_estimator.fit_fetch(biased_data)
Eln_res=Eln_estimator.fit_fetch(biased_data)
Ksl_res=Ksl_estimator.fit_fetch(biased_data)
Drz_res=Drz_estimator.fit_fetch(biased_data)

res_vec = [KM_res,Eul_res,Eln_res,Ksl_res,Drz_res] # made a list of all the trained estimators

# Plot of OVerdamped( spline, constant)
fig, axs = plt.subplots(1, 2, figsize=(10, 6))
axs[0].set_title("Force")
axs[0].set_xlabel("$x$")
axs[0].set_ylabel("$F(x)$")
axs[0].grid()

axs[1].set_title("Diffusion Coefficient")
axs[1].set_xlabel("$x$")
axs[1].set_ylabel("$D(x)$")
axs[1].grid()
fig.suptitle('B-spline Fitting with '+str(n_knots)+ ' knots')

xfa = np.linspace(-7.0, 7.0, 75)

#Plot exact quantities
biased_model_simu.remove_bias()
axs[0].plot(xfa, biased_model_simu.drift(xfa.reshape(-1, 1)), label="Exact")
axs[1].plot(xfa, biased_model_simu.diffusion(xfa.reshape(-1, 1)), label="Exact")
#Plot inferred quantities
names = ["KM","Euler", "Elerian", "Kessler", "Drozdov"]
markers = ["x", "1","2","3","|"]
for i in range(len(names)):
    res_vec[i].remove_bias()
    axs[0].plot(xfa, res_vec[i].drift(xfa.reshape(-1, 1)), markers[i],label=names[i] )
    axs[1].plot(xfa, res_vec[i].diffusion(xfa.reshape(-1, 1)), markers[i],label=names[i])
    print(names[i],res_vec[i].coefficients)
axs[0].legend()
axs[1].legend()

with zip

fig, axb = plt.subplots(1, 2, figsize=(10, 6))

n_knots= 4
domain = fl.MeshedDomain.create_from_range(np.linspace(biased_data.stats.min , biased_data.stats.max , n_knots).ravel())
bias_spline_trainmodel = fl.models.Overdamped(fl.functions.BSplinesFunction(domain), fl.functions.Polynomial(deg=0,coefficients=np.asarray([0.0])), has_bias=True)

name = "KM"
estimator = fl.KramersMoyalEstimator(deepcopy(bias_spline_trainmodel))
res = estimator.fit_fetch(biased_data)
print('has bias True',name,res.coefficients)
axb[0].plot(xfa, res.drift(xfa.reshape(-1, 1)), label=name)
axb[1].plot(xfa, res.diffusion(xfa.reshape(-1, 1)), label=name)


for name, marker, transitioncls in zip(
   ["Euler", "Elerian", "Kessler", "Drozdov"],
    ["+","1","2","3","|","x"],
   [
       fl.EulerDensity,
       fl.ElerianDensity,
       fl.KesslerDensity,
       fl.DrozdovDensity,
   ],
):
    estimator = fl.LikelihoodEstimator(transitioncls(deepcopy(bias_spline_trainmodel)), n_jobs=4)
    res = estimator.fit_fetch(biased_data)

    axb[0].plot(xfa, res.drift(xfa.reshape(-1, 1)),marker, label=name)
    axb[1].plot(xfa, res.diffusion(xfa.reshape(-1, 1)),marker, label=name)
    print('has bias true',name,res.coefficients)

axb[0].legend()
axb[1].legend()
plt.show()