Add args to TPS MD baseline

Author: plainerman

tps/second_order.py

@@ -15,17 +15,15 @@ def __init__(self, start_state, target_state, step_forward, step_backward, sampl
         self.sample_velocity = sample_velocity
 
 
-def one_way_shooting(system, trajectory, fixed_length, key):
-    raise NotImplementedError('Not implemented for multi step systems')
+def one_way_shooting(system, trajectory, fixed_length, dt, key):
     key = jax.random.split(key)
 
     # pick a random point along the trajectory
     point_idx = jax.random.randint(key[0], (1,), 1, len(trajectory) - 1)[0]
     # pick a random direction, either forward or backward
     direction = jax.random.randint(key[1], (1,), 0, 2)[0]
 
-    # TODO: Fix correct dt in ps / pass previous velocities
-    new_velocities = [(trajectory[point_idx] - trajectory[point_idx - 1]) / 0.001]
+    new_velocities = [(trajectory[point_idx] - trajectory[point_idx - 1]) / dt]
 
     if direction == 0:
         trajectory = trajectory[:point_idx + 1]
@@ -40,30 +38,42 @@ def one_way_shooting(system, trajectory, fixed_length, key):
     while len(trajectory) < steps:
         key, iter_key = jax.random.split(key)
         point, velocity = step_function(trajectory[-1], new_velocities[-1], iter_key)
-        trajectory.append(point)
-        new_velocities.append(velocity)
 
-        if jnp.isnan(point).any() or jnp.isnan(velocity).any():
-            return False, trajectory, new_velocities
+        nan_filter = jnp.isnan(point).any(axis=-1).flatten() | jnp.isnan(velocity).any(axis=-1).flatten()
+        too_big_filter = (jnp.abs(point) > MAX_ABS_VALUE).any(axis=-1).flatten()
 
-        # ensure that our trajectory does not explode
-        if (jnp.abs(point) > MAX_ABS_VALUE).any():
-            return False, trajectory, new_velocities
+        start_state_filter = system.start_state(point)
+        target_state_filter = system.target_state(point)
 
-        if system.start_state(trajectory[0]) and system.target_state(trajectory[-1]):
-            if fixed_length == 0 or len(trajectory) == fixed_length:
-                return True, trajectory, new_velocities
-            return False, trajectory, new_velocities
+        all_filters_combined = start_state_filter | target_state_filter | nan_filter | too_big_filter
+
+        limit = jnp.argmax(all_filters_combined) + 1 if all_filters_combined.any() else len(all_filters_combined)
+        trajectory.extend(point[:limit])
+        new_velocities.extend(velocity[:limit])
 
-        if system.target_state(trajectory[0]) and system.start_state(trajectory[-1]):
-            if fixed_length == 0 or len(trajectory) == fixed_length:
-                return True, trajectory[::-1], new_velocities[::-1]
+        if (nan_filter | too_big_filter)[:limit].any():
             return False, trajectory, new_velocities
 
+        if (start_state_filter | target_state_filter)[:limit].any():
+            break
+
+    # throw away the trajectory if it's not the right length
+    if len(trajectory) > steps:
+        return False, trajectory[:steps], new_velocities[:steps]
+
+    if fixed_length != 0 and len(trajectory) != fixed_length:
+        return False, trajectory, new_velocities
+
+    if system.start_state(trajectory[0]) and system.target_state(trajectory[-1]):
+        return True, trajectory, new_velocities
+
+    if system.target_state(trajectory[0]) and system.start_state(trajectory[-1]):
+        return True, trajectory[::-1], new_velocities[::-1]
+
     return False, trajectory, new_velocities
 
 
-def two_way_shooting(system, trajectory, fixed_length, key):
+def two_way_shooting(system, trajectory, fixed_length, _dt, key):
     key = jax.random.split(key)
 
     # pick a random point along the trajectory
@@ -139,7 +149,7 @@ def two_way_shooting(system, trajectory, fixed_length, key):
     return False, new_trajectory, new_velocities
 
 
-def mcmc_shooting(system, proposal, initial_trajectory, num_paths, key, fixed_length=0, warmup=50):
+def mcmc_shooting(system, proposal, initial_trajectory, num_paths, dt, key, fixed_length=0, warmup=50):
     # pick an initial trajectory
     trajectories = [initial_trajectory]
     velocities = []
@@ -162,12 +172,12 @@ def mcmc_shooting(system, proposal, initial_trajectory, num_paths, key, fixed_le
                 if len(trajectories) > warmup:
                     pbar.set_description('')
 
-                key, traj_idx_key, iter_key, accept_key = jax.random.split(key, 4)
+                key, traj_idx_key, ikey, accept_key = jax.random.split(key, 4)
                 traj_idx = jax.random.randint(traj_idx_key, (1,), warmup + 1, len(trajectories))[0]
                 # during warmup, we want an iterative scheme
                 traj_idx = traj_idx if traj_idx < len(trajectories) else -1
 
-                found, new_trajectory, new_velocities = proposal(system, trajectories[traj_idx], fixed_length, iter_key)
+                found, new_trajectory, new_velocities = proposal(system, trajectories[traj_idx], fixed_length, dt, ikey)
                 statistics['num_force_evaluations'] += len(new_trajectory) - 1
 
                 if not found:

tps_baseline.py

@@ -27,6 +27,14 @@
 from utils.animation import save_trajectory, to_md_traj
 from utils.rmsd import kabsch_align, kabsch_rmsd
 
+from argparse import ArgumentParser
+
+parser = ArgumentParser()
+parser.add_argument('--mechanism', type=str, choices=['one-way-shooting', 'two-way-shooting'], required=True)
+parser.add_argument('--fixed_length', type=int, default=0)
+parser.add_argument('--num_steps', type=int, default=10,
+                    help='The number of MD steps taken at once. More takes longer to compile but runs faster in the end.')
+
 
 def human_format(num):
     """https://stackoverflow.com/a/45846841/4417954"""
@@ -121,12 +129,14 @@ def step_n(step, _x, _v, n, _key):
 
 
 if __name__ == '__main__':
+    args = parser.parse_args()
+
     init_pdb = app.PDBFile("./files/AD_A.pdb")
     target_pdb = app.PDBFile("./files/AD_B.pdb")
     mdtraj_topology = md.Topology.from_openmm(init_pdb.topology)
     phis_psis = phi_psi_from_mdtraj(mdtraj_topology)
 
-    savedir = f"out/baselines/alanine"
+    savedir = f"out/baselines/alanine-{args.mechanism}"
     os.makedirs(savedir, exist_ok=True)
 
     # Construct the mass matrix
@@ -171,13 +181,8 @@ def U(_x):
 
     @jax.jit
     @jax.vmap
-    def dUdx_fn_unscaled(_x):
-        return jax.grad(lambda _x: U(_x).sum())(_x)
-
-
-    @jax.jit
     def dUdx_fn(_x):
-        return dUdx_fn_unscaled(_x) / mass / gamma
+        return jax.grad(lambda _x: U(_x).sum())(_x) / mass / gamma_in_ps
 
 
     @jax.jit
@@ -188,10 +193,10 @@ def step(_x, _key):
 
     @jax.jit
     def step_langevin_forward(_x, _v, _key):
-        """Perform one step of forward langevin"""
+        """Perform one step of forward langevin as implemented in openmm"""
         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_det = alpha * _v + f_scale * -dUdx_fn(_x)
         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
@@ -201,7 +206,7 @@ def step_langevin_forward(_x, _v, _key):
     def step_langevin_log_density(_x, _v, _new_x, _new_v):
         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_det = alpha * _v + f_scale * -dUdx_fn(_x)
         new_v_rand = new_v_det - _new_v
 
         return jax.scipy.stats.norm.logpdf(new_v_rand, 0, jnp.sqrt(kbT * (1 - alpha ** 2) / mass)).sum()
@@ -225,7 +230,7 @@ def step_langevin_backward(_x, _v, _key):
         alpha = jnp.exp(-gamma_in_ps * dt_in_ps)
         f_scale = (1 - alpha) / gamma_in_ps
         prev_x = _x - dt_in_ps * _v
-        prev_v = 1 / alpha * (_v + f_scale * dUdx_fn_unscaled(prev_x) / mass - jnp.sqrt(
+        prev_v = 1 / alpha * (_v + f_scale * dUdx_fn(prev_x) - jnp.sqrt(
             kbT * (1 - alpha ** 2) / mass) * jax.random.normal(_key, _x.shape))
 
         return prev_x, prev_v
@@ -281,8 +286,8 @@ def step_langevin_backward(_x, _v, _key):
         # jax.jit(jax.vmap(lambda s: kabsch_rmsd(B.reshape(22, 3), s.reshape(22, 3)) <= 7.5e-2)),
         # step_langevin_forward,
         # step_langevin_backward,
-        jax.jit(lambda _x, _v, _key: step_n(step_langevin_forward, _x, _v, 40, _key)),
-        jax.jit(lambda _x, _v, _key: step_n(step_langevin_backward, _x, _v, 40, _key)),
+        jax.jit(lambda _x, _v, _key: step_n(step_langevin_forward, _x, _v, args.num_steps, _key)),
+        jax.jit(lambda _x, _v, _key: step_n(step_langevin_backward, _x, _v, args.num_steps, _key)),
         jax.jit(lambda key: jnp.sqrt(kbT / mass) * jax.random.normal(key, (1, 66)))
     )
 
@@ -311,9 +316,17 @@ def step_langevin_backward(_x, _v, _key):
         with open(f'{savedir}/stats.json', 'r') as fp:
             statistics = json.load(fp)
     else:
+        if args.mechanism == 'one-way-shooting':
+            mechanism = tps2.one_way_shooting
+        elif args.mechanism == 'two-way-shooting':
+            mechanism = tps2.two_way_shooting
+        else:
+            raise ValueError(f"Unknown mechanism {args.mechanism}")
+
         try:
-            paths, velocities, statistics = tps2.mcmc_shooting(system, tps2.two_way_shooting, initial_trajectory,
-                                                               100, jax.random.PRNGKey(1), warmup=0, fixed_length=1000)
+            paths, velocities, statistics = tps2.mcmc_shooting(system, mechanism, initial_trajectory,
+                                                               100, dt_in_ps, jax.random.PRNGKey(1), warmup=0,
+                                                               fixed_length=args.fixed_length)
             # paths = tps2.unguided_md(system, B, 1, key)
             paths = [jnp.array(p) for p in paths]
             velocities = [jnp.array(p) for p in velocities]
@@ -324,6 +337,7 @@ def step_langevin_backward(_x, _v, _key):
             with open(f'{savedir}/stats.json', 'w') as fp:
                 json.dump(statistics, fp)
         except Exception as e:
+            print(e)
             breakpoint()
 
     print(statistics)