Update baselines, shooting, and plotting

Author: plainerman

eval/path_metrics.py

@@ -0,0 +1,14 @@
+import jax.numpy as jnp
+import matplotlib.pyplot as plt
+from tqdm import tqdm
+
+
+def plot_path_energy(paths, U, reduce=jnp.max, already_ln=False):
+    reduced = jnp.array([reduce(U(path)) for path in tqdm(paths, 'Computing path metric')])
+
+    if already_ln:
+        # Convert reduced to log10
+        reduced = reduced / jnp.log(10)
+        plt.plot(jnp.arange(0, len(reduced), 1), reduced)
+    else:
+        plt.semilogy(jnp.arange(0, len(reduced), 1), reduced)

tps/second_order.py

@@ -23,8 +23,8 @@ def one_way_shooting(system, trajectory, fixed_length, key):
     # pick a random direction, either forward or backward
     direction = jax.random.randint(key[1], (1,), 0, 2)[0]
 
-    # TODO: Fix correct dt in ps
-    velocity = (trajectory[point_idx] - trajectory[point_idx - 1]) / 0.001
+    # TODO: Fix correct dt in ps / pass previous velocities
+    new_velocities = [(trajectory[point_idx] - trajectory[point_idx - 1]) / 0.001]
 
     if direction == 0:
         trajectory = trajectory[:point_idx + 1]
@@ -38,27 +38,28 @@ def one_way_shooting(system, trajectory, fixed_length, key):
     key, iter_key = jax.random.split(key[3])
     while len(trajectory) < steps:
         key, iter_key = jax.random.split(key)
-        point, velocity = step_function(trajectory[-1], velocity, iter_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
+            return False, trajectory, new_velocities
 
         # ensure that our trajectory does not explode
         if (jnp.abs(point) > MAX_ABS_VALUE).any():
-            return False, trajectory
+            return False, trajectory, new_velocities
 
         if system.start_state(trajectory[0]) and system.target_state(trajectory[-1]):
             if fixed_length == 0 or len(trajectory) == fixed_length:
-                return True, trajectory
-            return False, trajectory
+                return True, trajectory, new_velocities
+            return False, trajectory, new_velocities
 
         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]
-            return False, trajectory
+                return True, trajectory[::-1], new_velocities[::-1]
+            return False, trajectory, new_velocities
 
-    return False, trajectory
+    return False, trajectory, new_velocities
 
 
 def two_way_shooting(system, trajectory, fixed_length, key):
@@ -71,62 +72,74 @@ def two_way_shooting(system, trajectory, fixed_length, key):
 
     steps = MAX_STEPS if fixed_length == 0 else fixed_length
 
-    initial_velocity = system.sample_velocity(key[1])
-
-    key, iter_key = jax.random.split(key[2])
     new_trajectory = [point]
+    new_velocities = [system.sample_velocity(key[1])]
 
-    velocity = initial_velocity
+    key, iter_key = jax.random.split(key[2])
     while len(new_trajectory) < steps:
         key, iter_key = jax.random.split(key)
-        point, velocity = system.step_forward(new_trajectory[-1], velocity, iter_key)
+        point, velocity = system.step_forward(new_trajectory[-1], new_velocities[-1], iter_key)
         new_trajectory.append(point)
+        new_velocities.append(velocity)
 
         if jnp.isnan(point).any() or jnp.isnan(velocity).any():
-            return False, trajectory
+            return False, new_trajectory, new_velocities
 
         # ensure that our trajectory does not explode
         if (jnp.abs(point) > MAX_ABS_VALUE).any():
-            return False, trajectory
+            return False, new_trajectory, new_velocities
 
         if system.start_state(point) or system.target_state(point):
             break
 
-    velocity = initial_velocity
     while len(new_trajectory) < steps:
         key, iter_key = jax.random.split(key)
-        point, velocity = system.step_backward(new_trajectory[0], velocity, iter_key)
+        point, velocity = system.step_backward(new_trajectory[0], new_velocities[0], iter_key)
         new_trajectory.insert(0, point)
+        new_velocities.insert(0, velocity)
 
         if jnp.isnan(point).any() or jnp.isnan(velocity).any():
-            return False, trajectory
+            return False, new_trajectory, new_velocities
 
         # ensure that our trajectory does not explode
         if (jnp.abs(point) > MAX_ABS_VALUE).any():
-            return False, trajectory
+            return False, new_trajectory, new_velocities
 
         if system.start_state(point) or system.target_state(point):
             break
 
     # throw away the trajectory if it's not the right length
     if fixed_length != 0 and len(new_trajectory) != fixed_length:
-        return False, trajectory
+        return False, new_trajectory, new_velocities
 
     if system.start_state(new_trajectory[0]) and system.target_state(new_trajectory[-1]):
