Added option to add a fixed number of solvent molecules

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,6 +438,21 @@ class Modeller(object):
 
                                 addedWaters.append((residue.index, atomPos))
 
+        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)

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
@@ -359,6 +359,14 @@ class TestModeller(unittest.TestCase):
         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. """