Merge pull request #15 from choderalab/analyzer

Adding offline mbar analysis

chiron/integrators.py

@@ -29,6 +29,7 @@ def __init__(
         collision_rate=1.0 / unit.picoseconds,
         save_frequency: int = 100,
         reporter: Optional[SimulationReporter] = None,
+        save_traj_in_memory: bool = False,
     ) -> None:
         """
         Initialize the LangevinIntegrator object.
@@ -56,7 +57,8 @@ def __init__(
             log.info(f"Using reporter {reporter} saving to {reporter.filename}")
             self.reporter = reporter
         self.save_frequency = save_frequency
-
+        self.save_traj_in_memory = save_traj_in_memory
+        self.traj = []
         self.velocities = None
 
     def set_velocities(self, vel: unit.Quantity) -> None:
@@ -109,10 +111,10 @@ def run(
         temperature = thermodynamic_state.temperature
         x0 = sampler_state.x0
 
-        log.info("Running Langevin dynamics")
-        log.info(f"n_steps = {n_steps}")
-        log.info(f"temperature = {temperature}")
-        log.info(f"Using seed: {key}")
+        log.debug("Running Langevin dynamics")
+        log.debug(f"n_steps = {n_steps}")
+        log.debug(f"temperature = {temperature}")
+        log.debug(f"Using seed: {key}")
 
         kbT_unitless = (self.kB * temperature).value_in_unit_system(unit.md_unit_system)
         mass_unitless = jnp.array(mass.value_in_unit_system(unit.md_unit_system))[
@@ -169,6 +171,9 @@ def run(
             if step % self.save_frequency == 0:
                 # log.debug(f"Saving at step {step}")
                 # check if reporter is attribute of the class
+                # log.debug(f"step {step} energy {potential.compute_energy(x, nbr_list)}")
+                # log.debug(f"step {step} force {F}")
+
                 if hasattr(self, "reporter") and self.reporter is not None:
                     d = {
                         "traj": x,
@@ -180,6 +185,8 @@ def run(
 
                     # log.debug(d)
                     self.reporter.report(d)
+                if self.save_traj_in_memory:
+                    self.traj.append(x)
 
         log.debug("Finished running Langevin dynamics")
         # save the final state of the simulation in the sampler_state object

chiron/mcmc.py

@@ -71,6 +71,7 @@ def __init__(
         simulation_reporter: Optional[SimulationReporter] = None,
         nr_of_steps=1_000,
         seed: int = 1234,
+        save_traj_in_memory: bool = False,
     ):
         """
         Initialize the LangevinDynamicsMove with a molecular system.
@@ -88,13 +89,15 @@ def __init__(
         self.stepsize = stepsize
         self.collision_rate = collision_rate
         self.simulation_reporter = simulation_reporter
-
+        self.save_traj_in_memory = save_traj_in_memory
+        self.traj = []
         from chiron.integrators import LangevinIntegrator
 
         self.integrator = LangevinIntegrator(
             stepsize=self.stepsize,
             collision_rate=self.collision_rate,
             reporter=self.simulation_reporter,
+            save_traj_in_memory=save_traj_in_memory,
         )
 
     def run(
@@ -122,6 +125,9 @@ def run(
             n_steps=self.nr_of_moves,
             key=self.key,
         )
+        if self.save_traj_in_memory:
+            self.traj.append(self.integrator.traj)
+            self.integrator.traj = []
 
 
 class MCMove(MCMCMove):