Save pdb trajectory

Author: plainerman

main.py

@@ -1,4 +1,6 @@
 from argparse import ArgumentParser
+
+from utils.animation import save_trajectory
 from utils.args import parse_args, str2bool
 from systems import System
 import matplotlib.pyplot as plt
@@ -203,11 +205,14 @@ def main():
     x_0 += args.base_sigma * eps
 
     x_t_det = setup.sample_paths(state_q, x_0, args.dt, args.T, args.BS, None)
+    # 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]]
 
-    if system.plot:
-        # 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]]
+    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')
 
+    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)
 
@@ -216,10 +221,13 @@ def main():
 
     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]]
 
-    if system.plot:
-        x_t_stoch_no_vel = x_t_stoch[:, :, :system.A.shape[0]]
+    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)
 

systems.py

@@ -16,7 +16,7 @@
 
 class System:
     def __init__(self, U: Callable[[ArrayLike], ArrayLike], A: ArrayLike, B: ArrayLike, mass: ArrayLike, plot,
-                 force_clip: float):
+                 force_clip: float, mdtraj_topology: Optional[md.Topology] = None):
         assert A.shape == B.shape == mass.shape
 
         self.U = jax.jit(U)
@@ -28,6 +28,7 @@ def __init__(self, U: Callable[[ArrayLike], ArrayLike], A: ArrayLike, B: ArrayLi
         self.mass = mass
 
         self.plot = plot
+        self.mdtraj_topology = mdtraj_topology
 
     @classmethod
     def from_name(cls, name: str, force_clip: float) -> Self:
@@ -109,4 +110,4 @@ def U(_x):
         else:
             raise ValueError(f"Unknown cv: {cv}")
 
-        return cls(U, A, B, mass, plot, force_clip)
+        return cls(U, A, B, mass, plot, force_clip, md.Topology.from_openmm(A_pdb.topology))