Put the correct facteor 2 where it is needed

examples/toy_models/plot_1D_Double_Well.py

@@ -45,12 +45,12 @@
 time_steps = 10000
 data = simulator.run(time_steps, q0, save_every=1)
 
-# # Plot resulting Trajectories
-# fig, axs = plt.subplots()
-# for n, trj in enumerate(data):
-#     axs.plot(trj["x"])
-#     axs.set_title("Trajectory")
-#
+# Plot resulting Trajectories
+fig, axs = plt.subplots()
+for n, trj in enumerate(data):
+    axs.plot(trj["x"])
+    axs.set_title("Trajectory")
+
 
 fig, axs = plt.subplots(1, 2)
 axs[0].set_title("Force")
@@ -68,27 +68,46 @@
 axs[1].plot(xfa, model_simu.diffusion(xfa.reshape(-1, 1)), label="Exact")
 trainforce = fl.functions.Polynomial(deg=3, coefficients=np.array([0, 0, 0, 0]))
 trainmodel = fl.models.Overdamped(force=trainforce, diffusion=fl.functions.Polynomial(deg=0, coefficients=np.asarray([0.9])), has_bias=False)
+
 for name, marker, transitioncls in zip(
     ["Euler", "Ozaki", "ShojiOzaki", "Elerian", "Kessler", "Drozdov"],
     ["x", "|", ".", "1", "2", "3"],
     [
         fl.EulerDensity,
+        fl.OzakiDensity,
+        fl.ShojiOzakiDensity,
+        fl.ElerianDensity,
+        fl.KesslerDensity,
+        fl.DrozdovDensity,
     ],
 ):
     estimator = fl.LikelihoodEstimator(transitioncls(trainmodel))
     res = estimator.fit_fetch(data)
     # print(res.coefficients)
     axs[0].plot(xfa, res.force(xfa.reshape(-1, 1)), marker=marker, label=name)
     axs[1].plot(xfa, res.diffusion(xfa.reshape(-1, 1)), marker=marker, label=name)
+
+
+name = "Kramers Moyal"
+res = fl.KramersMoyalEstimator(trainmodel).fit_fetch(data)
+axs[0].plot(xfa, res.force(xfa.reshape(-1, 1)), "--", label=name)
+axs[1].plot(xfa, res.diffusion(xfa.reshape(-1, 1)), "--", label=name)
+
+
 axs[0].legend()
 axs[1].legend()
 
+# Compute MFPT from one well to another
 plt.figure()
 
 x_mfpt, mfpt = fl.analysis.mfpt_1d(model_simu, -5.0, [-10.0, 10.0], Npoints=500)
-plt.plot(x_mfpt, mfpt)
+plt.plot(x_mfpt, mfpt, label="Right to left")
 x_mfpt, mfpt = fl.analysis.mfpt_1d(model_simu, 5.0, [-10.0, 10.0], Npoints=500)
-plt.plot(x_mfpt, mfpt)
-print(fl.analysis.mfpt_1d(model_simu, -5.0, [-10.0, 10.0], x_start=5.0, Npoints=500))
-print(fl.analysis.mfpt_1d(model_simu, 5.0, [-10.0, 10.0], x_start=-5.0, Npoints=500))
+plt.plot(x_mfpt, mfpt, label="Left to right")
+
+x_mfpt, mfpt = fl.analysis.mfpt_1d(res, -5.0, [-10.0, 10.0], Npoints=500)
+plt.plot(x_mfpt, mfpt, label="Right to left Estimation")
+x_mfpt, mfpt = fl.analysis.mfpt_1d(res, 5.0, [-10.0, 10.0], Npoints=500)
+plt.plot(x_mfpt, mfpt, label="Left to right Estimation")
+plt.legend()
 plt.show()

folie/analysis/overdamped_1d.py

@@ -19,7 +19,7 @@ def free_energy_profile_1d(model, x):
     force_val = model.force(x.reshape(-1, 1)).ravel()
     diff_val = model.diffusion(x.reshape(-1, 1)).ravel()
 
-    diff_U = diff_U = -force_val / diff_val + diff_prime_val
+    diff_U = -force_val / diff_val + diff_prime_val
 
     pmf = cumulative_trapezoid(diff_U, x, initial=0.0)
 

folie/estimation/direct_estimation.py

