Merge pull request #218 from choderalab/lrc

Add option to AbsoluteAlchemicalFactory to disable long-range dispersion correction for alchemical region

openmmtools/alchemy.py

@@ -673,10 +673,16 @@ class AbsoluteAlchemicalFactory(object):
         'direct-space' only models the direct space contribution
         'coulomb' includes switched Coulomb interaction
     alchemical_rf_treatment : str, optional, default = 'switched'
-         Controls how alchemical region electrostatics are treated when RF is used
-         Options are ['switched', 'shifted']
-         'switched' sets c_rf = 0 for all reaction-field interactions and ensures continuity with a switch
-         'shifted' retains c_rf != 0 but can give erroneous results for hydration free energies
+        Controls how alchemical region electrostatics are treated when RF is used
+        Options are ['switched', 'shifted']
+        'switched' sets c_rf = 0 for all reaction-field interactions and ensures continuity with a switch
+        'shifted' retains c_rf != 0 but can give erroneous results for hydration free energies
+    disable_alchemical_dispersion_correction : bool, optional, default=False
+        If True, the long-range dispersion correction will not be included for the alchemical
+        region to avoid the need to recompute the correction (a CPU operation that takes ~ 0.5 s)
+        every time 'lambda_sterics' is changed. If using nonequilibrium protocols, it is recommended
+        that this be set to True since this can lead to enormous (100x) slowdowns if the correction
+        must be recomputed every timestep.
 
     Examples
     --------
@@ -735,6 +741,15 @@ class AbsoluteAlchemicalFactory(object):
     >>> alchemical_state.set_alchemical_parameters(0.0)  # Set all lambda to 0
     >>> alchemical_state.apply_to_context(context)
 
+    Neglecting the long-range dispersion correction for the alchemical region
+    (for nonequilibrium switching, for example) requires instantiating a factory
+    with the appropriate options:
+
+    >>> new_factory = AbsoluteAlchemicalFactory(consistent_exceptions=False, disable_alchemical_dispersion_correction=True)
+    >>> reference_system = testsystems.WaterBox().system
+    >>> alchemical_region = AlchemicalRegion(alchemical_atoms=[0, 1, 2])
+    >>> alchemical_system = new_factory.create_alchemical_system(reference_system, alchemical_region)
+
     References
     ----------
     [1] Pham TT and Shirts MR. Identifying low variance pathways for free
@@ -747,11 +762,15 @@ class AbsoluteAlchemicalFactory(object):
     # Public interface
     # -------------------------------------------------------------------------
 
-    def __init__(self, consistent_exceptions=False, switch_width=1.0*unit.angstroms, alchemical_pme_treatment='direct-space', alchemical_rf_treatment='switched'):
+    def __init__(self, consistent_exceptions=False, switch_width=1.0*unit.angstroms,
+                 alchemical_pme_treatment='direct-space', alchemical_rf_treatment='switched',
+                 disable_alchemical_dispersion_correction=False):
+
         self.consistent_exceptions = consistent_exceptions
         self.switch_width = switch_width
         self.alchemical_pme_treatment = alchemical_pme_treatment
         self.alchemical_rf_treatment = alchemical_rf_treatment
+        self.disable_alchemical_dispersion_correction = disable_alchemical_dispersion_correction
 
     def create_alchemical_system(self, reference_system, alchemical_regions):
         """Create an alchemically modified version of the reference system.
@@ -1526,7 +1545,10 @@ def _alchemically_modify_NonbondedForce(self, reference_force, alchemical_region
             force.setUseSwitchingFunction(nonbonded_force.getUseSwitchingFunction())
             force.setCutoffDistance(nonbonded_force.getCutoffDistance())
             force.setSwitchingDistance(nonbonded_force.getSwitchingDistance())
-            force.setUseLongRangeCorrection(nonbonded_force.getUseDispersionCorrection())
+            if self.disable_alchemical_dispersion_correction:
+                force.setUseLongRangeCorrection(False)
+            else:
+                force.setUseLongRangeCorrection(nonbonded_force.getUseDispersionCorrection())
 
             if is_method_periodic:
                 force.setNonbondedMethod(openmm.CustomNonbondedForce.CutoffPeriodic)

openmmtools/tests/test_alchemy.py

@@ -1289,6 +1289,76 @@ def test_overlap(self):
             f.description = "Testing reference/alchemical overlap for {}".format(test_name)
             yield f
 
+class TestDispersionlessAlchemicalFactory(object):
+    """
+    Only test overlap for dispersionless alchemical factory, since energy agreement
+    will be poor.
+    """
+    @classmethod
+    def setup_class(cls):
+        """Create test systems and shared objects."""
+        cls.define_systems()
+        cls.define_regions()
+        cls.generate_cases()
+
+    @classmethod
+    def define_systems(cls):
+        """Create test systems and shared objects."""
+        cls.test_systems = dict()
+        cls.test_systems['LennardJonesFluid with dispersion correction'] = \
+            testsystems.LennardJonesFluid(nparticles=100, dispersion_correction=True)
+
+    @classmethod
+    def define_regions(cls):
+        """Create shared AlchemicalRegions for test systems in cls.test_regions."""
+        cls.test_regions = dict()
+        cls.test_regions['LennardJonesFluid'] = AlchemicalRegion(alchemical_atoms=range(10))
+
+    @classmethod
+    def generate_cases(cls):
+        """Generate all test cases in cls.test_cases combinatorially."""
+        cls.test_cases = dict()
+        factory = AbsoluteAlchemicalFactory(disable_alchemical_dispersion_correction=True)
+
+        # We generate all possible combinations of annihilate_sterics/electrostatics
+        # for each test system. We also annihilate bonds, angles and torsions every
+        # 3 test cases so that we test it at least one for each test system and for
+        # each combination of annihilate_sterics/electrostatics.
+        n_test_cases = 0
+        for test_system_name, test_system in cls.test_systems.items():
+
+            # Find standard alchemical region.
+            for region_name, region in cls.test_regions.items():
+                if region_name in test_system_name:
+                    break
+            assert region_name in test_system_name
+
+            # Create all combinations of annihilate_sterics.
+            for annihilate_sterics in itertools.product((True, False), repeat=1):
+                region = region._replace(annihilate_sterics=annihilate_sterics,
+                                         annihilate_electrostatics=True)
+
+                # Create test name.
+                test_case_name = test_system_name[:]
+                if annihilate_sterics:
+                    test_case_name += ', annihilated sterics'
+
+                # Pre-generate alchemical system
+                alchemical_system = factory.create_alchemical_system(test_system.system, region)
+                cls.test_cases[test_case_name] = (test_system, alchemical_system, region)
+
+                n_test_cases += 1
+
+    def test_overlap(self):
+        """Tests overlap between reference and alchemical systems."""
+        for test_name, (test_system, alchemical_system, alchemical_region) in self.test_cases.items():
+            #cached_trajectory_filename = os.path.join(os.environ['HOME'], '.cache', 'alchemy', 'tests',
+            #                                           test_name + '.pickle')
+            cached_trajectory_filename = None
+            f = partial(overlap_check, test_system.system, alchemical_system, test_system.positions,
+                        cached_trajectory_filename=cached_trajectory_filename, name=test_name)
+            f.description = "Testing reference/alchemical overlap for no alchemical dispersion {}".format(test_name)
+            yield f
 
 @attr('slow')
 class TestAbsoluteAlchemicalFactorySlow(TestAbsoluteAlchemicalFactory):