addExtraParticles() can use bonds to select positions

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

@@ -35,7 +35,7 @@ __author__ = "Peter Eastman"
 __version__ = "1.0"
 
 from simtk.openmm.app import Topology, PDBFile, ForceField
-from simtk.openmm.app.forcefield import HAngles, _createResidueSignature, _matchResidue, DrudeGenerator
+from simtk.openmm.app.forcefield import HAngles, AllBonds, _createResidueSignature, _matchResidue, DrudeGenerator
 from simtk.openmm.app.topology import Residue
 from simtk.openmm.vec3 import Vec3
 from simtk.openmm import System, Context, NonbondedForce, CustomNonbondedForce, HarmonicBondForce, HarmonicAngleForce, VerletIntegrator, LocalEnergyMinimizer
@@ -929,6 +929,7 @@ class Modeller(object):
         newTopology.setPeriodicBoxVectors(self.topology.getPeriodicBoxVectors())
         newAtoms = {}
         newPositions = []*nanometer
+        missingPositions = set()
         for chain in self.topology.chains():
             newChain = newTopology.addChain(chain.id)
             for residue in chain.residues():
@@ -986,9 +987,10 @@ class Modeller(object):
                     for index, atom in enumerate(template.atoms):
                         if atom in matchingAtoms:
                             templateAtomPositions[index] = self.positions[matchingAtoms[atom].index].value_in_unit(nanometer)
+                    newExtraPoints = {}
                     for index, atom in enumerate(template.atoms):
                         if atom.element is None:
-                            newTopology.addAtom(atom.name, None, newResidue)
+                            newExtraPoints[atom] = newTopology.addAtom(atom.name, None, newResidue)
                             position = None
                             for site in template.virtualSites:
                                 if site.index == index:
@@ -1010,14 +1012,61 @@ class Modeller(object):
                                     if atom2.type in drudeTypeMap[atom.type]:
                                         position = deepcopy(pos)
                             if position is None:
-                                # We couldn't figure out the correct position.  As a wild guess, just put it at the center of the residue
-                                # and hope that energy minimization will fix it.
+                                # We couldn't figure out the correct position.  Put it at a random position near the center of the residue,
+                                # and we'll try to fix it later based on bonds.
 
                                 knownPositions = [x for x in templateAtomPositions if x is not None]
-                                position = unit.sum(knownPositions)/len(knownPositions)
+                                position = Vec3(random.gauss(0, 1), random.gauss(0, 1), random.gauss(0, 1))+(unit.sum(knownPositions)/len(knownPositions))
+                                missingPositions.add(len(newPositions))
                             newPositions.append(position*nanometer)
+
+                    # Add bonds involving the extra points.
+
+                    for atom1, atom2 in template.bonds:
+                        atom1 = template.atoms[atom1]
+                        atom2 = template.atoms[atom2]
+                        if atom1 in newExtraPoints or atom2 in newExtraPoints:
+                            if atom1 in newExtraPoints:
+                                a1 = newExtraPoints[atom1]
+                            else:
+                                a1 = matchingAtoms[atom1]
+                            if atom2 in newExtraPoints:
+                                a2 = newExtraPoints[atom2]
+                            else:
+                                a2 = matchingAtoms[atom2]
+                            newTopology.addBond(a1, a2)
+                            
         for bond in self.topology.bonds():
             if bond[0] in newAtoms and bond[1] in newAtoms:
                 newTopology.addBond(newAtoms[bond[0]], newAtoms[bond[1]])
