Refactor tps_baseline.py

Author: plainerman

README.md

@@ -28,7 +28,10 @@ python eval/evaluate_mueller.py
 and 
 
 ```bash
-python tps_baseline.py
+python tps_baseline.py --mechanism two-way-shooting --num_paths 1000 --states phi-psi
+# num_steps compiles multiple MD steps into a single one, making sampling faster. But this makes startup longer. Only really worth it for long running simulations
+python tps_baseline.py --mechanism two-way-shooting --num_paths 1000 --fixed_length 1000 --states phi-psi --num_steps 50
+python tps_baseline.py --mechanism two-way-shooting --num_paths 1000 --states rmsd
 python eval/evaluate_tps.py
 ```
 

eval/path_metrics.py

@@ -14,3 +14,17 @@ def plot_path_energy(paths, U, reduce=jnp.max, add=0, already_ln=False, **kwargs
         plt.semilogy(jnp.arange(0, len(reduced), 1), reduced, **kwargs)
 
     return reduced
+
+
+def plot_iterative_min_max_energy(paths, U, potential_calls):
+    reduced = jnp.array([jnp.max(U(path)) for path in tqdm(paths)])
+
+    iterative_min = [reduced[0]]
+    for c in reduced[1:]:
+        iterative_min.append(min(c, iterative_min[-1]))
+
+    plt.xlabel('Number of potential calls')
+    plt.ylabel('Minimum energy of best path so far')
+    plt.semilogy(jnp.cumsum(jnp.array(potential_calls)), iterative_min)
+
+    return iterative_min

systems.py

@@ -105,8 +105,6 @@ def U(_x):
                            yticklabels=[r'$-\pi$', r'$-\frac {\pi} {2}$', '0', r'$\frac {\pi} {2}$', r'$\pi$'],
                            square=True, periodic=True,
                            )
-
-
         else:
             raise ValueError(f"Unknown cv: {cv}")
 

tps/second_order.py

@@ -158,7 +158,7 @@ def mcmc_shooting(system, proposal, initial_trajectory, num_paths, dt, key, fixe
     trajectories = [initial_trajectory]
     velocities = []
     statistics = {
-        'num_force_evaluations': 0,
+        'num_force_evaluations': [],
         'num_tries': 0,
         'num_metropolis_rejected': 0,
         'warmup': warmup,
@@ -193,7 +193,8 @@ def mcmc_shooting(system, proposal, initial_trajectory, num_paths, dt, key, fixe
                 # 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,
+                                                                 velocities[traj_idx - 1] if len(
+                                                                     trajectories) > 1 else None,
                                                                  fixed_length, dt, ikey)
                 num_force_evaluations += len(new_trajectory) - 1
 
@@ -206,7 +207,7 @@ def mcmc_shooting(system, proposal, initial_trajectory, num_paths, dt, key, fixe
                     # 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_force_evaluations'].append(num_force_evaluations)
                     statistics['num_metropolis_rejected'] += num_metropolis_rejected
                     num_tries = 0
                     num_force_evaluations = 0