From 038e24658adf1c93ca200d378bb8bac474902dca Mon Sep 17 00:00:00 2001
From: chrisiacovella <chris.iacovella@gmail.com>
Date: Wed, 27 Dec 2023 16:28:26 -0800
Subject: [PATCH] Displacement moves now accept nbr_list.

---
 Examples/LJ_mcmove.py |  6 +++---
 chiron/mcmc.py        | 13 ++++++++++---
 chiron/states.py      |  9 +++++++--
 3 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/Examples/LJ_mcmove.py b/Examples/LJ_mcmove.py
index 44b6d3d..0ed407e 100644
--- a/Examples/LJ_mcmove.py
+++ b/Examples/LJ_mcmove.py
@@ -59,9 +59,9 @@
 
 mc_move = MetropolisDisplacementMove(
     seed=1234,
-    displacement_sigma=0.1 * unit.nanometer,
-    nr_of_moves=10,
+    displacement_sigma=0.01 * unit.nanometer,
+    nr_of_moves=1000,
     simulation_reporter=reporter,
 )
 
-mc_move.run(sampler_state, thermodynamic_state, True)
+mc_move.run(sampler_state, thermodynamic_state, nbr_list, True)
diff --git a/chiron/mcmc.py b/chiron/mcmc.py
index 1189d80..b649afa 100644
--- a/chiron/mcmc.py
+++ b/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,6 +593,7 @@ def run(
         self,
         sampler_state: SamplerState,
         thermodynamic_state: ThermodynamicState,
+        nbr_list=None,
         progress_bar=True,
     ):
         from tqdm import tqdm
@@ -596,7 +601,9 @@ def run(
         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:
diff --git a/chiron/states.py b/chiron/states.py
index f533e07..edfae43 100644
--- a/chiron/states.py
+++ b/chiron/states.py
@@ -236,7 +236,9 @@ 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.
 
@@ -244,6 +246,8 @@ def get_reduced_potential(self, sampler_state: SamplerState) -> float:
         ----------
         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
         -------
@@ -266,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}")