Implemented Modeller.addExtraParticles(). Also some cleanup in ForceField.

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

@@ -114,15 +114,11 @@ class ForceField(object):
                         template.externalBonds.append(b)
                         template.atoms[b].externalBonds += 1
             for template in self._templates.values():
-                template.signature = _createResidueSignature([atom.element for atom in template.atoms])
-                if template.signature is None:
-                    sigString = None
+                signature = _createResidueSignature([atom.element for atom in template.atoms])
+                if signature in self._templateSignatures:
+                    self._templateSignatures[signature].append(template)
                 else:
-                    sigString = _signatureToString(template.signature)
-                if sigString in self._templateSignatures:
-                    self._templateSignatures[sigString].append(template)
-                else:
-                    self._templateSignatures[sigString] = [template]
+                    self._templateSignatures[signature] = [template]
             
             # Build sets of every atom type belonging to each class
             
@@ -260,20 +256,13 @@ class ForceField(object):
            particular force fields.
         Returns: the newly created System
         """
-        
-        # Record atom indices
-        
         data = ForceField._SystemData()
-        atomIndices = {}
-        for index, atom in enumerate(topology.atoms()):
-            data.atoms.append(atom)
-            atomIndices[atom] = index
+        data.atoms = list(topology.atoms())
 
         # Make a list of all bonds
         
         for bond in topology.bonds():
-            if bond[0] in atomIndices and bond[1] in atomIndices:
-                data.bonds.append(ForceField._BondData(atomIndices[bond[0]], atomIndices[bond[1]]))
+            data.bonds.append(ForceField._BondData(bond[0].index, bond[1].index))
 
         # Record which atoms are bonded to each other atom
         
@@ -294,19 +283,10 @@ class ForceField(object):
             for res in chain.residues():
                 template = None
                 matches = None
-                sig = _createResidueSignature([atom.element for atom in res.atoms()])
-                if sig is not None:
-                    signature = _signatureToString(sig)
-                    if signature in self._templateSignatures:
-                        for t in self._templateSignatures[signature]:
-                            matches = _matchResidue(res, t, bondedToAtom, atomIndices)
-                            if matches is not None:
-                                template = t
-                                break
-                if matches is None:
-                    # Check templates involving virtual sites
-                    for t in self._templateSignatures[None]:
-                        matches = _matchResidue(res, t, bondedToAtom, atomIndices)
+                signature = _createResidueSignature([atom.element for atom in res.atoms()])
+                if signature in self._templateSignatures:
+                    for t in self._templateSignatures[signature]:
+                        matches = _matchResidue(res, t, bondedToAtom)
                         if matches is not None:
                             template = t
                             break
@@ -421,7 +401,7 @@ class ForceField(object):
         
         for atom in data.virtualSites:
             site = data.virtualSites[atom]
-            index = atomIndices[atom]
+            index = atom.index
             if site.type == 'average2':
                 sys.setVirtualSite(index, mm.TwoParticleAverageSite(index+site.atoms[0], index+site.atoms[1], site.weights[0], site.weights[1]))
             elif site.type == 'average3':
@@ -448,8 +428,8 @@ def _createResidueSignature(elements):
     counts = {}
     for element in elements:
         if element is None:
-            return None # This residue contains "atoms" (probably virtual sites) that should match any element
-        if element in counts:
+            pass # This residue contains "atoms" (probably virtual sites) that should match any element
+        elif element in counts:
             counts[element] += 1
         else:
             counts[element] = 1
@@ -457,33 +437,28 @@ def _createResidueSignature(elements):
     for c in counts:
         sig.append((c, counts[c]))
     sig.sort(key=lambda x: -x[0].mass)
-    return sig
-
+    
+    # Convert it to a string.
 
-def _signatureToString(signature):
-    """Convert the signature returned by _createResidueSignature() to a string."""
     s = ''
-    for element, count in signature:
+    for element, count in sig:
         s += element.symbol+str(count)
     return s
 
 
-def _matchResidue(res, template, bondedToAtom, atomIndices):
+def _matchResidue(res, template, bondedToAtom):
     """Determine whether a residue matches a template and return a list of corresponding atoms.
     
     Parameters:
      - res (Residue) The residue to check
      - template (_TemplateData) The template to compare it to
      - bondedToAtom (list) Enumerates which other atoms each atom is bonded to