-        return True, new_trajectory
+        return True, new_trajectory, new_velocities
 
     if system.target_state(new_trajectory[0]) and system.start_state(new_trajectory[-1]):
-        return True, new_trajectory[::-1]
+        return True, new_trajectory[::-1], new_velocities[::-1]
 
-    return False, trajectory
+    return False, new_trajectory, new_velocities
 
 
 def mcmc_shooting(system, proposal, initial_trajectory, num_paths, key, fixed_length=0, warmup=50):
     # pick an initial trajectory
     trajectories = [initial_trajectory]
+    velocities = []
+    statistics = {
+        'num_force_evaluations': 0,
+        'num_tries': 0,
+        'num_metropolis_rejected': 0,
+        'warmup': warmup,
+        'num_paths': num_paths,
+        'max_steps': MAX_STEPS,
+        'max_abs_value': MAX_ABS_VALUE,
+    }
+    if fixed_length > 0:
+        statistics['fixed_length'] = fixed_length
 
     with tqdm(total=num_paths + warmup, desc='warming up' if warmup > 0 else '') as pbar:
         while len(trajectories) <= num_paths + warmup:
+            statistics['num_tries'] += 1
             if len(trajectories) > warmup:
                 pbar.set_description('')
 
@@ -135,7 +148,8 @@ def mcmc_shooting(system, proposal, initial_trajectory, num_paths, key, fixed_le
             # during warmup, we want an iterative scheme
             traj_idx = traj_idx if traj_idx < len(trajectories) else -1
 
-            found, new_trajectory = proposal(system, trajectories[traj_idx], fixed_length, iter_key)
+            found, new_trajectory, new_velocities = proposal(system, trajectories[traj_idx], fixed_length, iter_key)
+            statistics['num_force_evaluations'] += len(new_trajectory) - 1
 
             if not found:
                 continue
@@ -144,9 +158,12 @@ def mcmc_shooting(system, proposal, initial_trajectory, num_paths, key, fixed_le
             # 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:
                 trajectories.append(new_trajectory)
+                velocities.append(new_velocities)
                 pbar.update(1)
+            else:
+                statistics['num_metropolis_rejected'] += 1
 
-    return trajectories[warmup + 1:]
+    return trajectories[warmup + 1:], velocities[warmup:], statistics
 
 
 def unguided_md(system, initial_point, num_paths, key, fixed_length=0):

tps_baseline.py

@@ -1,9 +1,12 @@
 import os
+from functools import partial
+
 import jax
 import numpy as np
 import matplotlib.pyplot as plt
 import jax.numpy as jnp
 from tqdm import trange, tqdm
+import json
 
 # install openmm (from conda)
 import openmm.app as app
@@ -13,6 +16,7 @@
 # install mdtraj
 import mdtraj as md
 
+from eval.path_metrics import plot_path_energy
 from tps import first_order as tps1
 from tps import second_order as tps2
 from tps.plot import PeriodicPathHistogram
@@ -198,6 +202,26 @@ def step_langevin_forward(_x, _v, _key):
 
         return _x + dt_in_ps * new_v, new_v
 
+    @jax.jit
+    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_rand = new_v_det - _new_v
+
+        return jax.scipy.stats.norm.logpdf(new_v_rand, 0, jnp.sqrt(kbT * (1 - alpha ** 2) / mass)).sum()
+
+
+    def langevin_log_path_density(path_and_velocities):
+        path, velocities = path_and_velocities
+
+        log_prob = (-U(path[0]) / kbT).sum()
+        log_prob += jax.scipy.stats.norm.logpdf(velocities[0], 0, jnp.sqrt(kbT / mass)).sum()
+
+        for i in range(1, len(path)):
+            log_prob += step_langevin_log_density(path[i - 1], velocities[i - 1], path[i], velocities[i])
+
+        return log_prob
 
     @jax.jit
     def step_langevin_backward(_x, _v, _key):
@@ -235,8 +259,8 @@ def step_langevin_backward(_x, _v, _key):
 
     plt.title(f"{human_format(steps)} steps @ {temp} K, dt = {human_format(dt)}s")
     ramachandran(trajectory_phi_psi)
-    plt.scatter(phis_psis(A)[0], phis_psis(A)[1], color='red', marker='*')
-    plt.scatter(phis_psis(B)[0], phis_psis(B)[1], color='green', marker='*')
+    plt.scatter(phis_psis(A)[0, 0], phis_psis(A)[0, 1], color='red', marker='*')
+    plt.scatter(phis_psis(B)[0, 0], phis_psis(B)[0, 1], color='green', marker='*')
     plt.show()
 
     # Choose a system, either phi psi, or rmsd
@@ -255,10 +279,10 @@ def step_langevin_backward(_x, _v, _key):
     )
 
     system = tps2.SecondOrderSystem(
-        jax.jit(lambda s: is_within(phis_psis(s).reshape(-1, 2), phis_psis(A), radius)),
-        jax.jit(lambda s: is_within(phis_psis(s).reshape(-1, 2), phis_psis(B), radius)),
-        # jax.jit(jax.vmap(lambda s: kabsch_rmsd(A.reshape(22, 3), s.reshape(22, 3)) <= 7.5e-2)),
-        # jax.jit(jax.vmap(lambda s: kabsch_rmsd(B.reshape(22, 3), s.reshape(22, 3)) <= 7.5e-2)),
+        # jax.jit(lambda s: is_within(phis_psis(s).reshape(-1, 2), phis_psis(A), radius)),
+        # jax.jit(lambda s: is_within(phis_psis(s).reshape(-1, 2), phis_psis(B), radius)),
+        jax.jit(jax.vmap(lambda s: kabsch_rmsd(A.reshape(22, 3), s.reshape(22, 3)) <= 7.5e-2)),
+        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 key: jnp.sqrt(kbT / mass) * jax.random.normal(key, (1, 66)))
@@ -282,34 +306,58 @@ def step_langevin_backward(_x, _v, _key):
     initial_trajectory = [p for p in initial_trajectory]
     save_trajectory(mdtraj_topology, jnp.array(initial_trajectory), f'{savedir}/initial_trajectory.pdb')
 
-    load = False
+    load = True
     if load:
         paths = np.load(f'{savedir}/paths.npy', allow_pickle=True)
+        velocities = np.load(f'{savedir}/velocities.npy', allow_pickle=True)
+        with open(f'{savedir}/stats.json', 'r') as fp:
+            statistics = json.load(fp)
     else:
-        paths = tps2.mcmc_shooting(system, tps2.two_way_shooting, initial_trajectory, 100, jax.random.PRNGKey(1),
-                                   warmup=10)
+        paths, velocities, statistics = tps2.mcmc_shooting(system, tps2.two_way_shooting, initial_trajectory,
+                                               100, jax.random.PRNGKey(1), warmup=10)
         # paths = tps2.unguided_md(system, B, 1, key)
         paths = [jnp.array(p) for p in paths]
+        velocities = [jnp.array(p) for p in velocities]
         # store paths
         np.save(f'{savedir}/paths.npy', np.array(paths, dtype=object), allow_pickle=True)
+        np.save(f'{savedir}/velocities.npy', np.array(velocities, dtype=object), allow_pickle=True)
+        # save statistics, which is a dictionary
+        with open(f'{savedir}/stats.json', 'w') as fp:
+            json.dump(statistics, fp)
 
+    print(statistics)
     print([len(p) for p in paths])
     plt.hist([len(p) for p in paths], bins=jnp.sqrt(len(paths)).astype(int).item())
     plt.show()
 
     path_hist = PeriodicPathHistogram()
-    for i, path in tqdm(enumerate(paths)):
+    for i, path in tqdm(enumerate(paths), desc='Adding paths to histogram', total=len(paths)):
         path_hist.add_path(np.array(phis_psis(path)))
 
     plt.title(f"{human_format(len(paths))} paths @ {temp} K, dt = {human_format(dt)}s")
-    path_hist.plot(cmin=0.001)
+    path_hist.plot(cmin=0.01)
     ramachandran(None, states=[
         {'name': 'A', 'center': phis_psis(A).squeeze(), 'radius': radius},
         {'name': 'B', 'center': phis_psis(B).squeeze(), 'radius': radius},
     ], alpha=0.7)
     plt.savefig(f'{savedir}/paths.png', bbox_inches='tight')
     plt.show()
 
+    plot_path_energy(paths, jax.vmap(U))
+    plt.ylabel('Maximum energy')
+    plt.savefig(f'{savedir}/max_energy.png', bbox_inches='tight')
+    plt.show()
+
+    plot_path_energy(paths, jax.vmap(U), reduce=jnp.median)
+    plt.ylabel('Median energy')
+    plt.savefig(f'{savedir}/median_energy.png', bbox_inches='tight')
+    plt.show()
+
+    plot_path_energy(list(zip(paths, velocities)), langevin_log_path_density, reduce=lambda x: x, already_ln=True)
+    plt.ylabel('Path Density')
+    plt.savefig(f'{savedir}/path_density.png', bbox_inches='tight')
+    plt.show()
+
     for i, path in tqdm(enumerate(paths)):
         save_trajectory(mdtraj_topology, jnp.array([kabsch_align(p.reshape(-1, 3), B.reshape(-1, 3))[0] for p in path]),
                         f'{savedir}/trajectory_{i}.pdb')