+
+        if len(missingPositions) > 0:
+            # There were particles whose position we couldn't identify before, since they were neither virtual sites nor Drude particles.
+            # Try to figure them out based on bonds.  First, use the ForceField to create a list of every bond involving one of them.
+
+            system = forcefield.createSystem(newTopology, constraints=AllBonds)
+            bonds = []
+            for i in range(system.getNumConstraints()):
+                bond = system.getConstraintParameters(i)
+                if bond[0] in missingPositions or bond[1] in missingPositions:
+                    bonds.append(bond)
+
+            # Now run a few iterations of SHAKE to try to select reasonable positions.
+
+            for iteration in range(10):
+                for atom1, atom2, distance in bonds:
+                    if atom1 in missingPositions:
+                        if atom2 in missingPositions:
+                            weights = (0.5, 0.5)
+                        else:
+                            weights = (1.0, 0.0)
+                    else:
+                        weights = (0.0, 1.0)
+                    delta = newPositions[atom2]-newPositions[atom1]
+                    length = norm(delta)
+                    delta *= (distance-length)/length
+                    newPositions[atom1] -= weights[0]*delta
+                    newPositions[atom2] += weights[1]*delta
+                
         self.topology = newTopology
         self.positions = newPositions

wrappers/python/tests/TestModeller.py

@@ -3,6 +3,10 @@ from validateModeller import *
 from simtk.openmm.app import *
 from simtk.openmm import *
 from simtk.unit import *
+if sys.version_info >= (3, 0):
+    from io import StringIO
+else:
+    from cStringIO import StringIO
 
 class TestModeller(unittest.TestCase):
     """ Test the Modeller class. """
@@ -915,6 +919,70 @@ class TestModeller(unittest.TestCase):
             self.assertEqual(1, len(ep))
 
 
+    def test_multiSiteIon(self):
+        """Test adding extra particles whose positions are determined based on bonds."""
+        xml = """
+<ForceField>
+ <AtomTypes>
+  <Type name="Zn" class="Zn" element="Zn" mass="53.380"/>
+  <Type name="DA" class="DA" mass="3.0"/>
+ </AtomTypes>
+ <Residues>
+  <Residue name="ZN">
+   <Atom name="ZN" type="Zn"/>
+   <Atom name="D1" type="DA"/>
+   <Atom name="D2" type="DA"/>
+   <Atom name="D3" type="DA"/>
+   <Atom name="D4" type="DA"/>
+   <Bond from="0" to="2"/>
+   <Bond from="0" to="1"/>
+   <Bond from="0" to="3"/>
+   <Bond from="0" to="4"/>
+   <Bond from="1" to="2"/>
+   <Bond from="1" to="3"/>
+   <Bond from="1" to="4"/>
+   <Bond from="2" to="4"/>
+   <Bond from="2" to="3"/>
+   <Bond from="3" to="4"/>
+  </Residue>
+ </Residues>
+ <HarmonicBondForce>
+  <Bond class1="DA" class2="Zn" length="0.09" k="535552.0"/>
+  <Bond class1="DA" class2="DA" length="0.147" k="535552.0"/>
+ </HarmonicBondForce>
+</ForceField>"""
+        ff = ForceField(StringIO(xml))
+
+        # Create a zinc atom.
+
+        topology = Topology()
+        chain = topology.addChain()
+        residue = topology.addResidue('ZN', chain)
+        topology.addAtom('ZN', element.zinc, residue)
+
+        # Add the extra particles.
+
+        modeller = Modeller(topology, [Vec3(0.5, 1.0, 1.5)]*nanometers)
+        modeller.addExtraParticles(ff)
+        top = modeller.topology
+        pos = modeller.positions
+
+        # Check that the correct particles were added.
+
+        self.assertEqual(len(pos), 5)
+        for i, atom in enumerate(top.atoms()):
+            self.assertEqual(element.zinc if i == 0 else None, atom.element)
+
+        # Check that their positions are reasonable.
+
+        center = Vec3(0.5, 1.0, 1.5)*nanometers
+        self.assertEqual(center, modeller.positions[0])
+        for i in range(1, 5):
+            for j in range(i):
+                dist = norm(pos[i]-pos[j])
+                expectedDist = 0.09 if j == 0 else 0.147
+                self.assertTrue(dist > (expectedDist-0.01)*nanometers and dist < (expectedDist+0.01)*nanometers)
+
     def assertVecAlmostEqual(self, p1, p2, tol=1e-7):
         scale = max(1.0, norm(p1),)
         for i in range(3):