Merge pull request #13 from chrisiacovella/reduced-potential-update

Reduced potential update

Examples/LJ_mcmove.py

@@ -0,0 +1,67 @@
+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
+)
+
+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)
+
+from chiron.mcmc import MetropolisDisplacementMove
+
+mc_move = MetropolisDisplacementMove(
+    seed=1234,
+    displacement_sigma=0.01 * unit.nanometer,
+    nr_of_moves=1000,
+    simulation_reporter=reporter,
+)
+
+mc_move.run(sampler_state, thermodynamic_state, nbr_list, True)

chiron/mcmc.py

@@ -375,6 +375,7 @@ def apply(
         thermodynamic_state: ThermodynamicState,
         sampler_state: SamplerState,
         reporter: SimulationReporter,
+        nbr_list=None,
     ):
         """Apply a metropolized move to the sampler state.
 
@@ -388,12 +389,15 @@ def apply(
            The initial sampler state to apply the move to. This is modified.
         reporter: SimulationReporter
               The reporter to write the data to.
+        nbr_list: Neighbor List or Pair List routine,
+            The routine to use to calculate the interacting atoms.
+            Default is None and will use an unoptimized pairlist without PBC
         """
         import jax.numpy as jnp
 
         # Compute initial energy
         initial_energy = thermodynamic_state.get_reduced_potential(
-            sampler_state
+            sampler_state, nbr_list
         )  # NOTE: in kT
         log.debug(f"Initial energy is {initial_energy} kT.")
         # Store initial positions of the atoms that are moved.
@@ -417,7 +421,7 @@ def apply(
                 proposed_positions
             )
         proposed_energy = thermodynamic_state.get_reduced_potential(
-            sampler_state
+            sampler_state, nbr_list
         )  # NOTE: in kT
         # Accept or reject with Metropolis criteria.
         delta_energy = proposed_energy - initial_energy
@@ -589,14 +593,17 @@ def run(
         self,
         sampler_state: SamplerState,
         thermodynamic_state: ThermodynamicState,
+        nbr_list=None,
         progress_bar=True,
     ):
         from tqdm import tqdm
 
         for trials in (
             tqdm(range(self.nr_of_moves)) if progress_bar else range(self.nr_of_moves)
         ):
-            self.apply(thermodynamic_state, sampler_state, self.simulation_reporter)
+            self.apply(
+                thermodynamic_state, sampler_state, self.simulation_reporter, nbr_list
+            )
             if trials % 100 == 0:
                 log.debug(f"Acceptance rate: {self.n_accepted / self.n_proposed}")
                 if self.simulation_reporter is not None:

chiron/states.py

