Implemented matching of residues to single-residue patches

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

@@ -206,8 +206,8 @@ class ForceField(object):
                             template.addExternalBond(int(bond.attrib['from']))
                     for patch in residue.findall('AllowPatch'):
                         patchName = patch.attrib['name']
-                        if ':' in name:
-                            colonIndex = name.find(':')
+                        if ':' in patchName:
+                            colonIndex = patchName.find(':')
                             self.registerTemplatePatch(resName, patchName[:colonIndex], int(patchName[colonIndex+1:])-1)
                         else:
                             self.registerTemplatePatch(resName, patchName, 0)
@@ -488,6 +488,16 @@ class ForceField(object):
                 self.constraints[key] = distance
                 system.addConstraint(atom1, atom2, distance)
         
+        def recordMatchedAtomParameters(self, residue, template, matches):
+            """Record parameters for atoms based on having matched a residue to a template."""
+            matchAtoms = dict(zip(matches, residue.atoms()))
+            for atom, match in zip(residue.atoms(), matches):
+                self.atomType[atom] = template.atoms[match].type
+                self.atomParameters[atom] = template.atoms[match].parameters
+                for site in template.virtualSites:
+                    if match == site.index:
+                        self.virtualSites[atom] = (site, [matchAtoms[i].index for i in site.atoms], matchAtoms[site.excludeWith].index)
+
     class _TemplateData(object):
         """Inner class used to encapsulate data about a residue template definition."""
         def __init__(self, name):
@@ -611,15 +621,17 @@ class ForceField(object):
                 
                 for atom in template.atoms:
                     if not any(deleted.name == atom.name and deleted.residue == index for deleted in self.deletedAtoms):
-                        newTemplate.atoms.append(deepcopy(atom))
+                        newTemplate.atoms.append(ForceField._TemplateAtomData(atom.name, atom.type, atom.element, atom.parameters))
                 for atom in self.addedAtoms[index]:
-                    newTemplate.atoms.append(deepcopy(atom))
+                    if any(a.name == atom.name for a in newTemplate.atoms):
+                        raise ValueError("Patch '%s' adds an atom with the same name as an existing atom: %s" % (self.name, atom.name))
+                    newTemplate.atoms.append(ForceField._TemplateAtomData(atom.name, atom.type, atom.element, atom.parameters))
                 oldAtomIndex = dict([(atom.name, i) for i, atom in enumerate(template.atoms)])
                 newAtomIndex = dict([(atom.name, i) for i, atom in enumerate(newTemplate.atoms)])
                 for atom in self.changedAtoms[index]:
                     if atom.name not in newAtomIndex:
                         raise ValueError("Patch '%s' modifies nonexistent atom '%s' in template '%s'" % (self.name, atom.name, template.name))
-                    newTemplate.atoms[newAtomIndex[atom.name]] = deepcopy(atom)
+                    newTemplate.atoms[newAtomIndex[atom.name]] = ForceField._TemplateAtomData(atom.name, atom.type, atom.element, atom.parameters)
                 
                 # Copy over the virtual sites, translating the atom indices.
                 
@@ -643,7 +655,7 @@ class ForceField(object):
                 deletedExternalBonds = [atom.name for atom in self.deletedExternalBonds if atom.residue == index]
                 for atom in template.externalBonds:
                     if template.atoms[atom].name not in deletedExternalBonds:
-                        newTemplate.addExternalBond(atomMap[atom])
+                        newTemplate.addExternalBond(indexMap[atom])
                 for atom1, atom2 in self.addedBonds:
                     if atom1.residue == index and atom2.residue == index:
                         newTemplate.addBondByName(atom1.name, atom2.name)
@@ -745,7 +757,7 @@ class ForceField(object):
                 raise ValueError('%s: No parameters defined for atom type %s' % (self.forceName, t))
 
 
-    def _getResidueTemplateMatches(self, res, bondedToAtom):
+    def _getResidueTemplateMatches(self, res, bondedToAtom, templateSignatures=None):
         """Return the residue template matches, or None if none are found.
 
         Parameters
@@ -766,10 +778,12 @@ class ForceField(object):
         """
         template = None
         matches = None
+        if templateSignatures is None:
+            templateSignatures = self._templateSignatures
         signature = _createResidueSignature([atom.element for atom in res.atoms()])
-        if signature in self._templateSignatures:
+        if signature in templateSignatures:
             allMatches = []
-            for t in self._templateSignatures[signature]:
+            for t in templateSignatures[signature]:
                 match = _matchResidue(res, t, bondedToAtom)
                 if match is not None:
                     allMatches.append((t, match))
@@ -975,8 +989,8 @@ class ForceField(object):
             data.atomBonds[bond.atom2].append(i)
 
         # Find the template matching each residue and assign atom types.
-        # If no templates are found, attempt to use residue template generators to create new templates (and potentially atom types/parameters).
 
+        unmatchedResidues = []
         for chain in topology.chains():
             for res in chain.residues():
                 if res in residueTemplates:
