Allow progress continuation in tps_baselines.py

Author: plainerman

tps/second_order.py

@@ -149,7 +149,7 @@ def two_way_shooting(system, trajectory, fixed_length, _dt, key):
     return False, new_trajectory, new_velocities
 
 
-def mcmc_shooting(system, proposal, initial_trajectory, num_paths, dt, key, fixed_length=0, warmup=50):
+def mcmc_shooting(system, proposal, initial_trajectory, num_paths, dt, key, fixed_length=0, warmup=50, stored=None):
     # pick an initial trajectory
     trajectories = [initial_trajectory]
     velocities = []
@@ -165,8 +165,14 @@ def mcmc_shooting(system, proposal, initial_trajectory, num_paths, dt, key, fixe
     if fixed_length > 0:
         statistics['fixed_length'] = fixed_length
 
+    if stored is not None:
+        trajectories = stored['trajectories']
+        velocities = stored['velocities']
+        statistics = stored['statistics']
+
     try:
-        with tqdm(total=num_paths + warmup, desc='warming up' if warmup > 0 else '') as pbar:
+        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
                 if len(trajectories) > warmup:
@@ -194,7 +200,7 @@ def mcmc_shooting(system, proposal, initial_trajectory, num_paths, dt, key, fixe
     except KeyboardInterrupt:
         print('SIGINT received, stopping early')
         # Fix in case we stop when adding a trajectory
-        if len(trajectories) > len(velocities):
+        if len(trajectories) > len(velocities) + 1:
             velocities.append(new_velocities)
 
     return trajectories[warmup + 1:], velocities[warmup:], statistics

tps_baseline.py

@@ -34,6 +34,8 @@
 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.')
+parser.add_argument('--resume', action='store_true')
+parser.add_argument('--override', action='store_true')
 
 
 def human_format(num):
@@ -136,7 +138,7 @@ def step_n(step, _x, _v, n, _key):
     mdtraj_topology = md.Topology.from_openmm(init_pdb.topology)
     phis_psis = phi_psi_from_mdtraj(mdtraj_topology)
 
-    savedir = f"out/baselines/alanine-{args.mechanism}"
+    savedir = f"out/baselines/alanine-{args.mechanism}-{args.fixed_length}"
     os.makedirs(savedir, exist_ok=True)
 
     # Construct the mass matrix
@@ -203,7 +205,7 @@ def step_langevin_forward(_x, _v, _key):
 
 
     @jax.jit
-    def step_langevin_log_density(_x, _v, _new_x, _new_v):
+    def step_langevin_log_prob(_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(_x)
@@ -212,14 +214,14 @@ def step_langevin_log_density(_x, _v, _new_x, _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):
+    def langevin_log_path_likelihood(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])
+            log_prob += step_langevin_log_prob(path[i - 1], velocities[i - 1], path[i], velocities[i])
 
         return log_prob
 
@@ -236,34 +238,6 @@ def step_langevin_backward(_x, _v, _key):
         return prev_x, prev_v
 
 
-    key = jax.random.PRNGKey(1)
-    key, velocity_key = jax.random.split(key)
-    steps = 10_000
-
-    trajectory = [A]
-    _x = trajectory[-1]
-    _v = jnp.sqrt(kbT / mass) * jax.random.normal(velocity_key, (1, 66))
-
-    for i in trange(steps):
-        key, iter_key = jax.random.split(key)
-        _x, _v = step_langevin_forward(_x, _v, iter_key)
-
-        trajectory.append(_x)
-
-    trajectory = jnp.array(trajectory).reshape(-1, 66)
-
-    # save_trajectory(mdtraj_topology, trajectory[-1000:], 'simulation.pdb')
-
-    # 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)
-
-    plt.title(f"{human_format(steps)} steps @ {temp} K, dt = {human_format(dt)}s")
-    ramachandran(trajectory_phi_psi)
-    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
     # system = tps1.System(
     #     jax.jit(jax.vmap(lambda s: kabsch_rmsd(A.reshape(22, 3), s.reshape(22, 3)) < 0.1)),
@@ -273,72 +247,69 @@ def step_langevin_backward(_x, _v, _key):
 
     radius = 20 / deg
 
-    system = tps1.FirstOrderSystem(
-        lambda s: is_within(phis_psis(s).reshape(-1, 2), phis_psis(A), radius),
-        lambda s: is_within(phis_psis(s).reshape(-1, 2), phis_psis(B), radius),
-        step
-    )
+    # system = tps1.FirstOrderSystem(
+    #     lambda s: is_within(phis_psis(s).reshape(-1, 2), phis_psis(A), radius),
+    #     lambda s: is_within(phis_psis(s).reshape(-1, 2), phis_psis(B), radius),
+    #     step
+    # )
 
     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)),
-        # step_langevin_forward,
-        # step_langevin_backward,
         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)))
     )
 
-    print("A", phis_psis(A))
-    print("B", phis_psis(B))
-
-    filter1 = system.start_state(trajectory)
-    filter2 = 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 = md.load('./files/AD_A_B_500K_initial_trajectory.pdb').xyz.reshape(-1, 1, 66)
     initial_trajectory = [p for p in initial_trajectory]
     save_trajectory(mdtraj_topology, jnp.array(initial_trajectory), f'{savedir}/initial_trajectory.pdb')
 
-    load = False
-    if load:
-        paths = np.load(f'{savedir}/paths.npy', allow_pickle=True)
-        velocities = np.load(f'{savedir}/velocities.npy', allow_pickle=True)
+    if args.resume:
+        paths = [[x for x in p.astype(np.float32)] for p in np.load(f'{savedir}/paths.npy', allow_pickle=True)]
+        velocities = [[v for v in p.astype(np.float32)] for p in np.load(f'{savedir}/velocities.npy', allow_pickle=True)]
         with open(f'{savedir}/stats.json', 'r') as fp:
             statistics = json.load(fp)
+
+        stored = {
+            'trajectories': [initial_trajectory] + paths,
+            'velocities': velocities,
+            'statistics': statistics
+        }
+    else:
+        if os.path.exists(f'{savedir}/paths.npy') and not args.override:
+            print(f"The target directory is not empy."
+                  f"Please use --override to overwrite the existing data or --resume to continue.")
+            exit(1)
+
+        stored = None
+
+    if args.mechanism == 'one-way-shooting':
+        mechanism = tps2.one_way_shooting
+    elif args.mechanism == 'two-way-shooting':
+        mechanism = tps2.two_way_shooting
     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, 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]
-            # 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)
-        except Exception as e:
-            print(e)
-            breakpoint()
+        raise ValueError(f"Unknown mechanism {args.mechanism}")
+
+    try:
+        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,
+                                                           stored=stored)
+        # 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)
+    except Exception as e:
+        print(e)
+        breakpoint()
 
     print(statistics)
     print([len(p) for p in paths])
@@ -368,8 +339,8 @@ def step_langevin_backward(_x, _v, _key):
     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')
+    plot_path_energy(list(zip(paths, velocities)), langevin_log_path_likelihood, reduce=lambda x: x, already_ln=True)
+    plt.ylabel('Path Likelihood')
     plt.savefig(f'{savedir}/path_density.png', bbox_inches='tight')
     plt.show()