@@ -93,6 +93,6 @@ def fit(self, data, estimator=linear_model.LinearRegression(copy_X=False, fit_in
             dx_sq = dx**2
         else:
             dx_sq = dx[..., None] * dx[:, None, ...]
-        self.model.diffusion.fit(X, y=dx_sq / dt, estimator=estimator, **kwargs)
+        self.model.diffusion.fit(X, y=0.5 * dx_sq / dt, estimator=estimator, **kwargs)
         self.model.fitted_ = True
         return self

folie/estimation/overdamped_transitionDensity.py

@@ -65,12 +65,12 @@ def _logdensity1D(self, x, xt, dt: float, bias=0.0, **kwargs):
         :param dt: float, the time step between x and xt
         :return: probability (same dimension as x and xt)
         """
-        sig2t = (self._model.diffusion(x, **kwargs)).ravel() * dt
+        sig2t = 2 * (self._model.diffusion(x, **kwargs)).ravel() * dt
         mut = x.ravel() + self._model.meandispl(x, bias, **kwargs).ravel() * dt
         if not self.use_jac:
             return gaussian_likelihood_1D(xt, mut, sig2t), np.zeros(2)
 
-        jacV = (self._model.diffusion.grad_coeffs(x, **kwargs)) * dt
+        jacV = 2 * (self._model.diffusion.grad_coeffs(x, **kwargs)) * dt
         return gaussian_likelihood_derivative_1D(xt, mut, sig2t, self._model.meandispl.grad_coeffs(x, bias, **kwargs) * dt, jacV)
 
     def _logdensityND(self, x, xt, dt, bias=0.0, **kwargs):
@@ -84,11 +84,11 @@ def _logdensityND(self, x, xt, dt, bias=0.0, **kwargs):
 
         # TODO: Add correction terms
         E = x + self._model.meandispl(x, bias, **kwargs) * dt
-        V = (self._model.diffusion(x, **kwargs)) * dt
+        V = 2 * (self._model.diffusion(x, **kwargs)) * dt
         if not self.use_jac:
             ll = gaussian_likelihood_ND(xt, E, V)
             return ll, np.zeros(2)
-        jacV = (self._model.diffusion.grad_coeffs(x, **kwargs)) * dt
+        jacV = 2 * (self._model.diffusion.grad_coeffs(x, **kwargs)) * dt
         jacE = self._model.meandispl.grad_coeffs(x, bias, **kwargs) * dt
         return gaussian_likelihood_derivative_ND(xt, E, V, jacE, jacV)
 
@@ -107,7 +107,7 @@ def _hiddenvariance(self, x, xt, sigh, dt, **kwargs):
         """
 
         E2 = self._model.friction(x[:, : self._model.dim_x], **kwargs) * dt
-        V = (self._model.diffusion(x[:, : self._model.dim_x], **kwargs)) * dt
+        V = 2 * (self._model.diffusion(x[:, : self._model.dim_x], **kwargs)) * dt
         invV = np.linalg.inv(V)
         dh = self._model.dim_h
         dhdh = sigh[:, :dh, :dh] - sigh[:, :dh, dh:] - sigh[:, dh:, :dh] + sigh[:, dh:, dh:]
@@ -128,7 +128,7 @@ def _hiddenvariance(self, x, xt, sigh, dt, **kwargs):
 
         l_jac_E = -0.5 * np.einsum("tijc,tji->tc", jEV[:, :, dh:, :], dhh) - 0.5 * np.einsum("tijc,tji->tc", VjE[:, dh:, ...], hdh) + np.einsum("tijc,tji->tc", EVjE, hh)
 
-        jacV = self._model.diffusion.grad_coeffs(x[:, : self._model.dim_x], **kwargs) * dt
+        jacV = 2 * self._model.diffusion.grad_coeffs(x[:, : self._model.dim_x], **kwargs) * dt
         jacinvV = -np.einsum("tij,tjkc,tkl->tilc", invV, jacV, invV)
         EjVE = np.einsum("tdh,tdfc,tfg-> thgc", E2, jacinvV, E2)
         EjV = np.einsum("tdh,tdfc-> thfc", E2, jacinvV)
@@ -154,7 +154,7 @@ def e_step(self, weight, trj, coefficients, mu0, sig0):
             trj["xt"][:, : self._model.dim_x],
             self._model.force(trj["x"][:, : self._model.dim_x], trj["bias"][:, : self._model.dim_x]) * trj["dt"],
             self._model.friction(trj["x"][:, : self._model.dim_x]) * trj["dt"],
-            self._model.diffusion(trj["x"][:, : self._model.dim_x]) * trj["dt"],
+            2 * self._model.diffusion(trj["x"][:, : self._model.dim_x]) * trj["dt"],
             mu0,
             sig0,
         )
@@ -181,7 +181,7 @@ def _logdensity1D(self, x, xt, dt: float, bias=0.0, **kwargs):
         :param dt: float, the time step between x and xt
         :return: probability (same dimension as x and xt)
         """
-        sig = self._model.diffusion(x, **kwargs).ravel()
+        sig = 2 * self._model.diffusion(x, **kwargs).ravel()
         mu = self._model.meandispl(x, bias, **kwargs).ravel()
         mu_x = self._model.meandispl.grad_x(x, bias, **kwargs).ravel()
         temp = mu * (np.exp(mu_x * dt) - 1) / mu_x
@@ -209,19 +209,19 @@ def _logdensity1D(self, x, xt, dt: float, bias=0.0, **kwargs):
         :param dt: float, the time step between x and xt
         :return: probability (same dimension as x and xt)
         """
-        sig = np.sqrt(self._model.diffusion(x, **kwargs).ravel())
+        sig = np.sqrt(2 * self._model.diffusion(x, **kwargs).ravel())
         mu = self._model.meandispl(x, bias, **kwargs).ravel()
 
-        Mt = 0.5 * sig ** 2 * self._model.meandispl.hessian_x(x, bias, **kwargs).ravel()  # + self._model.meandispl_t(x)  #Time homogenous model
+        Mt = 0.5 * sig**2 * self._model.meandispl.hessian_x(x, bias, **kwargs).ravel()  # + self._model.meandispl_t(x)  #Time homogenous model
         Lt = self._model.meandispl.grad_x(x, bias, **kwargs).ravel()
         if (Lt == 0).any():  # TODO: need to fix this
             B = sig * np.sqrt(dt)
-            A = x.ravel() + mu * dt + Mt * dt ** 2 / 2
+            A = x.ravel() + mu * dt + Mt * dt**2 / 2
         else:
             B = sig * np.sqrt((np.exp(2 * Lt * dt) - 1) / (2 * Lt))
 
             elt = np.exp(Lt * dt) - 1
-            A = x.ravel() + mu / Lt * elt + Mt / (Lt ** 2) * (elt - Lt * dt)
+            A = x.ravel() + mu / Lt * elt + Mt / (Lt**2) * (elt - Lt * dt)
 
         return gaussian_likelihood_1D(xt, A, B)
 
@@ -245,17 +245,17 @@ def _logdensity1D(self, x, xt, dt: float, bias=0.0, **kwargs):
         :param dt: float, the time step between x and xt
         :return: probability (same dimension as x and xt)
         """
-        sig_x = self._model.diffusion.grad_x(x, **kwargs).ravel()
+        sig_x = 2 * self._model.diffusion.grad_x(x, **kwargs).ravel()
         if isinstance(x, np.ndarray) and (sig_x == 0).any:
             return super()._logdensity1D(x=x, xt=xt, dt=dt, bias=bias, **kwargs)[0]
 
-        sig = self._model.diffusion(x, **kwargs).ravel()
+        sig = 2 * self._model.diffusion(x, **kwargs).ravel()
         mu = self._model.meandispl(x, bias, **kwargs).ravel()
 
         A = sig * sig_x * dt * 0.5
         B = -0.5 * sig / sig_x + x.ravel() + mu * dt - A
         z = (xt.ravel() - B) / A
-        C = 1.0 / (sig_x ** 2 * dt)
+        C = 1.0 / (sig_x**2 * dt)
 
         scz = np.sqrt(C * z)
         cpz = -0.5 * (C + z)
@@ -286,17 +286,17 @@ def _logdensity1D(self, x, xt, dt: float, bias=0.0, **kwargs):
         :param dt: float, the time of observing Xt
         :return: probability (same dimension as x and xt)
         """
-        sig = self._model.diffusion(x, **kwargs).ravel()
-        sig_x = self._model.diffusion.grad_x(x, **kwargs).ravel()
-        sig_xx = self._model.diffusion.hessian_x(x, **kwargs).ravel()
+        sig = 2 * self._model.diffusion(x, **kwargs).ravel()
+        sig_x = 2 * self._model.diffusion.grad_x(x, **kwargs).ravel()
+        sig_xx = 2 * self._model.diffusion.hessian_x(x, **kwargs).ravel()
         mu = self._model.meandispl(x, bias, **kwargs).ravel()
         mu_x = self._model.meandispl.grad_x(x, bias, **kwargs).ravel()
         mu_xx = self._model.meandispl.hessian_x(x, bias, **kwargs).ravel()
         x = x.ravel()
-        d = dt ** 2 / 2
+        d = dt**2 / 2
         E = x + mu * dt + (mu * mu_x + 0.5 * sig * mu_xx) * d
 
-        V = x ** 2 + (2 * mu * x + sig) * dt + (2 * mu * (mu_x * x + mu + 0.5 * sig_x) + sig * (mu_xx * x + 2 * mu_x + 0.5 * sig_xx)) * d - E ** 2
+        V = x**2 + (2 * mu * x + sig) * dt + (2 * mu * (mu_x * x + mu + 0.5 * sig_x) + sig * (mu_xx * x + 2 * mu_x + 0.5 * sig_xx)) * d - E**2
         V = np.abs(V)
         return gaussian_likelihood_1D(xt, E, V)
 
@@ -317,14 +317,14 @@ def _logdensity1D(self, x, xt, dt: float, bias=0.0, **kwargs):
         :param dt: float, the time of observing Xt
         :return: probability (same dimension as x and xt)
         """
-        sig = self._model.diffusion(x, **kwargs).ravel()
-        sig_x = self._model.diffusion.grad_x(x, **kwargs).ravel()
-        sig_xx = self._model.diffusion.hessian_x(x, **kwargs).ravel()
+        sig = 2 * self._model.diffusion(x, **kwargs).ravel()
+        sig_x = 2 * self._model.diffusion.grad_x(x, **kwargs).ravel()
+        sig_xx = 2 * self._model.diffusion.hessian_x(x, **kwargs).ravel()
         mu = self._model.meandispl(x, bias, **kwargs).ravel()
         mu_x = self._model.meandispl.grad_x(x, bias, **kwargs).ravel()
         mu_xx = self._model.meandispl.hessian_x(x, bias, **kwargs).ravel()
 
-        d = dt ** 2 / 2
+        d = dt**2 / 2
         E = x.ravel() + mu * dt + (mu * mu_x + 0.5 * sig * mu_xx) * d
 
         V = sig * dt + (mu * sig_x + 2 * mu_x * sig + sig * sig_xx) * d

folie/models/overdamped.py

@@ -168,9 +168,9 @@ class Overdamped(BaseModelOverdamped):
 
     .. math::
 
-        \mathrm{d}X(t) = F(X)\mathrm{d}t + sigma(X,t)\mathrm{d}W_t
+        \mathrm{d}X(t) = F(X)\mathrm{d}t + \sigma(X,t)\mathrm{d}W_t
 
-    The components of the overdamped model are the force profile F(X) as well as the diffusion :math: `D(x) = \sigma(X)\sigma(X)^\T`
+    The components of the overdamped model are the force profile F(X) as well as the diffusion :math: `D(x) =  \frac{1}{2} \sigma(X)\sigma(X)^\T`
 
     When considering equilibrium model, the force and diffusion profile are related to the free energy profile V(X) via
 

folie/simulations/stepper.py

@@ -39,19 +39,18 @@ def run_step(self, x, dt, dW, bias=0.0):
 
 class EulerStepper(Stepper):
     def run_step_1D(self, x, dt, dW, bias=0.0):
-        sig_sq_dt = np.sqrt(self.model.diffusion(x) * dt)
+        sig_sq_dt = np.sqrt(2 * self.model.diffusion(x) * dt)
         return (x.T + (self.model.meandispl(x, bias)) * dt + sig_sq_dt * dW.T).T  # Weird transpose for broadcasting
 
     def run_step_ND(self, x, dt, dW, bias=0.0):  # New function
-        sig_sq_dt = np.sqrt(self.model.diffusion(x) * dt)  # Work only for diagonal diffusion that should a cholesky instead
+        sig_sq_dt = np.sqrt(2 * self.model.diffusion(x) * dt)  # Work only for diagonal diffusion that should a cholesky instead
         return x + self.model.meandispl(x, bias) * dt + np.einsum("ijk,ik->ij", sig_sq_dt, dW)
 
 
-
 class MilsteinStepper(Stepper):
     def run_step_1D(self, x, dt, dW, bias=0.0):
-        sig_sq_dt = np.sqrt(self.model.diffusion(x) * dt)
-        return (x.T + (self.model.meandispl(x, bias)) * dt + sig_sq_dt * dW.T + 0.25 * self.model.diffusion.grad_x(x)[..., 0] * ((dW.T)**2 - 1) * dt).T  # same problem of dW.T as in EulerStepper class
+        sig_sq_dt = np.sqrt(2 * self.model.diffusion(x) * dt)
+        return (x.T + (self.model.meandispl(x, bias)) * dt + sig_sq_dt * dW.T + 0.5 * self.model.diffusion.grad_x(x)[..., 0] * ((dW.T) ** 2 - 1) * dt).T  # same problem of dW.T as in EulerStepper class
 
 
 class VECStepper(Stepper):