Resolved unit issue in barostat.

Examples/LJ_mcbarostat.py

@@ -0,0 +1,70 @@
+from openmmtools.testsystems import LennardJonesFluid
+
+# Use the LennardJonesFluid example from openmmtools to initialize particle positions and topology
+# For this example, the topology provides the masses for the particles
+# The default LennardJonesFluid example considers the system to be Argon with 39.9 amu
+lj_fluid = LennardJonesFluid(reduced_density=0.1, nparticles=1000)
+
+
+from chiron.potential import LJPotential
+from openmm import unit
+
+# initialize the LennardJones potential in chiron
+#
+sigma = 0.34 * unit.nanometer
+epsilon = 0.238 * unit.kilocalories_per_mole
+cutoff = 3.0 * sigma
+
+lj_potential = LJPotential(
+    lj_fluid.topology, sigma=sigma, epsilon=epsilon, cutoff=cutoff
+)
+
+from chiron.states import SamplerState, ThermodynamicState
+
+# define the sampler state
+sampler_state = SamplerState(
+    x0=lj_fluid.positions, box_vectors=lj_fluid.system.getDefaultPeriodicBoxVectors()
+)
+
+# define the thermodynamic state
+thermodynamic_state = ThermodynamicState(
+    potential=lj_potential,
+    temperature=300 * unit.kelvin,
+    pressure=1.0 * unit.atmosphere,
+)
+
+from chiron.neighbors import NeighborListNsqrd, OrthogonalPeriodicSpace
+
+# define the neighbor list for an orthogonal periodic space
+skin = 0.5 * unit.nanometer
+
+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)
+
+from chiron.reporters import SimulationReporter
+
+# initialize a reporter to save the simulation data
+filename = "test_lj.h5"
+import os
+
+if os.path.isfile(filename):
+    os.remove(filename)
+reporter = SimulationReporter("test_mc_lj.h5", lj_fluid.topology, 1)
+
+print(thermodynamic_state.pressure)
+from chiron.mcmc import MCBarostatMove
+
+barostat_move = MCBarostatMove(
+    seed=1234,
+    volume_max_scale=0.01,
+    nr_of_moves=1000,
+    simulation_reporter=reporter,
+)
+
+barostat_move.run(sampler_state, thermodynamic_state, nbr_list, True)

chiron/mcmc.py