@@ -30,9 +30,7 @@ def __init__(
         # NOTE: all units are internally in the openMM units system as documented here:
         # http://docs.openmm.org/latest/userguide/theory/01_introduction.html#units
         if not isinstance(x0, unit.Quantity):
-            raise TypeError(
-                f"x0 must be a unit.Quantity, got {type(x0)} instead."
-            )
+            raise TypeError(f"x0 must be a unit.Quantity, got {type(x0)} instead.")
         if velocities is not None and not isinstance(velocities, unit.Quantity):
             raise TypeError(
                 f"velocities must be a unit.Quantity, got {type(velocities)} instead."
@@ -42,17 +40,13 @@ def __init__(
                 try:
                     box_vectors = self._convert_from_openmm_box(box_vectors)
                 except:
-                    raise TypeError(
-                        f"Unable to parse box_vectors {box_vectors}."
-                    )
+                    raise TypeError(f"Unable to parse box_vectors {box_vectors}.")
             else:
                 raise TypeError(
                     f"box_vectors must be a unit.Quantity or openMM box, got {type(box_vectors)} instead."
                 )
         if not x0.unit.is_compatible(unit.nanometer):
-            raise ValueError(
-                f"x0 must have units of distance, got {x0.unit} instead."
-            )
+            raise ValueError(f"x0 must have units of distance, got {x0.unit} instead.")
         if velocities is not None and not velocities.unit.is_compatible(
             unit.nanometer / unit.picosecond
         ):
@@ -75,7 +69,6 @@ def __init__(
         self._box_vectors = box_vectors
         self._distance_unit = unit.nanometer
 
-
     @property
     def x0(self) -> jnp.array:
         return self._convert_to_jnp(self._x0)
@@ -112,13 +105,14 @@ def _convert_to_jnp(self, array: unit.Quantity) -> jnp.array:
         array_ = array.value_in_unit_system(unit.md_unit_system)
         return jnp.array(array_)
 
-    def _convert_from_openmm_box(self, openmm_box_vectors:List)->unit.Quantity:
-
+    def _convert_from_openmm_box(self, openmm_box_vectors: List) -> unit.Quantity:
         box_vec = []
         for i in range(0, 3):
             layer = []
             for j in range(0, 3):
-                layer.append(openmm_box_vectors[i][j].value_in_unit(openmm_box_vectors[0].unit))
+                layer.append(
+                    openmm_box_vectors[i][j].value_in_unit(openmm_box_vectors[0].unit)
+                )
             box_vec.append(layer)
         return unit.Quantity(jnp.array(box_vec), openmm_box_vectors[0].unit)
 
@@ -148,13 +142,42 @@ def __init__(
         pressure: Optional[unit.Quantity] = None,
     ):
         self.potential = potential
-        self.temperature = temperature
-        if temperature:
-            self.beta = 1.0 / (
-                unit.BOLTZMANN_CONSTANT_kB * (self.temperature * unit.kelvin)
+
+        if temperature is not None and not isinstance(temperature, unit.Quantity):
+            raise TypeError(
+                f"temperature must be a unit.Quantity, got {type(temperature)} instead."
             )
+        elif temperature is not None:
+            if not temperature.unit.is_compatible(unit.kelvin):
+                raise ValueError(
+                    f"temperature must have units of temperature, got {temperature.unit} instead."
+                )
+
+        if volume is not None and not isinstance(volume, unit.Quantity):
+            raise TypeError(
+                f"volume must be a unit.Quantity, got {type(volume)} instead."
+            )
+        elif volume is not None:
+            if not volume.unit.is_compatible(unit.nanometer**3):
+                raise ValueError(
+                    f"volume must have units of distance**3, got {volume.unit} instead."
+                )
+        if pressure is not None and not isinstance(pressure, unit.Quantity):
+            raise TypeError(
+                f"pressure must be a unit.Quantity, got {type(pressure)} instead."
+            )
+        elif pressure is not None:
+            if not pressure.unit.is_compatible(unit.atmosphere):
+                raise ValueError(
+                    f"pressure must have units of pressure, got {pressure.unit} instead."
+                )
+
+        self.temperature = temperature
+        if temperature is not None:
+            self.beta = 1.0 / (unit.BOLTZMANN_CONSTANT_kB * (self.temperature))
         else:
             self.beta = None
+
         self.volume = volume
         self.pressure = pressure
 
@@ -213,14 +236,18 @@ def are_states_compatible(cls, state1, state2):
         """
         pass
 
-    def get_reduced_potential(self, sampler_state: SamplerState) -> float:
+    def get_reduced_potential(
+        self, sampler_state: SamplerState, nbr_list=None
+    ) -> float:
         """
         Compute the reduced potential for the given sampler state.
 
         Parameters
         ----------
         sampler_state : SamplerState
             The sampler state for which to compute the reduced potential.
+        nbr_list : NeighborList or PairList, optional
+            The neighbor list or pair list routine to use for calculating the reduced potential.
 
         Returns
         -------
@@ -243,7 +270,8 @@ def get_reduced_potential(self, sampler_state: SamplerState) -> float:
         log.debug(f"sample state: {sampler_state.x0}")
         reduced_potential = (
             unit.Quantity(
-                self.potential.compute_energy(sampler_state.x0), unit.kilojoule_per_mole
+                self.potential.compute_energy(sampler_state.x0, nbr_list),
+                unit.kilojoule_per_mole,
             )
         ) / unit.AVOGADRO_CONSTANT_NA
         log.debug(f"reduced potential: {reduced_potential}")

chiron/tests/test_integrators.py

@@ -1,19 +1,21 @@
 import pytest
 
+
 @pytest.fixture(scope="session")
 def prep_temp_dir(tmpdir_factory):
     """Create a temporary directory for the test."""
     tmpdir = tmpdir_factory.mktemp("test_langevin")
     return tmpdir
 
+
 def test_langevin_dynamics(prep_temp_dir, provide_testsystems_and_potentials):
     """
     Test the Langevin integrator with a set of test systems.
 
     This function initializes openmmtools.testsystems,
     sets up a potential, and uses the Langevin integrator to run the simulation.
     """
-    from openmm.unit import kelvin
+    from openmm import unit
 
     i = 0
     for testsystem, potential in provide_testsystems_and_potentials:
@@ -24,7 +26,7 @@ def test_langevin_dynamics(prep_temp_dir, provide_testsystems_and_potentials):
         from chiron.states import SamplerState, ThermodynamicState
 
         thermodynamic_state = ThermodynamicState(
-            potential=potential, temperature=300 * kelvin
+            potential=potential, temperature=300 * unit.kelvin
         )
 
         sampler_state = SamplerState(testsystem.positions)

chiron/tests/test_mcmc.py

@@ -101,7 +101,9 @@ def test_sample_from_harmonic_osciallator_with_MCMC_classes_and_LangevinDynamics
     from chiron.states import ThermodynamicState, SamplerState
 
     thermodynamic_state = ThermodynamicState(
-        harmonic_potential, temperature=300, volume=30 * (unit.angstrom**3)
+        harmonic_potential,
+        temperature=300 * unit.kelvin,
+        volume=30 * (unit.angstrom**3),
     )
     sampler_state = SamplerState(ho.positions)
 
@@ -151,7 +153,9 @@ def test_sample_from_harmonic_osciallator_with_MCMC_classes_and_MetropolisDispla
     from chiron.states import ThermodynamicState, SamplerState
 
     thermodynamic_state = ThermodynamicState(
-        harmonic_potential, temperature=300, volume=30 * (unit.angstrom**3)
+        harmonic_potential,
+        temperature=300 * unit.kelvin,
+        volume=30 * (unit.angstrom**3),
     )
     sampler_state = SamplerState(ho.positions)
 
@@ -204,7 +208,9 @@ def test_sample_from_harmonic_osciallator_array_with_MCMC_classes_and_Metropolis
     from chiron.states import ThermodynamicState, SamplerState
 
     thermodynamic_state = ThermodynamicState(
-        harmonic_potential, temperature=300, volume=30 * (unit.angstrom**3)
+        harmonic_potential,
+        temperature=300 * unit.kelvin,
+        volume=30 * (unit.angstrom**3),
     )
     sampler_state = SamplerState(ho.positions)
 
@@ -231,6 +237,44 @@ def test_sample_from_harmonic_osciallator_array_with_MCMC_classes_and_Metropolis
     sampler.run(n_iterations=2)  # how many times to repeat
 
 
+def test_thermodynamic_state_inputs():
+    from chiron.states import ThermodynamicState
+    from openmm import unit
+    from openmmtools.testsystems import HarmonicOscillatorArray
+
+    ho = HarmonicOscillatorArray()
+
+    # Initalize the potential
+    from chiron.potential import HarmonicOscillatorPotential
+
+    harmonic_potential = HarmonicOscillatorPotential(ho.topology, ho.K)
+
+    with pytest.raises(TypeError):
+        ThermodynamicState(potential=harmonic_potential, temperature=300)
+
+    with pytest.raises(ValueError):
+        ThermodynamicState(
+            potential=harmonic_potential, temperature=300 * unit.angstrom
+        )
+
+    ThermodynamicState(potential=harmonic_potential, temperature=300 * unit.kelvin)
+
+    with pytest.raises(TypeError):
+        ThermodynamicState(potential=harmonic_potential, volume=1000)
+    with pytest.raises(ValueError):
+        ThermodynamicState(potential=harmonic_potential, volume=1000 * unit.kelvin)
+
+    ThermodynamicState(potential=harmonic_potential, volume=1000 * (unit.angstrom**3))
+
+    with pytest.raises(TypeError):
+        ThermodynamicState(potential=harmonic_potential, pressure=100)
+
+    with pytest.raises(ValueError):
+        ThermodynamicState(potential=harmonic_potential, pressure=100 * unit.kelvin)
+
+    ThermodynamicState(potential=harmonic_potential, pressure=100 * unit.atmosphere)
+
+
 def test_sample_from_joint_distribution_of_two_HO_with_local_moves_and_MC_updates():
     # define two harmonic oscillators with different spring constants and equilibrium positions
     # sample from the joint distribution of the two HO using local langevin moves