Merge pull request #894 from peastman/numadded

Added option to add a fixed number of solvent molecules

docs-source/usersguide/application.rst

@@ -1405,8 +1405,13 @@ size, you can specify one:
 
     modeller.addSolvent(forcefield, boxSize=Vec3(5.0, 3.5, 3.5)*nanometers)
 
-This requests a 5 nm by 3.5 nm by 3.5 nm box.  Another option is to specify a
-padding distance:
+This requests a 5 nm by 3.5 nm by 3.5 nm box.  For a non-rectangular box, you
+can specify the three box vectors defining the unit cell:
+::
+
+    modeller.addSolvent(forcefield, boxVectors=(avec, bvec, cvec))
+
+Another option is to specify a padding distance:
 ::
 
     modeller.addSolvent(forcefield, padding=1.0*nanometers)
@@ -1416,6 +1421,14 @@ then creates a cubic box of width (solute size)+2*(padding).  The above line
 guarantees that no part of the solute comes closer than 1 nm to any edge of the
 box.
 
+Finally, you can specify the exact number of solvent molecules (including both
+water and ions) to add.  This is useful when you want to solvate several different
+conformations of the same molecule while guaranteeing they all have the same
+amount of solvent:
+::
+
+    modeller.addSolvent(forcefield, numAdded=5000)
+
 By default, :meth:`addSolvent` creates TIP3P water molecules, but it also supports other
 water models:
 ::

wrappers/python/simtk/openmm/app/modeller.py

@@ -240,7 +240,7 @@ class Modeller(object):
         self.topology = newTopology
         self.positions = newPositions
 
-    def addSolvent(self, forcefield, model='tip3p', boxSize=None, boxVectors=None, padding=None, positiveIon='Na+', negativeIon='Cl-', ionicStrength=0*molar):
+    def addSolvent(self, forcefield, model='tip3p', boxSize=None, boxVectors=None, padding=None, numAdded=None, positiveIon='Na+', negativeIon='Cl-', ionicStrength=0*molar):
         """Add solvent (both water and ions) to the model to fill a rectangular box.
 
         The algorithm works as follows:
@@ -250,11 +250,15 @@ class Modeller(object):
            randomly selecting a water molecule and replacing it with the ion.
         4. Ion pairs are added to give the requested total ionic strength.
 
-        The box size can be specified in four ways.  First, you can explicitly give the vectors defining the periodic box to
-        use.  Alternatively, for a rectangular box you can simply give the dimensions of the unit cell.  Third, you can
-        give a padding distance.  The largest dimension of the solute (along the x, y, or z axis) is determined, and a cubic
-        box of size (largest dimension)+2*padding is used.  Finally, if neither box vectors, box size, nor padding distance is specified,
-        the existing Topology's box vectors are used.
+        The box size can be specified in any of several ways:
+        
+        1. You can explicitly give the vectors defining the periodic box to use.
+        2. Alternatively, for a rectangular box you can simply give the dimensions of the unit cell.
+        3. You can give a padding distance.  The largest dimension of the solute (along the x, y, or z axis) is determined, and a cubic
+        box of size (largest dimension)+2*padding is used.
+        4. You can specify the total number of molecules (both waters and ions) to add.  A cubic box is then created whose size is
+        just large enough hold the specified amount of solvent.
+        5. Finally, if none of the above options is specified, the existing Topology's box vectors are used.
 
         Parameters:
          - forcefield (ForceField) the ForceField to use for determining van der Waals radii and atomic charges
@@ -262,14 +266,43 @@ class Modeller(object):
          - boxSize (Vec3=None) the size of the box to fill with water
          - boxVectors (tuple of Vec3=None) the vectors defining the periodic box to fill with water
          - padding (distance=None) the padding distance to use
+         - numAdded (int=None) the total number of molecules (waters and ions) to add
          - positiveIon (string='Na+') the type of positive ion to add.  Allowed values are 'Cs+', 'K+', 'Li+', 'Na+', and 'Rb+'
          - negativeIon (string='Cl-') the type of negative ion to add.  Allowed values are 'Cl-', 'Br-', 'F-', and 'I-'. Be aware
            that not all force fields support all ion types.
          - ionicStrength (concentration=0*molar) the total concentration of ions (both positive and negative) to add.  This
            does not include ions that are added to neutralize the system.
         """
+        if len([x for x in (boxSize, boxVectors, padding, numAdded) if x is not None]) > 1:
+            raise ValueError('At most one of the following arguments may be specified: boxSize, boxVectors, padding, numAdded')
+
+        # Load the pre-equilibrated water box.
+
+        vdwRadiusPerSigma = 0.5612310241546864907
+        if model == 'tip3p':
+            waterRadius = 0.31507524065751241*vdwRadiusPerSigma
+        elif model == 'spce':
+            waterRadius = 0.31657195050398818*vdwRadiusPerSigma
+        elif model == 'tip4pew':
+            waterRadius = 0.315365*vdwRadiusPerSigma
+        elif model == 'tip5p':
+            waterRadius = 0.312*vdwRadiusPerSigma
+        else:
+            raise ValueError('Unknown water model: %s' % model)
+        pdb = PDBFile(os.path.join(os.path.dirname(__file__), 'data', model+'.pdb'))
+        pdbTopology = pdb.getTopology()
+        pdbPositions = pdb.getPositions().value_in_unit(nanometer)
+        pdbResidues = list(pdbTopology.residues())
+        pdbBoxSize = pdbTopology.getUnitCellDimensions().value_in_unit(nanometer)            
+        
         # Pick a unit cell size.
 
