Add kabsch and second order integration

Author: plainerman

environment.yml

@@ -16,3 +16,4 @@ dependencies:
     - dmff @ git+https://github.com/deepmodeling/[email protected]
     - matplotlib==3.8.2
     - rdkit==2023.3.3
+    - ParmEd==4.2.2

tps_baseline.py

@@ -21,6 +21,9 @@
 from utils.PlotPathsAlanine_jax import PlotPathsAlanine
 from matplotlib import colors
 
+from utils.rmsd import kabsch
+from scipy.constants import physical_constants
+
 
 def human_format(num):
     """https://stackoverflow.com/a/45846841/4417954"""
@@ -88,7 +91,6 @@ def draw_path(_path, **kwargs):
         masked_path_x, masked_path_y = np.ma.MaskedArray(_path[:, 0], mask), np.ma.MaskedArray(_path[:, 1], mask)
         plt.plot(masked_path_x, masked_path_y, **kwargs)
 
-
     if path is not None:
         draw_path(path, color='red')
 
@@ -98,13 +100,21 @@ def draw_path(_path, **kwargs):
 
 
 T = 2.0
-dt = 1.0 * unit.microsecond
-dt = dt.value_in_unit(unit.second)
+dt_as_unit = unit.Quantity(value=1.0, unit=unit.femtoseconds)
+dt_in_ps = dt_as_unit.value_in_unit(unit.picosecond)
+dt = dt_as_unit.value_in_unit(unit.second)
+
+gamma_as_unit = 1.0 / unit.picosecond
+# actually gamma is 1/s, but we are working without units and just need the correct scaling
+# TODO: try to get rid of this duplicate definition
+gamma = 1.0 * unit.picosecond
+gamma_in_ps = gamma.value_in_unit(unit.picosecond)
+gamma = gamma.value_in_unit(unit.second)
 
 temp = 298.15
-temp = 10000  # TODO: remove this
 kbT = 1.380649 * 6.02214076 * 1e-3 * temp
 
+
 if __name__ == '__main__':
     init_pdb = app.PDBFile("./files/AD_c7eq.pdb")
     target_pdb = app.PDBFile("./files/AD_c7ax.pdb")
@@ -121,7 +131,7 @@ def draw_path(_path, **kwargs):
             new_mass.append(mass_)
     mass = jnp.array(new_mass)
     # Obtain sigma, gamma is by default 1
-    sigma = jnp.sqrt(2 * kbT / mass)
+    sigma = jnp.sqrt(2 * kbT / mass / gamma)
 
     # Initial and target shape [BS, 66]
     A = jnp.array(init_pdb.getPositions(asNumpy=True).value_in_unit(unit.nanometer)).reshape(1, -1)
@@ -135,44 +145,76 @@ def draw_path(_path, **kwargs):
                                     constraints=None,
                                     ewaldErrorTolerance=0.0005)
     # Create a box used when calling
-    # Calling U by U(x, box, pairs, ff.paramset.parameters), x is [22, 3] and output the energy, if it is batched, use vmap
     box = np.array([[50.0, 0.0, 0.0], [0.0, 50.0, 0.0], [0.0, 0.0, 50.0]])
     nbList = NeighborList(box, 4.0, potentials.meta["cov_map"])
     nbList.allocate(init_pdb.getPositions(asNumpy=True).value_in_unit(unit.nanometer))
     pairs = nbList.pairs
 
 
     def U(_x):
+        """
+        Calling U by U(x, box, pairs, ff.paramset.parameters), x is [22, 3] and output the energy, if it is batched, use vmap
+        """
         _U = potentials.getPotentialFunc()
 
         return _U(_x.reshape(22, 3), box, pairs, ff.paramset.parameters)
 
 
-    #TODO: we can introduce gamma here
     def dUdx_fn(_x):
-        return jax.grad(lambda _x: U(_x).sum())(_x) / mass
+        return jax.grad(lambda _x: U(_x).sum())(_x) / mass / gamma
 
 
     dUdx_fn = jax.vmap(dUdx_fn)
     dUdx_fn = jax.jit(dUdx_fn)
 