@@ -989,7 +1003,23 @@ class ForceField(object):
                     # Attempt to match one of the existing templates.
                     [template, matches] = self._getResidueTemplateMatches(res, bondedToAtom)
                 if matches is None:
-                    # No existing templates match.  Try any registered residue template generators.
+                    unmatchedResidues.append(res)
+                else:
+                    data.recordMatchedAtomParameters(res, template, matches)
+        
+        # Try to apply patches to find matches for any unmatched residues.
+        
+        if len(unmatchedResidues) > 0:
+            unmatchedResidues = _applyPatchesToMatchResidues(self, data, unmatchedResidues, bondedToAtom)
+        
+        # If we still haven't found a match for a residue, attempt to use residue template generators to create
+        # new templates (and potentially atom types/parameters).
+
+        for res in unmatchedResidues:
+            # A template might have been generated on an earlier iteration of this loop.
+            [template, matches] = self._getResidueTemplateMatches(res, bondedToAtom)
+            if matches is None:
+                # Try all generators.
                 for generator in self._templateGenerators:
                     if generator(self, res):
                         # This generator has registered a new residue template that should match.
@@ -1000,19 +1030,10 @@ class ForceField(object):
                         else:
                             # We successfully generated a residue template.  Break out of the for loop.
                             break
-
-                # Raise an exception if we have found no templates that match.
             if matches is None:
                 raise ValueError('No template found for residue %d (%s).  %s' % (res.index+1, res.name, _findMatchErrors(self, res)))
-
-                # Store parameters for the matched residue template.
-                matchAtoms = dict(zip(matches, res.atoms()))
-                for atom, match in zip(res.atoms(), matches):
-                    data.atomType[atom] = template.atoms[match].type
-                    data.atomParameters[atom] = template.atoms[match].parameters
-                    for site in template.virtualSites:
-                        if match == site.index:
-                            data.virtualSites[atom] = (site, [matchAtoms[i].index for i in site.atoms], matchAtoms[site.excludeWith].index)
+            else:
+                data.recordMatchedAtomParameters(res, template, matches)
 
         # Create the System and add atoms
 
