Added simple LJ langevin example.

Examples/LJ_langevin.py

@@ -0,0 +1,82 @@
+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
+)
+# 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_lj.h5", lj_fluid.topology, 1)
+
+from chiron.integrators import LangevinIntegrator
+
+# initialize the Langevin integrator
+integrator = LangevinIntegrator(reporter=reporter, save_frequency=100)
+
+integrator.run(
+    sampler_state,
+    thermodynamic_state,
+    n_steps=2000,
+    nbr_list=nbr_list,
+    progress_bar=True,
+)
+
+import h5py
+
+# read the data from the reporter
+with h5py.File("test_lj.h5", "r") as f:
+    energies = f["energy"][:]
+    steps = f["step"][:]
+
+
+# plot the energy
+import matplotlib.pyplot as plt
+
+plt.plot(steps, energies)
+plt.xlabel("Step (fs)")
+plt.ylabel("Energy (kj/mol)")
+plt.show()

chiron/potential.py

@@ -25,7 +25,8 @@ def compute_force(self, positions, nbr_list=None) -> jnp.ndarray:
 
     def compute_pairlist(self, positions, cutoff) -> jnp.array:
         """
-        Compute the pairlist for a given set of positions and a cutoff distance.
+        Compute the pairlist for a given set of positions and a cutoff distance
+        without using periodic boundary conditions or any specific optimizations.
 
         Parameters
         ----------
@@ -89,13 +90,21 @@ def __init__(
         if not isinstance(topology, Topology):
             if not isinstance(topology, property):
                 if topology is not None:
-                    raise TypeError(f"Topology must be a Topology object or None, type(topology) = {type(topology)}")
+                    raise TypeError(
+                        f"Topology must be a Topology object or None, type(topology) = {type(topology)}"
+                    )
         if not isinstance(sigma, unit.Quantity):
-            raise TypeError(f"sigma must be a unit.Quantity, type(sigma) = {type(sigma)}")
+            raise TypeError(
+                f"sigma must be a unit.Quantity, type(sigma) = {type(sigma)}"
+            )
         if not isinstance(epsilon, unit.Quantity):
-            raise TypeError(f"epsilon must be a unit.Quantity, type(epsilon) = {type(epsilon)}")
+            raise TypeError(
+                f"epsilon must be a unit.Quantity, type(epsilon) = {type(epsilon)}"
+            )
         if not isinstance(cutoff, unit.Quantity):
-            raise TypeError(f"cutoff must be a unit.Quantity, type(cutoff) = {type(cutoff)}")
+            raise TypeError(
+                f"cutoff must be a unit.Quantity, type(cutoff) = {type(cutoff)}"
+            )
 
         if not sigma.unit.is_compatible(unit.angstrom):
             raise ValueError(f"sigma must have units of distance, got {sigma.unit}")
@@ -104,8 +113,6 @@ def __init__(
         if not cutoff.unit.is_compatible(unit.nanometer):
             raise ValueError(f"cutoff must have units of distance, got {cutoff.unit}")
 
-
-
         self.sigma = sigma.value_in_unit_system(
             unit.md_unit_system
         )  # The distance at which the potential is zero
@@ -141,22 +148,17 @@ def _compute_energy_masked(self, distance, mask):
         )
         return energy.sum()
 
-    def compute_energy(
-        self,
-        positions: jnp.array,
-        nbr_list=None,
-    ):
+    def compute_energy(self, positions: jnp.array, nbr_list=None, debug_mode=False):
         """
         Compute the LJ energy.
 
         Parameters
         ----------
         positions : jnp.array
             The positions of the particles in the system
-        nbr_list : NeighborList, optional
-            Instance of the neighborlist class to use. By default, set to None, which will use an N^2 pairlist
-        shift : bool, optional
-            Whether to shift the potential energy at the cutoff, by default False
+        nbr_list : NeighborList, default=None
+            Instance of a neighbor list or pair list class to use.
+            If None, an unoptimized N^2 pairlist will be used without PBC conditions.
         Returns
         -------
         potential_energy : float
@@ -166,6 +168,9 @@ def compute_energy(
         # Compute the pair distances and displacement vectors
 
         if nbr_list is None:
+            log.debug(
+                "nbr_list is None, computing pairlist using N^2 method without PBC."
+            )
             # Compute the pairlist for a given set of positions and a cutoff distance
             # Note in this case, we do not need the pairs or displacement vectors
             # Since we already calculate the distance in the pairlist computation
@@ -194,8 +199,10 @@ def compute_energy(
                 raise ValueError("Neighborlist must be built before use")
 
             # ensure that the cutoff in the neighbor list is the same as the cutoff in the potential
-            if  nbr_list.cutoff != self.cutoff:
-                raise ValueError(f"Neighborlist cutoff ({nbr_list.cutoff}) must be the same as the potential cutoff ({self.cutoff})")
+            if nbr_list.cutoff != self.cutoff:
+                raise ValueError(
+                    f"Neighborlist cutoff ({nbr_list.cutoff}) must be the same as the potential cutoff ({self.cutoff})"
+                )
 
             n_neighbors, mask, dist, displacement_vectors = nbr_list.calculate(
                 positions
@@ -273,9 +280,13 @@ def __init__(
         U0: unit.Quantity = 0.0 * unit.kilocalories_per_mole,
     ):
         if not isinstance(topology, Topology):
-            if not isinstance(topology, property): #importing from the topology from the model results in it being a property object
+            if not isinstance(
+                topology, property
+            ):  # importing from the topology from the model results in it being a property object
                 if topology is not None:
-                    raise TypeError(f"Topology must be a Topology object or None, type(topology) = {type(topology)}")
+                    raise TypeError(
+                        f"Topology must be a Topology object or None, type(topology) = {type(topology)}"
+                    )
         if not isinstance(k, unit.Quantity):
             raise TypeError(f"k must be a unit.Quantity, type(k) = {type(k)}")
         if not isinstance(x0, unit.Quantity):
@@ -284,11 +295,17 @@ def __init__(
             raise TypeError(f"U0 must be a unit.Quantity, type(U0) = {type(U0)}")
 
         if not k.unit.is_compatible(unit.kilocalories_per_mole / unit.angstrom**2):
-            raise ValueError(f"k must be a unit.Quantity with units of energy per distance squared, k.unit = {k.unit}")
+            raise ValueError(
+                f"k must be a unit.Quantity with units of energy per distance squared, k.unit = {k.unit}"
+            )
         if not x0.unit.is_compatible(unit.angstrom):
-            raise ValueError(f"x0 must be a unit.Quantity with units of distance, x0.unit = {x0.unit}")
+            raise ValueError(
+                f"x0 must be a unit.Quantity with units of distance, x0.unit = {x0.unit}"
+            )
         if not U0.unit.is_compatible(unit.kilocalories_per_mole):
-            raise ValueError(f"U0 must be a unit.Quantity with units of energy, U0.unit = {U0.unit}")
+            raise ValueError(
+                f"U0 must be a unit.Quantity with units of energy, U0.unit = {U0.unit}"
+            )
 
         log.info("Initializing HarmonicOscillatorPotential")
         log.info(f"k = {k}")