+        if numAdded is not None:
+            # Select a padding distance which is guaranteed to give more than the specified number of molecules.
+            
+            padding = 1.1*(numAdded/((len(pdbResidues)/pdbBoxSize[0]**3)*8))**(1.0/3.0)
+            if padding < 0.5:
+                padding = 0.5 # Ensure we have enough when adding very small numbers of molecules
         if boxVectors is not None:
             if is_quantity(boxVectors[0]):
                 boxVectors = (boxVectors[0].value_in_unit(nanometer), boxVectors[1].value_in_unit(nanometer), boxVectors[2].value_in_unit(nanometer))
@@ -305,25 +338,6 @@ class Modeller(object):
         positiveElement = posIonElements[positiveIon]
         negativeElement = negIonElements[negativeIon]
 
-        # Load the pre-equilibrated water box.
-
-        vdwRadiusPerSigma = 0.5612310241546864907
-        if model == 'tip3p':
-            waterRadius = 0.31507524065751241*vdwRadiusPerSigma
-        elif model == 'spce':
-            waterRadius = 0.31657195050398818*vdwRadiusPerSigma
-        elif model == 'tip4pew':
-            waterRadius = 0.315365*vdwRadiusPerSigma
-        elif model == 'tip5p':
-            waterRadius = 0.312*vdwRadiusPerSigma
-        else:
-            raise ValueError('Unknown water model: %s' % model)
-        pdb = PDBFile(os.path.join(os.path.dirname(__file__), 'data', model+'.pdb'))
-        pdbTopology = pdb.getTopology()
-        pdbPositions = pdb.getPositions().value_in_unit(nanometer)
-        pdbResidues = list(pdbTopology.residues())
-        pdbBoxSize = pdbTopology.getUnitCellDimensions().value_in_unit(nanometer)
-
         # Have the ForceField build a System for the solute from which we can determine van der Waals radii.
 
         system = forcefield.createSystem(self.topology)
@@ -424,27 +438,42 @@ class Modeller(object):
 
                                 addedWaters.append((residue.index, atomPos))
 
-        # There could be clashes between water molecules at the box edges.  Find ones to remove.
-
-        upperCutoff = center+box/2-Vec3(waterCutoff, waterCutoff, waterCutoff)
-        lowerCutoff = center-box/2+Vec3(waterCutoff, waterCutoff, waterCutoff)
-        lowerSkinPositions = [pos for index, pos in addedWaters if pos[0] < lowerCutoff[0] or pos[1] < lowerCutoff[1] or pos[2] < lowerCutoff[2]]
-        filteredWaters = []
-        cells = {}
-        for i in range(len(lowerSkinPositions)):
-            cell = tuple((int(floor(lowerSkinPositions[i][j]/cellSize[j]))%numCells[j] for j in range(3)))
-            if cell in cells:
-                cells[cell].append(i)
-            else:
-                cells[cell] = [i]
-        for entry in addedWaters:
-            pos = entry[1]
-            if pos[0] < upperCutoff[0] and pos[1] < upperCutoff[1] and pos[2] < upperCutoff[2]:
-                filteredWaters.append(entry)
-            else:
-                if not any((periodicDistance(lowerSkinPositions[i], pos) < waterCutoff and norm(lowerSkinPositions[i]-pos) > waterCutoff for i in neighbors(pos))):
+        if numAdded is not None:
+            # We added many more waters than we actually want.  Sort them based on distance to the nearest box edge and
+            # only keep the ones in the middle.
+            
+            lowerBound = center-box/2
+            upperBound = center+box/2
+            distToEdge = (min(min(pos-lowerBound), min(upperBound-pos)) for index, pos in addedWaters)
+            sortedIndex = [i[0] for i in sorted(enumerate(distToEdge), key=lambda x: -x[1])]
+            addedWaters = [addedWaters[i] for i in sortedIndex[:numAdded]]
+            
+            # Compute a new periodic box size.
+            
+            maxSize = max(max((pos[i] for index, pos in addedWaters))-min((pos[i] for index, pos in addedWaters)) for i in range(3))
+            newTopology.setUnitCellDimensions(Vec3(maxSize, maxSize, maxSize))
+        else:
+            # There could be clashes between water molecules at the box edges.  Find ones to remove.
+    
+            upperCutoff = center+box/2-Vec3(waterCutoff, waterCutoff, waterCutoff)
+            lowerCutoff = center-box/2+Vec3(waterCutoff, waterCutoff, waterCutoff)
+            lowerSkinPositions = [pos for index, pos in addedWaters if pos[0] < lowerCutoff[0] or pos[1] < lowerCutoff[1] or pos[2] < lowerCutoff[2]]
+            filteredWaters = []
+            cells = {}
+            for i in range(len(lowerSkinPositions)):
+                cell = tuple((int(floor(lowerSkinPositions[i][j]/cellSize[j]))%numCells[j] for j in range(3)))
+                if cell in cells:
+                    cells[cell].append(i)
+                else:
+                    cells[cell] = [i]
+            for entry in addedWaters:
+                pos = entry[1]
+                if pos[0] < upperCutoff[0] and pos[1] < upperCutoff[1] and pos[2] < upperCutoff[2]:
                     filteredWaters.append(entry)
