Merge pull request #11 from chrisiacovella/update_neighboring

Update neighboring class output, additional input checking, more tests, simple LJ Langevin example system.

Examples/LJ_langevin.py

@@ -39,8 +39,6 @@
 nbr_list = NeighborListNsqrd(
     OrthogonalPeriodicSpace(), cutoff=cutoff, skin=skin, n_max_neighbors=180
 )
-from chiron.neighbors import PairList
-
 
 # build the neighbor list from the sampler state
 nbr_list.build_from_state(sampler_state)

chiron/integrators.py

@@ -12,6 +12,17 @@
 
 
 class LangevinIntegrator:
+    """
+    Langevin dynamics integrator for molecular dynamics simulation using the BAOAB splitting scheme [1].
+
+    References:
+    [1] Benedict Leimkuhler, Charles Matthews;
+        Robust and efficient configurational molecular sampling via Langevin dynamics.
+        J. Chem. Phys. 7 May 2013; 138 (17): 174102. https://doi.org/10.1063/1.4802990
+
+
+    """
+
     def __init__(
         self,
         stepsize=1.0 * unit.femtoseconds,
@@ -25,21 +36,28 @@ def __init__(
         Parameters
         ----------
         stepsize : unit.Quantity, optional
-            Time step size for the integration.
+            Time step of integration with units of time. Default is 1.0 * unit.femtoseconds.
         collision_rate : unit.Quantity, optional
-            Collision rate for the Langevin dynamics.
+            Collision rate for the Langevin dynamics, with units 1/time. Default is 1.0 / unit.picoseconds.
+        save_frequency : int, optional
+            Frequency of saving the simulation data. Default is 100.
+        reporter : SimulationReporter, optional
+            Reporter object for saving the simulation data. Default is None.
         """
 
         self.kB = unit.BOLTZMANN_CONSTANT_kB * unit.AVOGADRO_CONSTANT_NA
         log.info(f"stepsize = {stepsize}")
         log.info(f"collision_rate = {collision_rate}")
+        log.info(f"save_frequency = {save_frequency}")
+
         self.stepsize = stepsize
         self.collision_rate = collision_rate
         if reporter is not None:
             log.info(f"Using reporter {reporter} saving to {reporter.filename}")
             self.reporter = reporter
         self.save_frequency = save_frequency
 
+        self.velocities = None
     def set_velocities(self, vel: unit.Quantity) -> None:
         """
         Set the initial velocities for the Langevin Integrator.
@@ -73,6 +91,8 @@ def run(
             Number of simulation steps to perform.
         key : jax.random.PRNGKey, optional
             Random key for generating random numbers.
+        nbr_list : NeighborListNsqrd, optional
+            Neighbor list for the system.
         progress_bar : bool, optional
             Flag indicating whether to display a progress bar during integration.
 
@@ -85,7 +105,6 @@ def run(
 
         self.box_vectors = sampler_state.box_vectors
         self.progress_bar = progress_bar
-        self.velocities = None
         temperature = thermodynamic_state.temperature
         x0 = sampler_state.x0
 

chiron/mcmc.py

@@ -412,6 +412,7 @@ def apply(
         log.debug(f"Initial positions are {initial_positions} nm.")
         # Propose perturbed positions. Modifying the reference changes the sampler state.
         proposed_positions = self._propose_positions(initial_positions)
+
         log.debug(f"Proposed positions are {proposed_positions} nm.")
         # Compute the energy of the proposed positions.
         if atom_subset is None:
@@ -420,6 +421,10 @@ def apply(
             sampler_state.x0 = sampler_state.x0.at[jnp.array(atom_subset)].set(
                 proposed_positions
             )
+        if nbr_list is not None:
+            if nbr_list.check(sampler_state.x0):
+                nbr_list.build(sampler_state.x0, sampler_state.box_vectors)
+
         proposed_energy = thermodynamic_state.get_reduced_potential(
             sampler_state, nbr_list
         )  # NOTE: in kT