Allow custom template matchers for residues (#3161)

wrappers/python/openmm/app/forcefield.py

@@ -208,6 +208,7 @@ class ForceField(object):
         self._atomClasses = {'':set()}
         self._forces = []
         self._scripts = []
+        self._templateMatchers = []
         self._templateGenerators = []
         self.loadFile(files)
 
@@ -478,6 +479,21 @@ class ForceField(object):
     def registerScript(self, script):
         """Register a new script to be executed after building the System."""
         self._scripts.append(script)
+    
+    def registerTemplateMatcher(self, matcher):
+        """Register an object that can override the default logic for matching templates to residues.
+
+        A template matcher is a callable object that can be invoked as::
+        
+            template = f(forcefield, residue)
+        
+        where ``forcefield`` is the ForceField invoking it and ``residue`` is a openmm.app.Residue object.
+        It should return a _TemplateData object that matches the residue.  Alternatively it may return
+        None, in which case the standard logic will be used to find a template for the residue.
+
+        .. CAUTION:: This method is experimental, and its API is subject to change.        
+        """
+        self._templateMatchers.append(matcher)
 
     def registerTemplateGenerator(self, generator):
         """Register a residue template generator that can be used to parameterize residues that do not match existing forcefield templates.
@@ -957,6 +973,13 @@ class ForceField(object):
         """
         template = None
         matches = None
+        for matcher in self._templateMatchers:
+            template = matcher(self, res)
+            if template is not None:
+                match = compiled.matchResidueToTemplate(res, template, bondedToAtom, ignoreExternalBonds, ignoreExtraParticles)
+                if match is None:
+                    raise ValueError('A custom template matcher returned a template for residue %s, but it does not match the residue.' % res.name)
+                return [template, match]
         if templateSignatures is None:
             templateSignatures = self._templateSignatures
         signature = _createResidueSignature([atom.element for atom in res.atoms()])

wrappers/python/tests/TestForceField.py

@@ -391,6 +391,64 @@ class TestForceField(unittest.TestCase):
                 self.assertEqual(params[1], 1.0*nanometers)
                 self.assertEqual(params[2], 0.0*kilojoule_per_mole)
 
+    def test_residueMatcher(self):
+        """Test using a custom template matcher to select templates."""
+        xml = """
+<ForceField>
+ <AtomTypes>
+  <Type name="tip3p-O" class="OW" element="O" mass="15.99943"/>
+  <Type name="tip3p-H" class="HW" element="H" mass="1.007947"/>
+ </AtomTypes>
+ <Residues>
+  <Residue name="HOH">
+   <Atom name="O" type="tip3p-O" charge="-0.834"/>
+   <Atom name="H1" type="tip3p-H" charge="0.417"/>
+   <Atom name="H2" type="tip3p-H" charge="0.417"/>
+   <Bond from="0" to="1"/>
+   <Bond from="0" to="2"/>
+   <Bond from="1" to="2"/>
+  </Residue>
+  <Residue name="HOH2">
+   <Atom name="O" type="tip3p-O" charge="0.834"/>
+   <Atom name="H1" type="tip3p-H" charge="-0.417"/>
+   <Atom name="H2" type="tip3p-H" charge="-0.417"/>
+   <Bond from="0" to="1"/>
+   <Bond from="0" to="2"/>
+   <Bond from="1" to="2"/>
+  </Residue>
+ </Residues>
+ <NonbondedForce coulomb14scale="0.833333" lj14scale="0.5">
+  <UseAttributeFromResidue name="charge"/>
+  <Atom type="tip3p-O" sigma="0.315" epsilon="0.635"/>
+  <Atom type="tip3p-H" sigma="1" epsilon="0"/>
+ </NonbondedForce>
+</ForceField>"""
+        ff = ForceField(StringIO(xml))
+
+        # Load a water box.
+        prmtop = AmberPrmtopFile('systems/water-box-216.prmtop')
+        top = prmtop.topology
+        
+        # Building a System should fail, because two templates match each residue.
+        self.assertRaises(Exception, lambda: ff.createSystem(top))
+        
+        # Register a template matcher that selects a particular one.
+        def matcher(ff, res):
+            return ff._templates['HOH2']
+        ff.registerTemplateMatcher(matcher)
+        
+        # It should now succeed in building a System.
+        system = ff.createSystem(top)
+        
+        # Make sure it used the correct parameters.
+        nb = [f for f in system.getForces() if isinstance(f, NonbondedForce)][0]
+        for atom in top.atoms():
+            charge, sigma, epsilon = nb.getParticleParameters(atom.index)
+            if atom.name == 'O':
+                self.assertEqual(0.834*elementary_charge, charge)
+            else:
+                self.assertEqual(-0.417*elementary_charge, charge)
+
     def test_residueTemplateGenerator(self):
         """Test the ability to add residue template generators to parameterize unmatched residues."""
         def simpleTemplateGenerator(forcefield, residue):