-     - atomIndices (map) Maps from atoms to their indices in the System
     Returns: a list specifying which atom of the template each atom of the residue corresponds to,
     or None if it does not match the template
     """
     atoms = list(res.atoms())
     if len(atoms) != len(template.atoms):
         return None
-    residueAtomBonds = []
-    templateAtomBonds = []
     matches = len(atoms)*[0]
     hasMatch = len(atoms)*[False]
     
@@ -491,13 +466,13 @@ def _matchResidue(res, template, bondedToAtom, atomIndices):
     
     renumberAtoms = {}
     for i in range(len(atoms)):
-        renumberAtoms[atomIndices[atoms[i]]] = i
+        renumberAtoms[atoms[i].index] = i
     bondedTo = []
     externalBonds = []
     for atom in atoms:
-        bonds = [renumberAtoms[x] for x in bondedToAtom[atomIndices[atom]] if x in renumberAtoms]
+        bonds = [renumberAtoms[x] for x in bondedToAtom[atom.index] if x in renumberAtoms]
         bondedTo.append(bonds)
-        externalBonds.append(len([x for x in bondedToAtom[atomIndices[atom]] if x not in renumberAtoms]))
+        externalBonds.append(len([x for x in bondedToAtom[atom.index] if x not in renumberAtoms]))
     if _findAtomMatches(atoms, template, bondedTo, externalBonds, matches, hasMatch, 0):
         return matches
     return None

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

@@ -33,11 +33,12 @@ from __future__ import division
 __author__ = "Peter Eastman"
 __version__ = "1.0"
 
-from simtk.openmm.app import Topology, PDBFile
-from simtk.openmm.app.forcefield import HAngles
+from simtk.openmm.app import Topology, PDBFile, ForceField
+from simtk.openmm.app.forcefield import HAngles, _createResidueSignature, _matchResidue
+from simtk.openmm.app.topology import Residue
 from simtk.openmm.vec3 import Vec3
 from simtk.openmm import System, Context, NonbondedForce, VerletIntegrator, LocalEnergyMinimizer
-from simtk.unit import nanometer, molar, elementary_charge, amu, gram, liter, degree, sqrt, acos, is_quantity, dot, norm
+from simtk.unit import nanometer, molar, elementary_charge, amu, gram, liter, degree, sqrt, acos, is_quantity, dot, norm, sum
 import element as elem
 import os
 import random
@@ -516,7 +517,7 @@ class Modeller(object):
                     terminal = hydrogen.attrib['terminal']
                 data.hydrogens.append(Modeller._Hydrogen(hydrogen.attrib['name'], hydrogen.attrib['parent'], maxph, atomVariants, terminal))
     
-    def addHydrogens(self, forcefield, pH=7.0, variants=None, hydrogenDefinitions=None):
+    def addHydrogens(self, forcefield, pH=7.0, variants=None):
         """Add missing hydrogens to the model.
         
         Some residues can exist in multiple forms depending on the pH and properties of the local environment.  These
@@ -763,3 +764,145 @@ class Modeller(object):
         self.topology = newTopology
         self.positions = context.getState(getPositions=True).getPositions()
         return actualVariants
