Fixed errors with residueTemplates arguments to Modeller methods (#4528)

wrappers/python/openmm/app/modeller.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-2023 Stanford University and the Authors.
+Portions copyright (c) 2012-2024 Stanford University and the Authors.
 Authors: Peter Eastman
 Contributors: 
 
@@ -856,12 +856,15 @@ class Modeller(object):
         newTopology.setPeriodicBoxVectors(self.topology.getPeriodicBoxVectors())
         newAtoms = {}
         newPositions = []*nanometer
+        newResidueTemplates = {}
         newIndices = []
         acceptors = [atom for atom in self.topology.atoms() if atom.element in (elem.oxygen, elem.nitrogen)]
         for chain in self.topology.chains():
             newChain = newTopology.addChain(chain.id)
             for residue in chain.residues():
                 newResidue = newTopology.addResidue(residue.name, newChain, residue.id, residue.insertionCode)
+                if residue in residueTemplates:
+                    newResidueTemplates[newResidue] = residueTemplates[residue]
                 isNTerminal = (residue == chain._residues[0])
                 isCTerminal = (residue == chain._residues[-1])
                 if residue.name in Modeller._residueHydrogens:
@@ -1012,8 +1015,7 @@ class Modeller(object):
         if forcefield is not None:
             # Use the ForceField the user specified.
 
-            system = forcefield.createSystem(newTopology, rigidWater=False, nonbondedMethod=CutoffNonPeriodic, residueTemplates=residueTemplates)
-            atoms = list(newTopology.atoms())
+            system = forcefield.createSystem(newTopology, rigidWater=False, nonbondedMethod=CutoffNonPeriodic, residueTemplates=newResidueTemplates)
             for i in range(system.getNumParticles()):
                 if i not in addedH:
                     # Existing atom, make it immobile.
@@ -1128,11 +1130,14 @@ class Modeller(object):
         newTopology.setPeriodicBoxVectors(self.topology.getPeriodicBoxVectors())
         newAtoms = {}
         newPositions = []*nanometer
+        newResidueTemplates = {}
         missingPositions = set()
         for chain in self.topology.chains():
             newChain = newTopology.addChain(chain.id)
             for residue in chain.residues():
                 newResidue = newTopology.addResidue(residue.name, newChain, residue.id, residue.insertionCode)
+                if residue in residueTemplates:
+                    newResidueTemplates[newResidue] = residueTemplates[residue]
                 template = templates[residue.index]
                 if len(template.atoms) == len(list(residue.atoms())):
                     # Just copy the residue over.
@@ -1233,7 +1238,7 @@ class Modeller(object):
             # There were particles whose position we couldn't identify before, since they were neither virtual sites nor Drude particles.
             # Try to figure them out based on bonds.  First, use the ForceField to create a list of every bond involving one of them.
 
-            system = forcefield.createSystem(newTopology, constraints=AllBonds, residueTemplates=residueTemplates)
+            system = forcefield.createSystem(newTopology, constraints=AllBonds, residueTemplates=newResidueTemplates)
             bonds = []
             for i in range(system.getNumConstraints()):
                 bond = system.getConstraintParameters(i)
@@ -1580,6 +1585,10 @@ class Modeller(object):
             if len(toDelete) > 0:
                 modeller.delete(toDelete)
 
+        newResidueTemplates = {}
+        for r1, r2 in zip(self.topology.residues(), modeller.topology.residues()):
+            if r1 in residueTemplates:
+                newResidueTemplates[r2] = residueTemplates[r1]
         self.topology = modeller.topology
         self.positions = modeller.positions
 
@@ -1620,7 +1629,7 @@ class Modeller(object):
             if lowerZBoundary < waterZ.value_in_unit(nanometer) < upperZBoundary:
                 del waterPos[wRes]
 
-        self._addIons(forcefield, numTotalWaters, waterPos, positiveIon=positiveIon, negativeIon=negativeIon, ionicStrength=ionicStrength, neutralize=neutralize, residueTemplates=residueTemplates)
+        self._addIons(forcefield, numTotalWaters, waterPos, positiveIon=positiveIon, negativeIon=negativeIon, ionicStrength=ionicStrength, neutralize=neutralize, residueTemplates=newResidueTemplates)
 
 
 class _CellList(object):

wrappers/python/tests/TestModeller.py

@@ -1270,6 +1270,49 @@ class TestModeller(unittest.TestCase):
             {(None, None), (Single, 1.0), (Double, 2.0)}
         )
 
+    def test_residueTemplates(self):
+        """Test the residueTemplates argument to Modeller methods"""
+
+        # Create a Topology and ForceField involving residues that match multiple templates.
+
+        topology = Topology()
+        chain = topology.addChain()
+        residue1 = topology.addResidue('Fe', chain)
+        topology.addAtom('Fe', element.iron, residue1)
+        residue2 = topology.addResidue('Fe', chain)
+        topology.addAtom('Fe', element.iron, residue2)
+        positions = [Vec3(0, 0, 0), Vec3(1, 0, 0)]*nanometers
+        ff = ForceField('amber14/tip3pfb.xml', 'amber14/lipid17.xml')
+        residueTemplates = {residue1: 'FE2', residue2: 'FE'}
+
+        # Test addSolvent().
+
+        modeller = Modeller(topology, positions)
+        with self.assertRaises(Exception):
+            modeller.addSolvent(ff, padding=1*nanometer)
+        modeller.addSolvent(ff, padding=1*nanometer, residueTemplates=residueTemplates)
+
+        # Test addHydrogens().
+
+        modeller = Modeller(topology, positions)
+        with self.assertRaises(Exception):
+            modeller.addHydrogens(ff)
+        modeller.addHydrogens(ff, residueTemplates=residueTemplates)
+
+        # Test addExtraParticles().
+
+        modeller = Modeller(topology, positions)
+        with self.assertRaises(Exception):
+            modeller.addExtraParticles(ff)
+        modeller.addExtraParticles(ff, residueTemplates=residueTemplates)
+
+        # Test addMembrane().
+
+        modeller = Modeller(topology, positions)
+        with self.assertRaises(Exception):
+            modeller.addMembrane(ff)
+        modeller.addMembrane(ff, residueTemplates=residueTemplates)
+
     def assertVecAlmostEqual(self, p1, p2, tol=1e-7):
         scale = max(1.0, norm(p1),)
         for i in range(3):