@@ -1289,6 +1310,58 @@ def _findAtomMatches(atoms, template, bondedTo, externalBonds, matches, hasMatch
     return False
 
 
+def _applyPatchesToMatchResidues(forcefield, data, residues, bondedToAtom):
+    """Try to apply patches to find matches for residues."""
+    # Start by creating all templates than can be created by applying a combination of one-residue patches
+    # to a single template.  The number of these is usually not too large, and they often cover a large fraction
+    # of residues.
+    
+    patchedTemplateSignatures = {}
+    patchedTemplates = {}
+    for name, template in forcefield._templates.items():
+        if name in forcefield._templatePatches:
+            patches = [forcefield._patches[patchName] for patchName, patchResidueIndex in forcefield._templatePatches[name] if forcefield._patches[patchName].numResidues == 1]
+            if len(patches) > 0:
+                newTemplates = []
+                patchedTemplates[name] = newTemplates
+                _generatePatchedSingleResidueTemplates(template, patches, 0, newTemplates)
+                for patchedTemplate in newTemplates:
+                    signature = _createResidueSignature([atom.element for atom in patchedTemplate.atoms])
+                    if signature in patchedTemplateSignatures:
+                        patchedTemplateSignatures[signature].append(patchedTemplate)
+                    else:
+                        patchedTemplateSignatures[signature] = [patchedTemplate]
+    
+    # Now see if any of those templates matches any of the residues.
+    
+    unmatchedResidues = []
+    for res in residues:
+        [template, matches] = forcefield._getResidueTemplateMatches(res, bondedToAtom, patchedTemplateSignatures)
+        if matches is None:
+            unmatchedResidues.append(res)
+        else:
+            data.recordMatchedAtomParameters(res, template, matches)
+    return unmatchedResidues
+
+
+def _generatePatchedSingleResidueTemplates(template, patches, index, newTemplates):
+    """Apply all possible combinations of a set of single-residue patches to a template."""
+    try:
+        patchedTemplate = patches[index].createPatchedTemplates([template])[0]
+        newTemplates.append(patchedTemplate)
+    except:
+        # This probably means the patch is inconsistent with another one that has already been applied,
+        # so just ignore it.
+        patchedTemplate = None
+    
+    # Call this function recursively to generate combinations of patches.
+    
+    if index+1 < len(patches):
+        _generatePatchedSingleResidueTemplates(template, patches, index+1, newTemplates)
+        if patchedTemplate is not None:
+            _generatePatchedSingleResidueTemplates(patchedTemplate, patches, index+1, newTemplates)
+
+
 def _findMatchErrors(forcefield, res):
     """Try to guess why a residue failed to match any template and return an error message."""
     residueCounts = _countResidueAtoms([atom.element for atom in res.atoms()])

wrappers/python/tests/TestPatches.py

@@ -176,5 +176,90 @@ class TestForceField(unittest.TestCase):
         self.assertEqual(indexMap['B'], v.atoms[0])
         self.assertEqual(indexMap['C'], v.atoms[1])
 
+    def testAlaAlaAla(self):
+        """Test constructing a System that involves two patches."""
+
+        xml = """
+<ForceField>
+ <AtomTypes>
+  <Type name="N" class="N" element="N" mass="14.00672"/>
+  <Type name="H" class="H" element="H" mass="1.007947"/>
+  <Type name="CT" class="CT" element="C" mass="12.01078"/>
+  <Type name="H1" class="H1" element="H" mass="1.007947"/>
+  <Type name="HC" class="HC" element="H" mass="1.007947"/>
+  <Type name="C" class="C" element="C" mass="12.01078"/>
+  <Type name="O" class="O" element="O" mass="15.99943"/>
+  <Type name="O2" class="O2" element="O" mass="15.99943"/>
+  <Type name="N3" class="N3" element="N" mass="14.00672"/>
+ </AtomTypes>
+ <Residues>
+  <Residue name="ALA">
+   <Atom name="N" type="N"/>
+   <Atom name="H" type="H"/>
+   <Atom name="CA" type="CT"/>
+   <Atom name="HA" type="H1"/>
+   <Atom name="CB" type="CT"/>
+   <Atom name="HB1" type="HC"/>
+   <Atom name="HB2" type="HC"/>
+   <Atom name="HB3" type="HC"/>
+   <Atom name="C" type="C"/>
+   <Atom name="O" type="O"/>
+   <Bond from="0" to="1"/>
+   <Bond from="0" to="2"/>
+   <Bond from="2" to="3"/>
+   <Bond from="2" to="4"/>
+   <Bond from="2" to="8"/>
+   <Bond from="4" to="5"/>
+   <Bond from="4" to="6"/>
+   <Bond from="4" to="7"/>
+   <Bond from="8" to="9"/>
+   <ExternalBond from="0"/>
+   <ExternalBond from="8"/>
+   <AllowPatch name="CTER"/>
+   <AllowPatch name="NTER"/>
+  </Residue>
+ </Residues>
+ <Patches>
+  <Patch name="CTER">
+    <AddAtom name="OXT" type="O2"/>
+    <ChangeAtom name="O" type="O2"/>
+    <AddBond atomName1="C" atomName2="OXT"/>
+    <RemoveExternalBond atomName="C"/>
+  </Patch>
+  <Patch name="NTER">
+    <RemoveAtom name="H"/>
+    <AddAtom name="H1" type="H"/>
+    <AddAtom name="H2" type="H"/>
+    <AddAtom name="H3" type="H"/>
+    <ChangeAtom name="N" type="N3"/>
+    <RemoveBond atomName1="N" atomName2="H"/>
+    <AddBond atomName1="N" atomName2="H1"/>
+    <AddBond atomName1="N" atomName2="H2"/>
+    <AddBond atomName1="N" atomName2="H3"/>
+    <RemoveExternalBond atomName="N"/>
+  </Patch>
+ </Patches>
+ <NonbondedForce coulomb14scale="0.833333" lj14scale="0.5">
+  <Atom type="N" charge="-0.4157" sigma="0.324999852378" epsilon="0.71128"/>
+  <Atom type="H" charge="0.2719" sigma="0.106907846177" epsilon="0.0656888"/>
+  <Atom type="CT" charge="0.0337" sigma="0.339966950842" epsilon="0.4577296"/>
+  <Atom type="H1" charge="0.0823" sigma="0.247135304412" epsilon="0.0656888"/>
+  <Atom type="HC" charge="0.0603" sigma="0.264953278775" epsilon="0.0656888"/>
+  <Atom type="C" charge="0.5973" sigma="0.339966950842" epsilon="0.359824"/>
+  <Atom type="O" charge="-0.5679" sigma="0.295992190115" epsilon="0.87864"/>
+  <Atom type="O2" charge="-0.8055" sigma="0.295992190115" epsilon="0.87864"/>
+  <Atom type="N3" charge="0.1414" sigma="0.324999852378" epsilon="0.71128"/>
+ </NonbondedForce>
+</ForceField>"""
+        ff = ForceField(StringIO(xml))
+        pdb = PDBFile(os.path.join('systems', 'ala_ala_ala.pdb'))
+        system = ff.createSystem(pdb.topology)
+        nb = system.getForce(0)
+        expectedCharges = [0.1414, 0.2719, 0.2719, 0.2719, 0.0337, 0.0823, 0.0337, 0.0603, 0.0603, 0.0603, 0.5973, -0.5679,
+                           -0.4157, 0.2719, 0.0337, 0.0823, 0.0337, 0.0603, 0.0603, 0.0603, 0.5973, -0.5679,
+                           0.5973, -0.8055, -0.8055, -0.4157, 0.2719, 0.0337, 0.0823, 0.0337, 0.0603, 0.0603, 0.0603]
+        for i in range(system.getNumParticles()):
+            self.assertEqual(expectedCharges[i], nb.getParticleParameters(i)[0].value_in_unit(elementary_charge))
+
 if __name__ == '__main__':
     unittest.main()