+
     @jax.jit
     def step(_x, _key):
         """Perform one step of forward euler"""
         return _x - dt * dUdx_fn(_x) + jnp.sqrt(dt) * sigma * jax.random.normal(_key, _x.shape)
 
+
+    def dUdx_fn_unscaled(_x):
+        return jax.grad(lambda _x: U(_x).sum())(_x)
+
+    dUdx_fn_unscaled = jax.vmap(dUdx_fn_unscaled)
+    dUdx_fn_unscaled = jax.jit(dUdx_fn_unscaled)
+
+    @jax.jit
+    def step_langevin(_x, _v, _key):
+        alpha = jnp.exp(-gamma_in_ps * dt_in_ps)
+        f_scale = (1 - alpha) / gamma_in_ps
+        new_v_det = alpha * _v + f_scale * -dUdx_fn_unscaled(_x) / mass
+        new_v = new_v_det + jnp.sqrt(kbT * (1 - alpha ** 2) / mass) * jax.random.normal(_key, _x.shape)
+
+        return _x + dt_in_ps * new_v, new_v
+
     key = jax.random.PRNGKey(1)
+    key, velocity_key = jax.random.split(key)
+    steps = 1_000_000
 
-    trajectory = [A]  # or [B]
+    trajectory = [A]
     _x = trajectory[-1]
-    steps = 20_000
+
+    velocity_variance = unit.Quantity(1 / mass, unit=1 / unit.dalton) * unit.BOLTZMANN_CONSTANT_kB * unit.Quantity(temp, unit=unit.kelvin)
+    # Although velocity+variance is of the unit J / Da = m^2 / s^2, openmm cannot handle this directly and we need to convert it
+    velocity_variance_in_si = 1 / physical_constants['unified atomic mass unit'][
+        0] * velocity_variance.value_in_unit(unit.joule / unit.dalton)
+    # velocity_variance_in_si = unit.Quantity(velocity_variance_in_si, unit.meter / unit.second)
+
+    _v = jnp.sqrt(velocity_variance_in_si) * jax.random.normal(velocity_key, _x.shape)
+    _v = unit.Quantity(_v, unit.meter / unit.second).value_in_unit(unit.nanometer / unit.picosecond)
+
     for i in trange(steps):
         key, iter_key = jax.random.split(key)
-        _x = step(_x, iter_key)
+        _x, _v = step_langevin(_x, _v, iter_key)
+
         trajectory.append(_x)
 
     trajectory = jnp.array(trajectory).reshape(-1, 66)
-    assert not jnp.isnan(trajectory).any()
+
+    # we only need to check whether the last frame contains nan, is it propagates
+    assert not jnp.isnan(trajectory[-1]).any()
     trajectory_phi_psi = phis_psis(trajectory, mdtraj_topology)
 
     trajs = None
@@ -189,15 +231,56 @@ def step(_x, _key):
     ramachandran(trajectory_phi_psi)
     plt.show()
 
-
     # TODO: this is work in progress. Get some baselines with tps