-        addedWaters = filteredWaters
+                else:
+                    if not any((periodicDistance(lowerSkinPositions[i], pos) < waterCutoff and norm(lowerSkinPositions[i]-pos) > waterCutoff for i in neighbors(pos))):
+                        filteredWaters.append(entry)
+            addedWaters = filteredWaters
 
         # Add ions to neutralize the system.
 

wrappers/python/tests/TestModeller.py

@@ -308,7 +308,7 @@ class TestModeller(unittest.TestCase):
                         self.assertTrue(len(matoms)==0 and len(m1atoms)==1 and len(m2atoms)==1)
     
     def test_addSolventPeriodicBox(self):
-        """ Test the addSolvent() method; test that the four ways of passing in the periodic box all work. """
+        """ Test the addSolvent() method; test that the five ways of passing in the periodic box all work. """
     
         # First way of passing in periodic box vectors:  set it in the original topology.
         topology_start = self.pdb.topology
@@ -358,6 +358,14 @@ class TestModeller(unittest.TestCase):
         self.assertVecAlmostEqual(dim3[0]/nanometers, Vec3(2.8802, 0, 0))
         self.assertVecAlmostEqual(dim3[1]/nanometers, Vec3(0, 2.8802, 0))
         self.assertVecAlmostEqual(dim3[2]/nanometers, Vec3(0, 0, 2.8802))
+        
+        # Fifth way: specify a number of molecules to add instead of a box size
+        topology_start = self.pdb.topology
+        modeller = Modeller(topology_start, self.positions)
+        modeller.deleteWater()
+        numInitial = len(list(modeller.topology.residues()))
+        modeller.addSolvent(self.forcefield, numAdded=1000)
+        self.assertEqual(numInitial+1000, len(list(modeller.topology.residues())))
 
     def test_addSolventNeutralSolvent(self):
         """ Test the addSolvent() method; test adding ions to neutral solvent. """
@@ -389,7 +397,7 @@ class TestModeller(unittest.TestCase):
 
     def test_addSolventNegativeSolvent(self):
         """ Test the addSolvent() method; test adding ions to a negatively charged solvent. """
-
+    
         topology_start = self.pdb.topology
         topology_start.setUnitCellDimensions(Vec3(3.5, 3.5, 3.5)*nanometers)
                
@@ -476,7 +484,7 @@ class TestModeller(unittest.TestCase):
        
         topology_start = self.pdb.topology
         topology_start.setUnitCellDimensions(Vec3(3.5, 3.5, 3.5)*nanometers)
- 
+
         # set up modeller with no solvent
         modeller = Modeller(topology_start, self.positions)
         modeller.deleteWater()
@@ -490,7 +498,7 @@ class TestModeller(unittest.TestCase):
             modeller = Modeller(topology_nowater, positions_nowater)
             modeller.addSolvent(self.forcefield, positiveIon=positiveIon, ionicStrength=1.0*molar)
             topology_after = modeller.getTopology()
-   
+
             water_count = 0 
             positive_ion_count = 0
             chlorine_count = 0
@@ -514,7 +522,7 @@ class TestModeller(unittest.TestCase):
             modeller.addSolvent(self.forcefield, negativeIon=negativeIon, ionicStrength=1.0*molar)
 
             topology_after = modeller.getTopology()
-   
+
             water_count = 0
             sodium_count = 0   
             negative_ion_count = 0
@@ -543,7 +551,7 @@ class TestModeller(unittest.TestCase):
         # remove hydrogens from the topology
         toDelete = [atom for atom in topology_start.atoms() if atom.element==Element.getBySymbol('H')]
         modeller.delete(toDelete)
- 
+
         # Create a variants list to force the one histidine to be of the right variation.
         residues = [residue for residue in topology_start.residues()]
         variants = [None]*len(residues)
@@ -557,7 +565,7 @@ class TestModeller(unittest.TestCase):
         topology_after = modeller.getTopology()
 
         validate_equivalence(self, topology_start, topology_after)
-   
+
     def test_addHydrogensPdb3(self):
         """ Test the addHydrogens() method on the metallothionein pdb file. """
         
@@ -650,7 +658,7 @@ class TestModeller(unittest.TestCase):
 
         # There should be extra hydrogens on the CYS residues.  Assert that they exist
         # on modeller2, then delete them and validate that the topologies match.
- 
+
        # These are the indices of the hydrogens to delete from CYS to make CYX.
         index_list_CYS = [31, 49, 110, 135, 171, 193, 229]
         atoms = [atom for atom in topology2.atoms()]