Merge pull request #2448 from JoaoRodrigues/restrain_existing_H

Restraint existing atoms during addHydrogens

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

@@ -712,6 +712,8 @@ class Modeller(object):
         automatically, this function will only add hydrogens.  It will never remove ones that are already present in the
         model, regardless of the specified pH.
 
+        In all cases, the positions of existing atoms (including existing hydrogens) are not modified.
+
         Definitions for standard amino acids and nucleotides are built in.  You can call loadHydrogenDefinitions() to load
         additional definitions for other residue types.
 
@@ -930,14 +932,16 @@ class Modeller(object):
 
         # The hydrogens were added at random positions.  Now perform an energy minimization to fix them up.
 
+        addedH = set(newIndices)  # keep track of Hs added
+
         if forcefield is not None:
             # Use the ForceField the user specified.
 
             system = forcefield.createSystem(newTopology, rigidWater=False, nonbondedMethod=CutoffNonPeriodic)
             atoms = list(newTopology.atoms())
             for i in range(system.getNumParticles()):
-                if atoms[i].element != elem.hydrogen:
-                    # This is a heavy atom, so make it immobile.
+                if i not in addedH:
+                    # Existing atom, make it immobile.
                     system.setParticleMass(i, 0)
         else:
             # Create a System that restrains the distance of each hydrogen from its parent atom
@@ -955,7 +959,7 @@ class Modeller(object):
             bondedTo = []
             for atom in newTopology.atoms():
                 nonbonded.addParticle([])
-                if atom.element != elem.hydrogen:
+                if atom.index not in addedH:  # make immobile
                     system.addParticle(0.0)
                 else:
                     system.addParticle(1.0)

wrappers/python/tests/TestModeller.py

@@ -593,6 +593,51 @@ class TestModeller(unittest.TestCase):
 
         validate_equivalence(self, topology_start, topology_after)
 
+    def test_addHydrogensPdb3_keepPositions(self):
+        """ Test addHydrogens() does not change existing Hs positions """
+
+        # build the Modeller
+        topology_start = self.topology_start3
+        positions = self.positions3.value_in_unit(nanometers)
+        modeller = Modeller(topology_start, positions)
+
+        # Record original hydrogen positions
+        oriH = [atom.index for atom in modeller.topology.atoms() if atom.element == element.hydrogen]
+        oriH_pos = [positions[i] for i in oriH]
+
+        # Remove hydrogens from last residue
+        res_list = list(topology_start.residues())
+        toDelete = [atom for atom in res_list[-1].atoms() if atom.element == element.hydrogen]
+        modeller.delete(toDelete)
+
+        n_deleted = len(toDelete)
+
+        # Add hydrogen atoms back.
+        modeller.addHydrogens(self.forcefield)
+        topology_after = modeller.getTopology()
+
+        # Fetch 'new' positions
+        new_positions = modeller.positions.value_in_unit(nanometers)
+        newH = [atom.index for atom in topology_after.atoms() if atom.element == element.hydrogen]
+        newH_pos = [new_positions[i] for i in newH]
+
+        # Did we add all Hs back in correctly?
+        self.assertEqual(len(newH), len(oriH))
+
+        # Are the old ones at the same position?
+        # Negative control
+        oriH_fixed = oriH_pos[:-1*n_deleted]
+        newH_fixed = newH_pos[:-1*n_deleted]
+        xyz_diff = any([norm(o-n) > 1e-6 for o, n in zip(oriH_fixed, newH_fixed)])
+        self.assertEqual(xyz_diff, False)
+
+        # Were the new ones optimized?
+        # Positive control
+        oriH_added = oriH_pos[-1*n_deleted:]
+        newH_added = newH_pos[-1*n_deleted:]
+        xyz_diff = all([norm(o-n) > 1e-6 for o, n in zip(oriH_added, newH_added)])
+        self.assertEqual(xyz_diff, True)
+
     def test_addHydrogensASH(self):
         """ Test of addHydrogens() in which we force ASH to be a variant using the variants parameter. """