-    system = tps.System(
-        jax.jit(
-            lambda s: jnp.all(jnp.linalg.norm(A.reshape(-1, 22, 3) - s.reshape(-1, 22, 3), axis=2) <= 5e-2, axis=1)),
-        jax.jit(
-            lambda s: jnp.all(jnp.linalg.norm(B.reshape(-1, 22, 3) - s.reshape(-1, 22, 3), axis=2) <= 5e-2, axis=1)),
-        step
-    )
+
+    # l2_system = tps.System(
+    #     jax.jit(
+    #         lambda s: jnp.all(jnp.linalg.norm(A.reshape(-1, 22, 3) - s.reshape(-1, 22, 3), axis=2) <= 5e-2, axis=1)),
+    #     jax.jit(
+    #         lambda s: jnp.all(jnp.linalg.norm(B.reshape(-1, 22, 3) - s.reshape(-1, 22, 3), axis=2) <= 5e-2, axis=1)),
+    #     step
+    # )
+    #
+    # rmsd_system = tps.System(
+    #     jax.jit(lambda s: kabsch(A.reshape(22, 3), s.reshape(22, 3)) < 0.15),
+    #     jax.jit(lambda s: kabsch(B.reshape(22, 3), s.reshape(22, 3)) < 0.15),
+    #     step
+    # )
+    #
+    # # @jax.jit
+    # def is_within_phi_psi(s, center, radius, period=2 * jnp.pi):
+    #     points = phis_psis(s, mdtraj_topology)
+    #     delta = jnp.abs(center - points)
+    #     delta = jnp.where(delta > period / 2, delta - period, delta)
+    #
+    #     return jnp.hypot(delta[:, 0], delta[:, 1]) < radius
+    #
+    #
+    # deg = 180.0 / jnp.pi
+    # # State('A', torch.tensor([-150, 150]) / deg, torch.tensor([20, 45, 65, 80]) / deg),
+    # # State('B', torch.tensor([-70, 135]) / deg, torch.tensor([20, 45, 65, 75]) / deg),
+    # # State('C', torch.tensor([-150, -65]) / deg, torch.tensor([20, 45, 60]) / deg),
+    # # State('D', torch.tensor([-70, -50]) / deg, torch.tensor([20, 45, 60]) / deg),
+    # # State('E', torch.tensor([50, -100]) / deg, torch.tensor([20, 45, 65, 80]) / deg),
+    # # State('F', torch.tensor([40, 65]) / deg, torch.tensor([20, 45, 65, 80]) / deg),
+    #
+    # phi_psi_system = tps.System(
+    #     lambda s: is_within_phi_psi(s, jnp.array([-150, 150]) / deg, 20 / deg),
+    #     lambda s: is_within_phi_psi(s, jnp.array([50, -100]) / deg, 20 / deg),
+    #     step
+    # )
+    #
+    # # TODO: fix vmap
+    # filter1 = jax.vmap(phi_psi_system.start_state)(trajectory)
+    # filter2 = jax.vmap(phi_psi_system.target_state)(trajectory)
+    #
+    # plt.title('start')
+    # ramachandran(trajectory_phi_psi[filter1])
+    # plt.show()
+    #
+    # plt.title('target')
+    # ramachandran(trajectory_phi_psi[filter2])
+    # plt.show()
 
     # initial_trajectory = [t.reshape(1, -1) for t in interpolate([A, B], 100)]
 

utils/rmsd.py

@@ -0,0 +1,45 @@
+import jax.numpy as jnp
+import jax
+
+"""
+From https://hunterheidenreich.com/posts/kabsch_algorithm/ and adapted
+"""
+
+
+@jax.jit
+def kabsch(P, Q):
+    """
+    Computes the optimal rotation and translation to align two sets of points (P -> Q),
+    and their RMSD.
+
+    :param P: A Nx3 matrix of points
+    :param Q: A Nx3 matrix of points
+    :return: A tuple containing the optimal rotation matrix, the optimal
+             translation vector, and the RMSD.
+    """
+    assert P.shape == Q.shape, "Matrix dimensions must match"
+
+    # Compute centroids
+    centroid_P = jnp.mean(P, axis=0)
+    centroid_Q = jnp.mean(Q, axis=0)
+
+    # Optimal translation
+    t = centroid_Q - centroid_P
+
+    # Center the points
+    p = P - centroid_P
+    q = Q - centroid_Q
+
+    # Compute the covariance matrix
+    H = jnp.dot(p.T, q)
+
+    # SVD
+    U, S, Vt = jnp.linalg.svd(H)
+
+    # Validate right-handed coordinate system
+    Vt = jnp.where(jnp.linalg.det(jnp.dot(Vt.T, U.T)) < 0.0, -Vt, Vt)
+
+    # Optimal rotation
+    R = jnp.dot(Vt.T, U.T)
+
+    return jnp.sqrt(jnp.sum(jnp.square(jnp.dot(p, R.T) - q)) / P.shape[0])