Fix shooting and TPS baseliens

Author: plainerman

eval/path_metrics.py

@@ -12,3 +12,5 @@ def plot_path_energy(paths, U, reduce=jnp.max, add=0, already_ln=False, **kwargs
         plt.plot(jnp.arange(0, len(reduced), 1), reduced, **kwargs)
     else:
         plt.semilogy(jnp.arange(0, len(reduced), 1), reduced, **kwargs)
+
+    return reduced

evaluate_mueller.py

@@ -0,0 +1,15 @@
+import jax.numpy as jnp
+from tqdm import tqdm
+
+
+def decorrelated(paths):
+    prev = paths[0]
+    decorrelated = [prev]
+
+    for x in tqdm(paths[1:]):
+        # check if the two arrays share a common value
+        if not jnp.in1d(prev, x).any():
+            decorrelated.append(x)
+            prev = x
+
+    return decorrelated

tps/paths.py

@@ -15,21 +15,25 @@ 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, dt, key):
+def one_way_shooting(system, trajectory, previous_velocities, fixed_length, dt, key):
     key = jax.random.split(key)
 
+    if previous_velocities is None:
+        previous_velocities = [(trajectory[i] - trajectory[i - 1]) / dt for i in range(1, len(trajectory))]
+        previous_velocities.insert(0, system.sample_velocity(key[0]))
+
     # pick a random point along the trajectory
-    point_idx = jax.random.randint(key[0], (1,), 1, len(trajectory) - 1)[0]
+    point_idx = jax.random.randint(key[1], (1,), 1, len(trajectory) - 1)[0]
     # pick a random direction, either forward or backward
-    direction = jax.random.randint(key[1], (1,), 0, 2)[0]
-
-    new_velocities = [(trajectory[point_idx] - trajectory[point_idx - 1]) / dt]
+    direction = jax.random.randint(key[2], (1,), 0, 2)[0]
 
     if direction == 0:
         trajectory = trajectory[:point_idx + 1]
+        new_velocities = previous_velocities[:point_idx + 1]
         step_function = system.step_forward
     else:  # direction == 1:
         trajectory = trajectory[point_idx:][::-1]
+        new_velocities = previous_velocities[point_idx:][::-1]
         step_function = system.step_backward
 
     steps = MAX_STEPS if fixed_length == 0 else fixed_length
@@ -73,7 +77,7 @@ def one_way_shooting(system, trajectory, fixed_length, dt, key):
     return False, trajectory, new_velocities
 
 
-def two_way_shooting(system, trajectory, fixed_length, _dt, key):
+def two_way_shooting(system, trajectory, _previous_velocities, fixed_length, _dt, key):
     key = jax.random.split(key)
 
     # pick a random point along the trajectory
