Merge pull request #259 from choderalab/virtual-sites

Virtual sites in alchemical factory and add ions in WaterBox

openmmtools/alchemy.py

@@ -802,39 +802,37 @@ def create_alchemical_system(self, reference_system, alchemical_regions):
         timer = utils.Timer()
         timer.start('Create alchemically modified system')
 
-        # Build alchemical system to modify.
-        alchemical_system = openmm.System()
+        # Build alchemical system to modify. This copies particles, vsites,
+        # constraints, box vectors and all the forces. We'll later remove
+        # the forces that we remodel to be alchemically modified.
+        alchemical_system = copy.deepcopy(reference_system)
 
-        # Set periodic box vectors.
-        box_vectors = reference_system.getDefaultPeriodicBoxVectors()
-        alchemical_system.setDefaultPeriodicBoxVectors(*box_vectors)
-
-        # TODO: Are we missing important components of the System here, such as vsites?
-        # Should we deepcopy instead?
-
-        # Add particles.
+        # Check that there are no virtual sites to alchemically modify.
         for particle_index in range(reference_system.getNumParticles()):
-            mass = reference_system.getParticleMass(particle_index)
-            alchemical_system.addParticle(mass)
-
-        # Add constraints.
-        for constraint_index in range(reference_system.getNumConstraints()):
-            atom_i, atom_j, r0 = reference_system.getConstraintParameters(constraint_index)
-            alchemical_system.addConstraint(atom_i, atom_j, r0)
-
-        # Modify forces as appropriate, copying other forces without modification.
-        for reference_force in reference_system.getForces():
+            if (reference_system.isVirtualSite(particle_index) and
+                        particle_index in alchemical_region.alchemical_atoms):
+                raise ValueError('Alchemically modified virtual sites are not supported')
+
+        # Modify forces as appropriate. We delete the forces that
+        # have been processed modified at the end of the for loop.
+        forces_to_remove = []
+        for force_index, reference_force in enumerate(reference_system.getForces()):
             # TODO switch to functools.singledispatch when we drop Python2 support
             reference_force_name = reference_force.__class__.__name__
             alchemical_force_creator_name = '_alchemically_modify_{}'.format(reference_force_name)
             try:
                 alchemical_force_creator_func = getattr(self, alchemical_force_creator_name)
             except AttributeError:
-                alchemical_forces = [copy.deepcopy(reference_force)]
+                pass
             else:
+                forces_to_remove.append(force_index)
                 alchemical_forces = alchemical_force_creator_func(reference_force, alchemical_region)
-            for alchemical_force in alchemical_forces:
-                alchemical_system.addForce(alchemical_force)
+                for alchemical_force in alchemical_forces:
+                    alchemical_system.addForce(alchemical_force)
+
+        # Remove original forces that have been alchemically modified.
+        for force_index in reversed(forces_to_remove):
+            alchemical_system.removeForce(force_index)
 
         # Record timing statistics.
         timer.stop('Create alchemically modified system')

openmmtools/tests/test_alchemy.py

@@ -1126,8 +1126,10 @@ def define_systems(cls):
         cls.test_systems['LennardJonesCluster'] = testsystems.LennardJonesCluster()
         cls.test_systems['LennardJonesFluid with dispersion correction'] = \
             testsystems.LennardJonesFluid(nparticles=100, dispersion_correction=True)
-        cls.test_systems['WaterBox with reaction field, no switch, no dispersion correction'] = \
+        cls.test_systems['TIP3P WaterBox with reaction field, no switch, no dispersion correction'] = \
             testsystems.WaterBox(dispersion_correction=False, switch=False, nonbondedMethod=openmm.app.CutoffPeriodic)
+        cls.test_systems['TIP4P-EW WaterBox and NaCl with PME'] = \
+            testsystems.WaterBox(nonbondedMethod=openmm.app.PME, model='tip4pew', ionic_strength=200*unit.millimolar)
 
         # Vacuum and implicit.
         cls.test_systems['AlanineDipeptideVacuum'] = testsystems.AlanineDipeptideVacuum()
