diff --git a/examples/plot_free_energy.py b/examples/plot_free_energy.py new file mode 100644 index 0000000..42ff95c --- /dev/null +++ b/examples/plot_free_energy.py @@ -0,0 +1,108 @@ +""" +================================ +ABMD biased dynamics +================================ + +Estimation of an overdamped Langevin in presence of biased dynamics. +""" + +import numpy as np +import folie as fl +import matplotlib.pyplot as plt +from scipy.stats import gaussian_kde + + +# First let's generate some biased trajectories +potential = fl.functions.MultiWell1D(V=4.0) + +x_range = np.linspace(0, 30, 150) +# plt.plot(x_range, potential.potential_plot(x_range.reshape(-1, 1))) +# plt.show() + +diff_function = fl.functions.Polynomial(deg=0, coefficients=np.array(1.0)) +model_simu = fl.Overdamped(potential, diffusion=diff_function) +simulator = fl.simulations.Simulator(fl.simulations.EulerStepper(model_simu), 1e-3) +data = simulator.run(250000, 30 * np.random.rand(25), 1) + +# Plot the resulting trajectories +# sphinx_gallery_thumbnail_number = 1 +# fig, axs = plt.subplots(1, 1) +# for n, trj in enumerate(data): +# axs[0].plot(trj["x"]) + +# axs[0].set_title("Trajectories") +# axs[0].set_xlabel("step") + + +fig, axs = plt.subplots(1, 3) +axs[0].set_title("Drift") +axs[0].set_xlabel("$x$") +axs[0].set_ylabel("$F(x)$") +axs[0].grid() + +axs[1].set_title("Diffusion") +axs[1].set_xlabel("$x$") +axs[1].set_ylabel("$D(x)$") +axs[1].grid() + +axs[2].set_title("Free energy") +axs[2].set_xlabel("$x$") +axs[2].set_ylabel("$V(x)$") +axs[2].grid() + + +xfa = np.linspace(0.0, 30.0, 75) +axs[0].plot(xfa, model_simu.drift(xfa.reshape(-1, 1)), label="Exact") +axs[1].plot(xfa, model_simu.diffusion(xfa.reshape(-1, 1)), label="Exact") + + +axs[2].plot(x_range, potential.potential_plot(x_range.reshape(-1, 1)), label="Exact") + + +name = "KramersMoyal" +model = fl.OverdampedSplines1D(fl.domains.MeshedDomain1D.create_from_range(np.linspace(-5, 35, 15))) +estimator = fl.KramersMoyalEstimator(model) +res = estimator.fit_fetch(data) +print(name, res.coefficients) +axs[0].plot(xfa, res.drift(xfa.reshape(-1, 1)), "--", label=name) +axs[1].plot(xfa, res.diffusion(xfa.reshape(-1, 1)), "--", label=name) + +# Compare Free energy +axs[2].plot(xfa, fl.analysis.free_energy_profile_1d(res, xfa), "--", label=name) + + +for name, marker, transitioncls in zip( + ["Euler", "Elerian", "Kessler", "Drozdov"], + ["+", "1", "2", "3"], + [ + fl.EulerDensity, + fl.ElerianDensity, + fl.KesslerDensity, + fl.DrozdovDensity, + ], +): + estimator = fl.KramersMoyalEstimator(model) + res = estimator.fit_fetch(data) + print(name, res.coefficients) + axs[0].plot(xfa, res.drift(xfa.reshape(-1, 1)), marker, label=name) + axs[1].plot(xfa, res.diffusion(xfa.reshape(-1, 1)), marker, label=name) + axs[2].plot(xfa, fl.analysis.free_energy_profile_1d(res, xfa), marker, label=name) + + +hist, bins = np.histogram(np.concatenate([trj["x"][25000:, 0] for trj in data]), bins=50) +x_bins = 0.5 * (bins[1:] + bins[:-1]) + +fe_hist = -np.log(hist) + +axs[2].plot(x_bins, fe_hist - np.min(fe_hist), label="Histogram") + +kde = gaussian_kde(np.concatenate([trj["x"][25000::100, 0] for trj in data])) + +fe_kde = -kde.logpdf(x_range) +axs[2].plot(x_range, fe_kde - np.min(fe_kde), label="KDE") + + +axs[0].legend() +axs[1].legend() +axs[2].legend() +plt.show() diff --git a/folie/analysis/overdamped_1d.py b/folie/analysis/overdamped_1d.py index 337e35f..68d7f37 100644 --- a/folie/analysis/overdamped_1d.py +++ b/folie/analysis/overdamped_1d.py @@ -3,7 +3,7 @@ """ from .._numpy import np -from scipy.integrate import cumulative_trapezoid +from scipy.integrate import cumulative_trapezoid, solve_ivp def free_energy_profile_1d(model, x): @@ -15,15 +15,22 @@ def free_energy_profile_1d(model, x): """ x = x.ravel() - diff_prime_val = model.diffusion.grad_x(x.reshape(-1, 1)).ravel() - drift_val = model.drift(x.reshape(-1, 1)).ravel() - diff_val = model.diffusion(x.reshape(-1, 1)).ravel() - diff_U = (-drift_val + diff_prime_val) / diff_val + def grad_V(x, _): + x = np.asarray(x).reshape(-1, 1) + diff_prime_val = model.diffusion.grad_x(x).ravel() + drift_val = model.drift(x).ravel() + diff_val = model.diffusion(x).ravel() - pmf = cumulative_trapezoid(diff_U, x, initial=0.0) + return (-drift_val + diff_prime_val) / diff_val - return pmf - np.min(pmf) + sol = solve_ivp(grad_V, [x.min() - 1e-10, x.max() + 1e10], np.array([0.0]), t_eval=x) # Add some epsilon to range to ensure inclusion of x + + V = sol.y.ravel() + + # V = cumulative_trapezoid(grad_V(x), x, initial=0.0) + + return V - np.min(V) def mfpt_1d(model, x_end: float, x_range, Npoints=500, x_start=None): diff --git a/folie/functions/analytical.py b/folie/functions/analytical.py index c11ceb3..a31667f 100644 --- a/folie/functions/analytical.py +++ b/folie/functions/analytical.py @@ -41,7 +41,8 @@ def potential_plot(self, *args): Get argument on the form of potential_plot(x,y) where x and y are result of np.meshgrid """ shapes = [a.shape[0] for a in args] - return self.potential(np.column_stack([a.ravel() for a in args])).reshape(*shapes) + pot = self.potential(np.column_stack([a.ravel() for a in args])).reshape(*shapes) + return pot - np.min(pot) class ConstantForce(PotentialFunction): @@ -121,6 +122,48 @@ def grad_V(self, x): return 2 * self.a * ((x - self.x0) * (x - 1.0)) * (2 * x - self.x0 - self.x1) +class MultiWell1D(PotentialFunction): + """ + A multiwell potential in 1D defined from a sum of Gaussian + """ + + def __init__(self, V=1.0, a=[0.1, 0.15, 0.04], x0=[5.0, 25.0, 15.0]): + self.V = V + self.a = a + self.x0 = x0 + self.n_well = len(self.a) + + self.dim = 1 + super().__init__() + + @property + def coefficients(self): + return np.array([self.V, *self.a, *self.x0]) + + def potential(self, X): + r""" + The potential is + + $$ V(x) =-\log(\sum_i e^{-a_i(x-x_0^i)^2/2}) + """ + pot = 0.0 + for n in range(self.n_well): + pot += np.exp(-0.5 * (self.a[n] * (X[:, 0] - self.x0[n]) ** 2)) + return -self.V * np.log(pot) + + def grad_V(self, X): + """ + Compute potential derivative + """ + x = X[:, 0] + dU = np.zeros(X.shape) + norm = 0.0 + for n in range(self.n_well): + dU[:, 0] += self.a[n] * (x - self.x0[n]) * np.exp(-0.5 * (self.a[n] * (x - self.x0[n]) ** 2)) + norm += np.exp(-0.5 * (self.a[n] * (x - self.x0[n]) ** 2)) + return self.V * np.divide(dU, norm[:, None], out=np.zeros_like(dU), where=norm[:, None] != 0) + + class Cosine(PotentialFunction): """ Cosine potential