choderalab · mikemhenry · Aug 2, 2022 · Jul 27, 2022 · Jul 27, 2022 · Jul 27, 2022
diff --git a/openmmtools/mcmc.py b/openmmtools/mcmc.py
@@ -1056,6 +1056,9 @@ class LangevinDynamicsMove(BaseIntegratorMove):
     reassign_velocities : bool, optional
         If True, the velocities will be reassigned from the Maxwell-Boltzmann
         distribution at the beginning of the move (default is False).
+    constraint_tolerance : float, optional
+        Fraction of the constrained distance within which constraints are maintained for the
+        integrator (default is 1e-8).
 
     Attributes
     ----------
@@ -1069,6 +1072,9 @@ class LangevinDynamicsMove(BaseIntegratorMove):
     reassign_velocities : bool
         If True, the velocities will be reassigned from the Maxwell-Boltzmann
         distribution at the beginning of the move.
+    constraint_tolerance : float
+        Fraction of the constrained distance within which constraints are maintained for the
+        integrator.
 
     Examples
     --------
@@ -1112,12 +1118,13 @@ class LangevinDynamicsMove(BaseIntegratorMove):
     """
 
     def __init__(self, timestep=1.0*unit.femtosecond, collision_rate=10.0/unit.picoseconds,
-                 n_steps=1000, reassign_velocities=False, **kwargs):
+                 n_steps=1000, reassign_velocities=False, constraint_tolerance=1e-8, **kwargs):
         super(LangevinDynamicsMove, self).__init__(n_steps=n_steps,
                                                    reassign_velocities=reassign_velocities,
                                                    **kwargs)
         self.timestep = timestep
         self.collision_rate = collision_rate
+        self.constraint_tolerance = constraint_tolerance
 
     def apply(self, thermodynamic_state, sampler_state, context_cache=None):
         """Apply the Langevin dynamics MCMC move.
@@ -1143,17 +1150,21 @@ def __getstate__(self):
         serialization = super(LangevinDynamicsMove, self).__getstate__()
         serialization['timestep'] = self.timestep
         serialization['collision_rate'] = self.collision_rate
+        serialization['constraint_tolerance'] = self.constraint_tolerance
         return serialization
 
     def __setstate__(self, serialization):
         super(LangevinDynamicsMove, self).__setstate__(serialization)
         self.timestep = serialization['timestep']
         self.collision_rate = serialization['collision_rate']
+        self.constraint_tolerance = serialization['constraint_tolerance']
 
     def _get_integrator(self, thermodynamic_state):
         """Implement BaseIntegratorMove._get_integrator()."""
-        return openmm.LangevinMiddleIntegrator(thermodynamic_state.temperature,
-                                               self.collision_rate, self.timestep)
+        integrator = openmm.LangevinMiddleIntegrator(thermodynamic_state.temperature,
+                                                     self.collision_rate, self.timestep)
+        integrator.setConstraintTolerance(self.constraint_tolerance)
+        return integrator
 
 
 class LangevinSplittingDynamicsMove(LangevinDynamicsMove):

diff --git a/openmmtools/tests/test_mcmc.py b/openmmtools/tests/test_mcmc.py
@@ -505,6 +505,40 @@ def test_langevin_splitting_move():
         sampler = MCMCSampler(thermodynamic_state, sampler_state, move=move)
         sampler.run(1)
 
+def test_langevin_dynamics_move_constraint_tolerance():
+    """Test constraint tolerance is properly set in LangevinDynamicsMove integrator."""
+    testsystem = testsystems.AlanineDipeptideVacuum()
+    thermodynamic_state = ThermodynamicState(testsystem.system, 300 * unit.kelvin)
+    # Check for default tolerance
+    default_move = LangevinDynamicsMove()
+    default_constraint_tolerance = 1e-8
+    move_tolerance = default_move.constraint_tolerance
+    assert move_tolerance == default_constraint_tolerance, f"LangevinDynamicsMove tolerance, {move_tolerance}, is" \
+                                                           f" not the same as the expected default tolerance," \
+                                                           f" {default_constraint_tolerance}."
+    default_integrator = default_move._get_integrator(thermodynamic_state)
+    default_integrator_tolerance = default_integrator.getConstraintTolerance()
+    assert default_integrator_tolerance == default_constraint_tolerance, f"LangevinDynamicsMove integrator tolerance," \
+                                                                         f" {default_integrator_tolerance}, is not " \
+                                                                         f"the same as the expected default " \
+                                                                         f"tolerance, {default_constraint_tolerance}."
+    # Now we change the tolerance in initializer and check
+    new_constraint_tolerance = 1e-5
+    new_move = LangevinDynamicsMove(constraint_tolerance=new_constraint_tolerance)
+    new_integrator = new_move._get_integrator(thermodynamic_state)
+    new_integrator_tolerance = new_integrator.getConstraintTolerance()
+    assert new_integrator_tolerance == new_constraint_tolerance, f"LangevinDynamicsMove integrator tolerance," \
+                                                                 f" {new_integrator_tolerance}, is not the same as" \
+                                                                 f" the specified value of {new_constraint_tolerance}."
+    # Test by changing public attribute
+    constraint_tolerance = 1e-7
+    move = LangevinDynamicsMove()  # create default move
+    move.constraint_tolerance = constraint_tolerance  # change the public attribute
+    integrator = move._get_integrator(thermodynamic_state)
+    integrator_tolerance = integrator.getConstraintTolerance()
+    assert integrator_tolerance == constraint_tolerance, f"LangevinDynamicsMove integrator tolerance," \
+                                                         f" {integrator_tolerance}, is not the same as" \
+                                                         f" the specified value of {constraint_tolerance}."
 
 
 # =============================================================================