Change free energy computation and add example

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()

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):

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