@@ -414,11 +414,27 @@ def apply(
         else:
             initial_box_vectors = None
 
+        # if we are running in NPT, we need to store the initial volume as well
         if thermodynamic_state.pressure is not None:
+            if thermodynamic_state.volume is None:
+                if sampler_state.box_vectors is None:
+                    raise ValueError(
+                        "Cannot run NPT without a volume or box vectors specified"
+                    )
+                # units are nm**3
+                initial_volume = (
+                    initial_box_vectors[0][0]
+                    * initial_box_vectors[1][1]
+                    * initial_box_vectors[2][2]
+                )
+                # set the volume in the thermodynamic state
+                thermodynamic_state.volume = initial_volume
+
             initial_volume = jnp.copy(thermodynamic_state.volume)
 
         log.debug(f"Initial positions are {initial_positions} nm.")
         # Propose perturbed positions. Modifying the reference changes the sampler state.
+        # This will also return box vectors for methodologies that change both position and box
         proposed_positions, proposed_box_vectors = self._propose_positions(
             initial_positions, initial_box_vectors
         )
@@ -431,16 +447,21 @@ def apply(
             sampler_state.x0 = sampler_state.x0.at[jnp.array(atom_subset)].set(
                 proposed_positions
             )
+        # set the box vectors in the sampler state  if they changed
         sampler_state.box_vectors = proposed_box_vectors
 
         if nbr_list is not None:
             # nbr_list box_vectors setter also sets space.box_vectors
             nbr_list.box_vectors = proposed_box_vectors
-            sampler_state.x0 = nbr_list.space.wrap(sampler_state.x0)
 
+            # this will call the space.wrap function
+            sampler_state.x0 = nbr_list.wrap(sampler_state.x0)
+
+            # if the box vectors changed, we need to rebuild the neighbor list
             if jnp.any(initial_box_vectors != proposed_box_vectors):
                 nbr_list.build(sampler_state.x0, sampler_state.box_vectors)
             else:
+                # check to see if any particles have moved too far
                 if nbr_list.check(sampler_state.x0):
                     nbr_list.build(sampler_state.x0, sampler_state.box_vectors)
 
@@ -450,9 +471,16 @@ def apply(
 
         # Accept or reject with Metropolis criteria.
         delta_energy = proposed_energy - initial_energy
-        if thermodynamic_state.pressure is not None:
-            proposed_volume = jnp.copy(thermodynamic_state.volume)
 
+        # if we are running NPT, we need to add in the N*ln(V/V0) term to delta energy
+        if thermodynamic_state.pressure is not None:
+            # units are nm**3
+            proposed_volume = (
+                proposed_box_vectors[0][0]
+                * proposed_box_vectors[1][1]
+                * proposed_box_vectors[2][2]
+            )
+            thermodynamic_state.volume = proposed_volume
             delta_energy += sampler_state.x0.shape[0] * jnp.log(
                 proposed_volume / initial_volume
             )
@@ -489,7 +517,7 @@ def apply(
                 )
             if thermodynamic_state.pressure is not None:
                 thermodynamic_state.volume = initial_volume
-                thermoydnamic_state.box_vectors = initial_box_vectors
+                thermodynamic_state.box_vectors = initial_box_vectors
 
             log.debug(
                 f"Move rejected. Energy change: {delta_energy:.3f} kT. Number of rejected moves: {self.n_proposed - self.n_accepted}."
@@ -756,17 +784,21 @@ def displace_positions(
         # log.debug(f"Number of atoms is {nr_of_atoms}.")
 
         volume = box_vectors[0][0] * box_vectors[1][1] * box_vectors[2][2]
+        log.debug(f"Volume is {volume}.")
         volume_scale_factor = volume_max_scale * volume
-
+        log.debug(f"Volume is {volume}.")
         delta_volume = jrandom.uniform(subkey, minval=-1) * volume_scale_factor
         log.debug(f"Delta volume is {delta_volume}.")
 
         new_volume = volume + delta_volume
-        length_scale = jnp.power(new_volume / delta_volume, 1.0 / 3.0)
+        log.debug(f"New volume is {new_volume}.")
+        length_scaled = jnp.power(new_volume / volume, 1.0 / 3.0)
+
+        log.debug(f"Length scaled is {length_scaled}.")
 
-        updated_position = positions * length_scale
+        updated_position = positions * length_scaled
         box_scaling = jnp.array(
-            [[length_scale, 0, 0], [0, length_scale, 0], [0, 0, length_scale]]
+            [[length_scaled, 0, 0], [0, length_scaled, 0], [0, 0, length_scaled]]
         )
         scaled_box = box_vectors * box_scaling
         return updated_position, scaled_box

chiron/neighbors.py

@@ -369,6 +369,23 @@ def box_vectors(self, box_vectors: jnp.array) -> None:
         self._box_vectors = box_vectors
         self.space.box_vectors = box_vectors
 
+    def wrap(self, coordinates: jnp.array) -> jnp.array:
+        """
+        Wrap the coordinates of the system.
+
+        Parameters
+        ----------
+        coordinates: jnp.array
+            Coordinates of the system
+
+        Returns
+        -------
+        jnp.array
+            Wrapped coordinates of the system
+
+        """
+        return self.space.wrap(coordinates)
+
 
 class NeighborListNsqrd(PairsBase):
     """

chiron/states.py

@@ -283,12 +283,15 @@ def get_reduced_potential(
         ) / unit.AVOGADRO_CONSTANT_NA
         log.debug(f"reduced potential: {reduced_potential}")
         if self.pressure is not None:
-            self.volume = (
-                sampler_state.box_vectors[0][0]
-                * sampler_state.box_vectors[1][1]
-                * sampler_state.box_vectors[2][2]
-            )
-            reduced_potential += self.pressure * self.volume
+            # in case volume is not set, calculate from the box vectors
+            if self.volume is None:
+                self.volume = (
+                    sampler_state.box_vectors[0][0]
+                    * sampler_state.box_vectors[1][1]
+                    * sampler_state.box_vectors[2][2]
+                )
+
+            reduced_potential += self.pressure * self.volume * unit.nanometer**3
 
         return self.beta * reduced_potential