Move path histogram plotting to end of inference

Author: plainerman

main.py

@@ -215,27 +215,25 @@ def main():
     # In case we have a second order integration scheme, we remove the velocity for plotting
     x_t_det_no_vel = x_t_det[:, :, :system.A.shape[0]]
 
+    key, path_key = jax.random.split(key)
+    x_t_stoch = setup.sample_paths(state_q, x_0, args.dt, args.T, args.BS, path_key)
+    x_t_stoch_no_vel = x_t_stoch[:, :, :system.A.shape[0]]
+    np.save(f'{args.save_dir}/stochastic_paths.npy', x_t_stoch_no_vel)
+
     if system.mdtraj_topology:
         save_trajectory(system.mdtraj_topology, x_t_det_no_vel[0].reshape(1, -1, 3), f'{args.save_dir}/det_0.pdb')
         save_trajectory(system.mdtraj_topology, x_t_det_no_vel[-1].reshape(1, -1, 3), f'{args.save_dir}/det_-1.pdb')
 
+        save_trajectory(system.mdtraj_topology, x_t_stoch_no_vel[0].reshape(1, -1, 3), f'{args.save_dir}/stoch_0.pdb')
+        save_trajectory(system.mdtraj_topology, x_t_stoch_no_vel[-1].reshape(1, -1, 3), f'{args.save_dir}/stoch_-1.pdb')
+
     if system.plot:
         plot_energy(system, [x_t_det_no_vel[0], x_t_det_no_vel[-1]], args.log_plots)
         show_or_save_fig(args.save_dir, 'path_energy_deterministic', args.extension)
 
         system.plot(title='Deterministic Paths', trajectories=x_t_det_no_vel)
         show_or_save_fig(args.save_dir, 'paths_deterministic', args.extension)
 
-    key, path_key = jax.random.split(key)
-    x_t_stoch = setup.sample_paths(state_q, x_0, args.dt, args.T, args.BS, path_key)
-    x_t_stoch_no_vel = x_t_stoch[:, :, :system.A.shape[0]]
-    np.save(f'{args.save_dir}/stochastic_paths.npy', x_t_stoch_no_vel)
-
-    if system.mdtraj_topology:
-        save_trajectory(system.mdtraj_topology, x_t_stoch_no_vel[0].reshape(1, -1, 3), f'{args.save_dir}/stoch_0.pdb')
-        save_trajectory(system.mdtraj_topology, x_t_stoch_no_vel[-1].reshape(1, -1, 3), f'{args.save_dir}/stoch_-1.pdb')
-
-    if system.plot:
         plot_energy(system, [x_t_stoch_no_vel[0], x_t_stoch_no_vel[-1]], args.log_plots)
         show_or_save_fig(args.save_dir, 'path_energy_stochastic', args.extension)