@@ -170,11 +174,14 @@ def mcmc_shooting(system, proposal, initial_trajectory, num_paths, dt, key, fixe
         velocities = stored['velocities']
         statistics = stored['statistics']
 
+    num_tries = 0
+    num_force_evaluations = 0
+    num_metropolis_rejected = 0
     try:
         with tqdm(total=num_paths + warmup, initial=len(trajectories) - 1,
                   desc='warming up' if warmup > 0 else '') as pbar:
             while len(trajectories) <= num_paths + warmup:
-                statistics['num_tries'] += 1
+                num_tries += 1
                 if len(trajectories) > warmup:
                     pbar.set_description('')
 
@@ -183,20 +190,33 @@ def mcmc_shooting(system, proposal, initial_trajectory, num_paths, dt, key, fixe
                 # 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, dt, ikey)
-                statistics['num_force_evaluations'] += len(new_trajectory) - 1
+                # trajectories and velocities are one off
+                found, new_trajectory, new_velocities = proposal(system,
+                                                                 trajectories[traj_idx],
+                                                                 velocities[traj_idx - 1] if len(trajectories) > 1 else None,
+                                                                 fixed_length, dt, ikey)
+                num_force_evaluations += len(new_trajectory) - 1
 
                 if not found:
                     continue
 
                 ratio = len(trajectories[-1]) / len(new_trajectory)
                 # The first trajectory might have a very unreasonable length, so we skip it
                 if len(trajectories) == 1 or jax.random.uniform(accept_key, shape=(1,)) < ratio:
+                    # only update them in the dictionary once accepted
+                    # this allows us to continue the progress
+                    statistics['num_tries'] += num_tries
+                    statistics['num_force_evaluations'] += num_force_evaluations
+                    statistics['num_metropolis_rejected'] += num_metropolis_rejected
+                    num_tries = 0
+                    num_force_evaluations = 0
+                    num_metropolis_rejected = 0
+
                     trajectories.append(new_trajectory)
                     velocities.append(new_velocities)
                     pbar.update(1)
                 else:
-                    statistics['num_metropolis_rejected'] += 1
+                    num_metropolis_rejected += 1
     except KeyboardInterrupt:
         print('SIGINT received, stopping early')
         # Fix in case we stop when adding a trajectory

tps/second_order.py

@@ -1,5 +1,4 @@
 import os
-from functools import partial
 import traceback
 import jax
 import numpy as np
@@ -33,6 +32,7 @@
 parser.add_argument('--mechanism', type=str, choices=['one-way-shooting', 'two-way-shooting'], required=True)
 parser.add_argument('--states', type=str, default='phi-psi', choices=['phi-psi', 'rmsd'])
 parser.add_argument('--fixed_length', type=int, default=0)
+parser.add_argument('--warmup', type=int, default=0)
 parser.add_argument('--num_paths', type=int, required=True)
 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.')
@@ -213,6 +213,18 @@ def step_langevin_forward(_x, _v, _key):
         return _x + dt_in_ps * new_v, new_v
 
 
+    @jax.jit
+    def step_langevin_backward(_x, _v, _key):
+        """Perform one step of backward langevin"""
+        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(prev_x) - jnp.sqrt(
+            kbT * (1 - alpha ** 2) / mass) * jax.random.normal(_key, _x.shape))
+
+        return prev_x, prev_v
+
+
     @jax.jit
     def step_langevin_log_prob(_x, _v, _new_x, _new_v):
         alpha = jnp.exp(-gamma_in_ps * dt_in_ps)
@@ -225,6 +237,8 @@ def step_langevin_log_prob(_x, _v, _new_x, _new_v):
 
     def langevin_log_path_likelihood(path_and_velocities):
         path, velocities = path_and_velocities
+        assert len(path) == len(velocities), \
+            f'path and velocities must have the same length, but got {len(path)} and {len(velocities)}'
 
         log_prob = (-U(path[0]) / kbT).sum()
         log_prob += jax.scipy.stats.norm.logpdf(velocities[0], 0, jnp.sqrt(kbT / mass)).sum()
@@ -235,18 +249,6 @@ def langevin_log_path_likelihood(path_and_velocities):
         return log_prob
 
 
-    @jax.jit
-    def step_langevin_backward(_x, _v, _key):
-        """Perform one step of backward langevin"""
-        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(prev_x) - jnp.sqrt(
-            kbT * (1 - alpha ** 2) / mass) * jax.random.normal(_key, _x.shape))
-
-        return prev_x, prev_v
-
-
     # Choose a system, either phi psi, or rmsd
     # system = tps1.System(
     #     jax.jit(jax.vmap(lambda s: kabsch_rmsd(A.reshape(22, 3), s.reshape(22, 3)) < 0.1)),
@@ -324,7 +326,8 @@ def step_langevin_backward(_x, _v, _key):
 
     try:
         paths, velocities, statistics = tps2.mcmc_shooting(system, mechanism, initial_trajectory,
-                                                           args.num_paths, dt_in_ps, jax.random.PRNGKey(1), warmup=0,
+                                                           args.num_paths, dt_in_ps, jax.random.PRNGKey(1),
+                                                           warmup=args.warmup,
                                                            fixed_length=args.fixed_length,
                                                            stored=stored)
         # paths = tps2.unguided_md(system, B, 1, key)