added in reset of box vectors and volume if move is rejected and pres…

…sure is set.
choderalab · Dec 29, 2023 · 5fa497a · 5fa497a
1 parent c1ea3fc
commit 5fa497a
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/chiron/mcmc.py b/chiron/mcmc.py
@@ -487,6 +487,10 @@ def apply(
                 sampler_state.x0 = sampler_state.x0.at[jnp.array([atom_subset])].set(
                     initial_positions
                 )
+            if thermodynamic_state.pressure is not None:
+                thermodynamic_state.volume = initial_volume
+                thermoydnamic_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}."
             )