+    
+    def addExtraParticles(self, forcefield):
+        """Add missing extra particles to the model that are required by a force field.
+        
+        Some force fields use "extra particles" that do not represent actual atoms, but still need to be included in
+        the System.  Examples include lone pairs, Drude particles, and the virtual sites used in some water models
+        to adjust the charge distribution.  Extra particles can be recognized by the fact that their element is None.
+        
+        This method is primarily used to add extra particles, but it can also remove them.  It tries to match every
+        residue in the Topology to a template in the force field.  If there is no match, it will both add and remove
+        extra particles as necessary to make it match.
+        
+        Parameters:
+         - forcefield (ForceField) the ForceField defining what extra particles should be present
+        """
+        # Create copies of all residue templates that have had all extra points removed.
+        
+        templatesNoEP = {}
+        for resName, template in forcefield._templates.iteritems():
+            if any(atom.element is None for atom in template.atoms):
+                index = 0
+                newIndex = {}
+                newTemplate = ForceField._TemplateData(resName)
+                for i, atom in enumerate(template.atoms):
+                    if atom.element is not None:
+                        newIndex[i] = index
+                        index += 1
+                        newTemplate.atoms.append(ForceField._TemplateAtomData(atom.name, atom.type, atom.element))
+                for b1, b2 in template.bonds:
+                    if b1 in newIndex and b2 in newIndex:
+                        newTemplate.bonds.append((newIndex[b1], newIndex[b2]))
+                        newTemplate.atoms[newIndex[b1]].bondedTo.append(newIndex[b2])
+                        newTemplate.atoms[newIndex[b2]].bondedTo.append(newIndex[b1])
+                for b in template.externalBonds:
+                    if b in newIndex:
+                        newTemplate.externalBonds.append(newIndex[b])
+                templatesNoEP[template] = newTemplate
+        
+        # Record which atoms are bonded to each other atom, with and without extra particles.
+        
+        bondedToAtom = []
+        bondedToAtomNoEP = []
+        for atom in self.topology.atoms():
+            bondedToAtom.append(set())
+            if atom.element is not None:
+                bondedToAtomNoEP.append(set())
+        for atom1, atom2 in self.topology.bonds():
+            bondedToAtom[atom1.index].add(atom2.index)
+            bondedToAtom[atom2.index].add(atom1.index)
+            if atom1.element is not None and atom2.element is not None:
+                bondedToAtomNoEP[atom1.index].add(atom2.index)
+                bondedToAtomNoEP[atom2.index].add(atom1.index)
+        
+        # Create the new Topology.
+        
+        newTopology = Topology()
+        newTopology.setUnitCellDimensions(deepcopy(self.topology.getUnitCellDimensions()))
+        newAtoms = {}
+        newPositions = []*nanometer
+        for chain in self.topology.chains():
+            newChain = newTopology.addChain()
+            for residue in chain.residues():
+                newResidue = newTopology.addResidue(residue.name, newChain)
+                
+                # Look for a matching template.
+                
+                matchFound = False
+                signature = _createResidueSignature([atom.element for atom in residue.atoms()])
+                if signature in forcefield._templateSignatures:
+                    for t in forcefield._templateSignatures[signature]:
+                        if _matchResidue(residue, t, bondedToAtom) is not None:
+                            matchFound = True
+                if matchFound:
+                    # Just copy the residue over.
+                    
+                    for atom in residue.atoms():
+                        newAtom = newTopology.addAtom(atom.name, atom.element, newResidue)
+                        newAtoms[atom] = newAtom
+                        newPositions.append(deepcopy(self.positions[atom.index]))
+                else:
+                    # There's no matching template.  Try to find one that matches based on everything except
+                    # extra points.
+                    
+                    template = None
+                    residueNoEP = Residue(residue.name, residue.index, residue.chain)
+                    residueNoEP._atoms = [atom for atom in residue.atoms() if atom.element is not None]
+                    if signature in forcefield._templateSignatures:
+                        for t in forcefield._templateSignatures[signature]:
+                            if t in templatesNoEP:
+                                matches = _matchResidue(residueNoEP, templatesNoEP[t], bondedToAtomNoEP)
+                                if matches is not None:
+                                    template = t;
+                                    # Record the corresponding atoms.
+                                    matchingAtoms = {}
+                                    for atom, match in zip(residueNoEP.atoms(), matches):
+                                        templateAtomName = t.atoms[match].name
+                                        for templateAtom in template.atoms:
+                                            if templateAtom.name == templateAtomName:
+                                                matchingAtoms[templateAtom] = atom
+                                    break
+                    if template is None:
+                        raise ValueError('Residue %d (%s) does not match any template defined by the ForceField.' % (residue.index+1, residue.name))
+                    
+                    # Add the regular atoms.
+                    
+                    for atom in residue.atoms():
+                        if atom.element is not None:
+                            newAtoms[atom] = newTopology.addAtom(atom.name, atom.element, newResidue)
+                            newPositions.append(deepcopy(self.positions[atom.index]))
+                    
+                    # Add the extra points.
+                    
+                    templateAtomPositions = len(template.atoms)*[None]
+                    for index, atom in enumerate(template.atoms):
+                        if atom in matchingAtoms:
+                            templateAtomPositions[index] = self.positions[matchingAtoms[atom].index]
+                    for index, atom in enumerate(template.atoms):
+                        if atom.element is None:
+                            newTopology.addAtom(atom.name, None, newResidue)
+                            position = None
+                            for site in template.virtualSites:
+                                if site.index == index:
+                                    # This is a virtual site.  Compute its position by the correct rule.
+                                    
+                                    if site.type == 'average2':
+                                        position = site.weights[0]*templateAtomPositions[index+site.atoms[0]] + site.weights[1]*templateAtomPositions[index+site.atoms[1]]
+                                    elif site.type == 'average3':
+                                        position = site.weights[0]*templateAtomPositions[index+site.atoms[0]] + site.weights[1]*templateAtomPositions[index+site.atoms[1]] + site.weights[2]*templateAtomPositions[index+site.atoms[2]]
+                                    elif site.type == 'outOfPlane':
+                                        position = site.weights[0]*templateAtomPositions[index+site.atoms[0]] + site.weights[1]*templateAtomPositions[index+site.atoms[1]] + site.weights[2]*templateAtomPositions[index+site.atoms[2]]
+                            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.
+                                
+                                knownPositions = [x for x in templateAtomPositions if x is not None]
+                                position = sum(knownPositions)/len(knownPositions)
+                            newPositions.append(position)
+        for bond in self.topology.bonds():
+            if bond[0] in newAtoms and bond[1] in newAtoms:
+                newTopology.addBond(newAtoms[bond[0]], newAtoms[bond[1]])
+        self.topology = newTopology
+        self.positions = newPositions
\ No newline at end of file