Residue templates can specify constraints (#5197)

* Residue templates can specify constraints

* Patched template generation preserves constraints

docs-source/usersguide/application/06_creating_ffs.rst

@@ -151,6 +151,15 @@ contains the following tags:
   (Alternatively, the deprecated :code:`from` tag may indicate the atom by
   index instead of name.)
 
+The :code:`<Residue>` tag may also contain :code:`<Constraint>` tags to specify
+pairs of atoms whose distance should be constrained.  Each one should include the
+attibutes :code:`atomName1` and :code:`atomName2` with the names of the two connected
+atoms, and a :code:`distance` attribute with the distance between them in nanometers.
+:code:`<Constraint>` tags are only used for constraints that are required by the
+force field and should always be included.  The more common situation is that a bond
+can be either flexible or constrained based on the options the user specifies when
+calling :code:`createSystem()`.  Bonds of that sort should not have :code:`<Constraint>`
+tags.
 
 The :code:`<Residue>` tag may also contain :code:`<VirtualSite>` tags,
 as in the following example:

wrappers/python/openmm/app/forcefield.py

@@ -325,6 +325,8 @@ class ForceField(object):
                             template.addExternalBondByName(bond.attrib['atomName'])
                         else:
                             template.addExternalBond(int(bond.attrib['from']))
+                    for constraint in residue.findall('Constraint'):
+                        template.addConstraintByName(constraint.attrib['atomName1'], constraint.attrib['atomName2'], float(constraint.attrib['distance']))
                     for patch in residue.findall('AllowPatch'):
                         patchName = patch.attrib['name']
                         if ':' in patchName:
@@ -700,6 +702,7 @@ class ForceField(object):
             self.virtualSites = []
             self.bonds = []
             self.externalBonds = []
+            self.constraints = []
             self.overrideLevel = 0
             self.rigidWater = True
             self.attributes = {}
@@ -741,6 +744,16 @@ class ForceField(object):
             atom = self.getAtomIndexByName(atom_name)
             self.addExternalBond(atom)
 
+        def addConstraint(self, atom1, atom2, distance):
+            """Add a constrained distance between two atoms in a template given their indices in the template."""
+            self.constraints.append((atom1, atom2, distance))
+
+        def addConstraintByName(self, atom1_name, atom2_name, distance):
+            """Add a constrained distance between two atoms in a template given their atom names."""
+            atom1 = self.getAtomIndexByName(atom1_name)
+            atom2 = self.getAtomIndexByName(atom2_name)
+            self.addConstraint(atom1, atom2, distance)
+
         def areParametersIdentical(self, template2, matchingAtoms, matchingAtoms2):
             """Get whether this template and another one both assign identical atom types and parameters to all atoms.
 
@@ -757,6 +770,8 @@ class ForceField(object):
             atoms2 = [template2.atoms[m] for m in matchingAtoms2]
             if any(a1.type != a2.type or a1.parameters != a2.parameters for a1,a2 in zip(atoms1, atoms2)):
                 return False
+            if set(self.constraints) != set(template2.constraints):
+                return False
             # Properly comparing virtual sites really needs a much more complicated analysis.  This simple check
             # could easily fail for templates containing vsites, even if they're actually identical.  Since we
             # currently have no force fields that include both patches and vsites, I'm not going to worry about it now.
@@ -913,6 +928,14 @@ class ForceField(object):
                 for atom in self.addedExternalBonds:
                     newTemplate.addExternalBondByName(atom.name)
 
+                # Build the list of constraints.
+
+                for atom1, atom2, distance in template.constraints:
+                    a1 = template.atoms[atom1]
+                    a2 = template.atoms[atom2]
+                    if a1 in atomMap and a2 in atomMap:
+                        newTemplate.addConstraint(atomMap[a1], atomMap[a2], distance)
+
                 # Add new virtual sites.
 
                 indexMap = dict((i, newAtomIndex[atom.name]) for i, atom in enumerate(self.addedAtoms[index]+self.changedAtoms[index]))
@@ -1221,7 +1244,8 @@ class ForceField(object):
             The cutoff distance to use for nonbonded interactions
         constraints : object=None
             Specifies which bonds and angles should be implemented with constraints.
-            Allowed values are None, HBonds, AllBonds, or HAngles.
+            Allowed values are None, HBonds, AllBonds, or HAngles.  Regardless of this value,
+            any constraints that are explicitly specified by the force field will be added.
         rigidWater : boolean=None
             If true, water molecules will be fully rigid regardless of the value
             passed for the constraints argument.  If None (the default), it uses the
@@ -1381,11 +1405,27 @@ class ForceField(object):
                 atom1 = data.atoms[bond.atom1]
                 atom2 = data.atoms[bond.atom2]
                 bond.isConstrained = atom1.element is elem.hydrogen or atom2.element is elem.hydrogen
+        matchedTemplateForResidue = {}
+        for residue, template in templateForResidue.items():
+            if isinstance(residue, MergedResidue):
+                for res in residue.residues:
+                    matchedTemplateForResidue[res] = template
+            else:
+                matchedTemplateForResidue[residue] = template
         for bond in data.bonds:
             atom1 = data.atoms[bond.atom1]
             atom2 = data.atoms[bond.atom2]
             if rigidResidue[atom1.residue.index] and rigidResidue[atom2.residue.index]:
                 bond.isConstrained = True
+            else:
+                if atom1.residue not in matchedTemplateForResidue:
+                    # This can happen with multi-residue patches, which would need more complex logic to handle.
+                    continue
+                t1 = data.atomTemplateIndexes[atom1]
+                t2 = data.atomTemplateIndexes[atom2]
+                for constraint in template.constraints:
+                    if sorted((t1, t2)) == sorted((constraint[0], constraint[1])):
+                        bond.isConstrained = True
 
         # Identify angles that should be implemented with constraints
 
@@ -1438,6 +1478,28 @@ class ForceField(object):
             if 'postprocessSystem' in dir(force):
                 force.postprocessSystem(sys, data, args)
 
+        # Add constraints specified in templates.  We do this at the very end so they will override
+        # constraints from bonds or angles.
+
+        existingConstraints = {}
+        for i in range(sys.getNumConstraints()):
+            p1, p2, d = sys.getConstraintParameters(i)
+            existingConstraints[(min(p1, p2), max(p1, p2))] = i
+        for residue, template in templateForResidue.items():
+            for constraint in template.constraints:
+                atoms = list(residue.atoms())
+                atom1 = [a for a in atoms if data.atomTemplateIndexes[a] == constraint[0]]
+                atom2 = [a for a in atoms if data.atomTemplateIndexes[a] == constraint[1]]
+                if len(atom1) != 1 or len(atom2) != 1:
+                    raise ValueError('Failed to identify atoms for constraint')
+                a1 = atom1[0].index
+                a2 = atom2[0].index
+                key = (min(a1, a2), max(a1, a2))
+                if key in existingConstraints:
+                    sys.setConstraintParameters(existingConstraints[key], a1, a2, constraint[2])
+                else:
+                    existingConstraints[key] = sys.addConstraint(a1, a2, constraint[2])
+
         # Execute scripts found in the XML files.
 
         for script in self._scripts:
@@ -1514,6 +1576,10 @@ class ForceField(object):
                 if recordParameters:
                     data.recordMatchedAtomParameters(matchedRes, template, matches)
                 templateForResidue[matchedRes] = template
+                if isinstance(matchedRes, MergedResidue):
+                    for res in matchedRes.residues:
+                        if res in templateForResidue:
+                            del templateForResidue[res]
         return templateForResidue
 
 
@@ -3601,21 +3667,6 @@ def getAtomPrint(data, atomIndex):
     return returnString
 
 
-def countConstraint(data):
-
-    bondCount = 0
-    angleCount = 0
-    for bond in data.bonds:
-        if bond.isConstrained:
-            bondCount += 1
-
-    angleCount = 0
-    for (angle, isConstrained) in zip(data.angles, data.isAngleConstrained):
-        if (isConstrained):
-            angleCount += 1
-
-    print("Constraints bond=%d angle=%d  total=%d" % (bondCount, angleCount, (bondCount+angleCount)))
-
 ## @private
 class AmoebaBondGenerator(object):
     """An AmoebaBondGenerator constructs a AmoebaBondForce."""

wrappers/python/tests/TestForceField.py

@@ -205,6 +205,63 @@ class TestForceField(unittest.TestCase):
         # Make sure flexibleConstraints yields just as many angles as no constraints
         self.assertEqual(af2.getNumAngles(), af3.getNumAngles())
 
+    def testTemplateConstraints(self):
+        """Test constraints defined by a residue template."""
+        xml = """
+<ForceField>
+ <AtomTypes>
+  <Type name="C" class="C" element="C" mass="12.01078"/>
+  <Type name="O" class="O" element="O" mass="15.99943"/>
+  <Type name="H" class="H" element="H" mass="1.007947"/>
+  <Type name="Na+" class="Na+" element="Na" mass="22.99"/>
+  <Type name="Cl-" class="Cl-" element="Cl" mass="35.45"/>
+ </AtomTypes>
+ <Residues>
+  <Residue name="MEOH">
+   <Atom name="CB" type="C"/>
+   <Atom name="OG" type="O"/>
+   <Atom name="HG" type="H"/>
+   <Atom name="HB1" type="H"/>
+   <Atom name="HB2" type="H"/>
+   <Atom name="HB3" type="H"/>
+   <Bond atomName1="CB" atomName2="OG"/>
+   <Bond atomName1="CB" atomName2="HB1"/>
+   <Bond atomName1="CB" atomName2="HB2"/>
+   <Bond atomName1="CB" atomName2="HB3"/>
+   <Bond atomName1="OG" atomName2="HG"/>
+   <Constraint atomName1="CB" atomName2="OG" distance="1.987"/>
+   <Constraint atomName1="OG" atomName2="HG" distance="1.123"/>
+  </Residue>
+  <Residue name="NA">
+    <Atom name="NA" type="Na+"/>
+  </Residue>
+  <Residue name="CL">
+    <Atom name="CL" type="Cl-"/>
+  </Residue>
+ </Residues>
+ <HarmonicBondForce>
+  <Bond class1="C" class2="O" k="100.0" length="2.0"/>
+  <Bond class1="C" class2="H" k="100.0" length="1.0"/>
+  <Bond class1="O" class2="H" k="100.0" length="1.1"/>
+ </HarmonicBondForce>
+</ForceField>"""
+        ff = ForceField(StringIO(xml))
+        pdb = PDBFile('systems/methanol_ions.pdb')
+        expected = {None:2, HBonds:5, AllBonds:5}
+        for constraints in [None, HBonds, AllBonds]:
+            system = ff.createSystem(pdb.topology, constraints=constraints)
+            self.assertEqual(expected[constraints], system.getNumConstraints())
+            lengths = {}
+            for i in range(system.getNumConstraints()):
+                p1, p2, length = system.getConstraintParameters(i)
+                lengths[(min(p1, p2), max(p1, p2))] = length.value_in_unit(nanometer)
+            self.assertEqual(1.987, lengths[(0, 1)])
+            self.assertEqual(1.123, lengths[(1, 2)])
+            if constraints is not None:
+                self.assertEqual(1.0, lengths[(0, 3)])
+                self.assertEqual(1.0, lengths[(0, 4)])
+                self.assertEqual(1.0, lengths[(0, 5)])
+
     def test_ImplicitSolvent(self):
         """Test the four types of implicit solvents using the implicitSolvent
         parameter.