Beginning support for force field patches

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

@@ -6,7 +6,7 @@ Simbios, the NIH National Center for Physics-Based Simulation of
 Biological Structures at Stanford, funded under the NIH Roadmap for
 Medical Research, grant U54 GM072970. See https://simtk.org.
 
-Portions copyright (c) 2012-2015 Stanford University and the Authors.
+Portions copyright (c) 2012-2016 Stanford University and the Authors.
 Authors: Peter Eastman, Mark Friedrichs
 Contributors:
 
@@ -115,6 +115,8 @@ class ForceField(object):
         """
         self._atomTypes = {}
         self._templates = {}
+        self._patches = {}
+        self._templatePatches = {}
         self._templateSignatures = {None:[]}
         self._atomClasses = {'':set()}
         self._forces = []
@@ -201,8 +203,81 @@ class ForceField(object):
                             template.addExternalBondByName(bond.attrib['atomName'])
                         else:
                             template.addExternalBond(int(bond.attrib['from']))
+                    for patch in patch.findall('AllowPatch'):
+                        patchName = patch.attrib['name']
+                        if ':' in name:
+                            colonIndex = name.find(':')
+                            self.registerTemplatePatch(resName, patchName[:colonIndex], int(patchName[colonIndex+1:])-1)
+                        else:
+                            self.registerTemplatePatch(resName, patchName, 0)
                     self.registerResidueTemplate(template)
 
+        # Load the patch defintions.
+
+        for tree in trees:
+            if tree.getroot().find('Patches') is not None:
+                for patch in tree.getroot().find('Patches').findall('Patch'):
+                    patchName = patch.attrib['name']
+                    if 'residues' in patch.attrib:
+                        numResidues = int(patch.attrib['residues'])
+                    else:
+                        numResidues = 1
+                    patchData = ForceField._PatchData(patchName, numResidues)
+                    allAtomNames = set()
+                    for atom in patch.findall('AddAtom'):
+                        params = {}
+                        for key in atom.attrib:
+                            if key not in ('name', 'type'):
+                                params[key] = _convertParameterToNumber(atom.attrib[key])
+                        atomName = atom.attrib['name']
+                        if atomName in allAtomNames:
+                            raise ValueError('Patch '+patchName+' contains multiple atoms named '+atomName)
+                        allAtomNames.add(atomName)
+                        atomDescription = ForceField._PatchAtomData(atomName)
+                        typeName = atom.attrib['type']
+                        patchData.addedAtoms.append(ForceField._TemplateAtomData(atomDescription, typeName, self._atomTypes[typeName].element, params))
+                    for atom in patch.findall('ChangeAtom'):
+                        params = {}
+                        for key in atom.attrib:
+                            if key not in ('name', 'type'):
+                                params[key] = _convertParameterToNumber(atom.attrib[key])
+                        atomName = atom.attrib['name']
+                        if atomName in allAtomNames:
+                            raise ValueError('Patch '+patchName+' contains multiple atoms named '+atomName)
+                        allAtomNames.add(atomName)
+                        atomDescription = ForceField._PatchAtomData(atomName)
+                        typeName = atom.attrib['type']
+                        patchData.changedAtoms.append(ForceField._TemplateAtomData(atomDescription, typeName, self._atomTypes[typeName].element, params))
+                    for atom in patch.findall('RemoveAtom'):
+                        atomName = atom.attrib['name']
+                        if atomName in allAtomNames:
+                            raise ValueError('Patch '+patchName+' contains multiple atoms named '+atomName)
+                        allAtomNames.add(atomName)
+                        atomDescription = ForceField._PatchAtomData(atomName)
+                        patchData.deletedAtoms.append(atomDescription)
+                    for bond in patch.findall('AddBond'):
+                        atom1 = ForceField._PatchAtomData(bond.attrib['atomName1'])
+                        atom2 = ForceField._PatchAtomData(bond.attrib['atomName2'])
+                        patchData.addedBonds.append((atom1, atom2))
+                    for bond in patch.findall('RemoveBond'):
+                        atom1 = ForceField._PatchAtomData(bond.attrib['atomName1'])
+                        atom2 = ForceField._PatchAtomData(bond.attrib['atomName2'])
+                        patchData.deletedBonds.append((atom1, atom2))
+                    for bond in patch.findall('AddExternalBond'):
+                        atom = ForceField._PatchAtomData(bond.attrib['atomName'])
+                        patchData.addedExternalBonds.append(atom)
+                    for bond in patch.findall('RemoveExternalBond'):
+                        atom = ForceField._PatchAtomData(bond.attrib['atomName'])
+                        patchData.deletedExternalBonds.append(atom)
+                    for residue in patch.findall('ApplyToResidue'):
+                        name = residue.attrib['name']
+                        if ':' in name:
+                            colonIndex = name.find(':')
+                            self.registerTemplatePatch(name[colonIndex+1:], patchName, int(name[:colonIndex])-1)
+                        else:
+                            self.registerTemplatePatch(name, patchName, 0)
+                    self.registerPatch(patchData)
+
         # Load force definitions
 
         for tree in trees:
@@ -272,6 +347,16 @@ class ForceField(object):
         else:
             self._templateSignatures[signature] = [template]
 
+    def registerPatch(self, patch):
+        """Register a new patch that can be applied to templates."""
+        self._patches[patch.name] = patch
+    
+    def registerTemplatePatch(self, residue, patch, patchResidueIndex):
+        """Register that a particular patch can be used with a particular residue."""
+        if residue not in self._templatePatches:
+            self._templatePatches[residue] = []
+        self._templatePatches[residue].append((patch, patchResidueIndex))
+
     def registerScript(self, script):
         """Register a new script to be executed after building the System."""
         self._scripts.append(script)
@@ -496,6 +581,30 @@ class ForceField(object):
             else:
                 self.excludeWith = self.atoms[0]
 
+    class _PatchData(object):
+        """Inner class used to encapsulate data about a patch definition."""
+        def __init__(self, name, numResidues):
+            self.name = name
+            self.numResidues = numResidues
+            self.addedAtoms = []
+            self.deletedAtoms = []
+            self.changedAtoms = []
+            self.addedBonds = []
+            self.deletedBonds = []
+            self.addedExternalBonds = []
+            self.deletedExternalBonds = []
+            
+    class _PatchAtomData(object):
+        """Inner class used to encapsulate data about an atom in a patch definition."""
+        def __init__(self, description):
+            if ':' in description:
+                colonIndex = description.find(':')
+                self.residue = int(description[:colonIndex])-1
+                self.name = description[colonIndex+1:]
+            else:
+                self.residue = 0
+                self.name = description
+
     class _AtomType(object):
         """Inner class used to record atom types and associated properties."""
         def __init__(self, name, atomClass, mass, element):

wrappers/python/tests/TestPatches.py

+import unittest
+from simtk.openmm.app import *
+from simtk.openmm import *
+from simtk.unit import *
+import simtk.openmm.app.forcefield as forcefield
+import math
+try:
+    from cStringIO import StringIO
+except ImportError:
+    from io import StringIO
+import os
+
+class TestForceField(unittest.TestCase):
+    """Test ForceFields that use patches."""
+
+    def testParsePatch(self):
+        """Test parsing a <Patch> tag."""
+
+        xml = """
+<ForceField>
+ <AtomTypes>
+  <Type name="A type" class="A class" element="O" mass="15.99943"/>
+  <Type name="B type" class="B class" element="H" mass="1.007947"/>
+ </AtomTypes>
+ <Patches>
+  <Patch name="Test">
+    <AddAtom name="A" type="A type"/>
+    <ChangeAtom name="B" type="B type"/>
+    <RemoveAtom name="C"/>
+    <AddBond atomName1="A" atomName2="B"/>
+    <RemoveBond atomName1="B" atomName2="C"/>
+    <AddExternalBond atomName="A"/>
+    <RemoveExternalBond atomName="C"/>
+    <ApplyToResidue name="RES"/>
+  </Patch>
+ </Patches>
+</ForceField>"""
+        ff = ForceField(StringIO(xml))
+        self.assertEqual(1, len(ff._patches))
+        patch = ff._patches['Test']
+        self.assertEqual(1, len(patch.addedAtoms))
+        self.assertEqual(1, len(patch.changedAtoms))
+        self.assertEqual(1, len(patch.deletedAtoms))
+        self.assertEqual(1, len(patch.addedBonds))
+        self.assertEqual(1, len(patch.deletedBonds))
+        self.assertEqual(1, len(patch.addedExternalBonds))
+        self.assertEqual(1, len(patch.deletedExternalBonds))
+        self.assertEqual(1, len(ff._templatePatches))
+        self.assertEqual(1, len(ff._templatePatches['RES']))
+        self.assertEqual('A', patch.addedAtoms[0].name.name)
+        self.assertEqual(0, patch.addedAtoms[0].name.residue)
+        self.assertEqual('A type', patch.addedAtoms[0].type)
+        self.assertEqual('B', patch.changedAtoms[0].name.name)
+        self.assertEqual(0, patch.changedAtoms[0].name.residue)
+        self.assertEqual('B type', patch.changedAtoms[0].type)
+        self.assertEqual('C', patch.deletedAtoms[0].name)
+        self.assertEqual(0, patch.deletedAtoms[0].residue)
+        self.assertEqual('A', patch.addedBonds[0][0].name)
+        self.assertEqual(0, patch.addedBonds[0][0].residue)
+        self.assertEqual('B', patch.addedBonds[0][1].name)
+        self.assertEqual(0, patch.addedBonds[0][1].residue)
+        self.assertEqual('B', patch.deletedBonds[0][0].name)
+        self.assertEqual(0, patch.deletedBonds[0][0].residue)
+        self.assertEqual('C', patch.deletedBonds[0][1].name)
+        self.assertEqual(0, patch.deletedBonds[0][1].residue)
+        self.assertEqual('A', patch.addedExternalBonds[0].name)
+        self.assertEqual(0, patch.addedExternalBonds[0].residue)
+        self.assertEqual('C', patch.deletedExternalBonds[0].name)
+        self.assertEqual(0, patch.deletedExternalBonds[0].residue)
+        self.assertEqual('Test', ff._templatePatches['RES'][0][0])
+        self.assertEqual(0, ff._templatePatches['RES'][0][1])
+
+if __name__ == '__main__':
+    unittest.main()