@@ -1149,7 +1151,8 @@ def define_regions(cls):
         cls.test_regions = dict()
         cls.test_regions['LennardJonesCluster'] = AlchemicalRegion(alchemical_atoms=range(2))
         cls.test_regions['LennardJonesFluid'] = AlchemicalRegion(alchemical_atoms=range(10))
-        cls.test_regions['WaterBox'] = AlchemicalRegion(alchemical_atoms=range(3))
+        cls.test_regions['TIP3P WaterBox'] = AlchemicalRegion(alchemical_atoms=range(3))
+        cls.test_regions['TIP4P-EW WaterBox and NaCl'] = AlchemicalRegion(alchemical_atoms=range(4))  # Modify ions.
         cls.test_regions['Toluene'] = AlchemicalRegion(alchemical_atoms=range(6))  # Only partially modified.
         cls.test_regions['AlanineDipeptide'] = AlchemicalRegion(alchemical_atoms=range(22))
         cls.test_regions['HostGuestExplicit'] = AlchemicalRegion(alchemical_atoms=range(126, 156))

openmmtools/testsystems.py

@@ -2470,7 +2470,10 @@ class WaterBox(TestSystem):
 
     """
 
-    def __init__(self, box_edge=25.0*unit.angstroms, cutoff=DEFAULT_CUTOFF_DISTANCE, model='tip3p', switch_width=DEFAULT_SWITCH_WIDTH, constrained=True, dispersion_correction=True, nonbondedMethod=app.PME, ewaldErrorTolerance=DEFAULT_EWALD_ERROR_TOLERANCE, **kwargs):
+    def __init__(self, box_edge=25.0*unit.angstroms, cutoff=DEFAULT_CUTOFF_DISTANCE, model='tip3p',
+                 switch_width=DEFAULT_SWITCH_WIDTH, constrained=True, dispersion_correction=True,
+                 nonbondedMethod=app.PME, ewaldErrorTolerance=DEFAULT_EWALD_ERROR_TOLERANCE,
+                 positive_ion='Na+', negative_ion='Cl-', ionic_strength=0*unit.molar, **kwargs):
         """
         Create a water box test system.
 
@@ -2493,6 +2496,13 @@ def __init__(self, box_edge=25.0*unit.angstroms, cutoff=DEFAULT_CUTOFF_DISTANCE,
            Sets the nonbonded method to use for the water box (one of app.CutoffPeriodic, app.Ewald, app.PME).
         ewaldErrorTolerance : float, optional, default=DEFAULT_EWALD_ERROR_TOLERANCE
            The Ewald or PME tolerance.  Used only if nonbondedMethod is Ewald or PME.
+        positive_ion : str, optional
+            The positive ion to add (default is Na+).
+        negative_ion : str, optional
+            The negative ion to add (default is Cl-).
+        ionic_strength : simtk.unit.Quantity, optional
+            The total concentration of ions (both positive and negative)
+            to add (default is 0.0*molar).
 
         Examples
         --------
@@ -2538,8 +2548,11 @@ def __init__(self, box_edge=25.0*unit.angstroms, cutoff=DEFAULT_CUTOFF_DISTANCE,
         if model not in supported_models:
             raise Exception("Specified water model '%s' is not in list of supported models: %s" % (model, str(supported_models)))
 
-        # Load forcefield for solvent model.
-        ff = app.ForceField(model + '.xml')
+        # Load forcefield for solvent model and ions.
+        force_fields = [model + '.xml']
+        if ionic_strength != 0.0*unit.molar:
+            force_fields.append('amber99sb.xml')  # For the ions.
+        ff = app.ForceField(*force_fields)
 
         # Create empty topology and coordinates.
         top = app.Topology()
@@ -2550,7 +2563,8 @@ def __init__(self, box_edge=25.0*unit.angstroms, cutoff=DEFAULT_CUTOFF_DISTANCE,
 
         # Add solvent to specified box dimensions.
         boxSize = unit.Quantity(numpy.ones([3]) * box_edge / box_edge.unit, box_edge.unit)
-        m.addSolvent(ff, boxSize=boxSize, model=model)
+        m.addSolvent(ff, boxSize=boxSize, model=model, positiveIon=positive_ion,
+                     negativeIon=negative_ion, ionicStrength=ionic_strength)
 
         # Get new topology and coordinates.
         newtop = m.getTopology()