Support removal of Hydrogens when using gaff forcefield

planckton/init.py

@@ -117,11 +117,6 @@ def __init__(
         remove_hydrogen_atoms=False,
         foyer_kwargs={"assert_dihedral_params": False},
     ):
-        if ff == FORCEFIELD["gaff"] and remove_hydrogen_atoms == True:
-            raise NotImplementedError(
-                "Removing hydrogens is not supported with the GAFF forcefield"
-            )
-
         if not isinstance(compound, (list, set)):
             self.compound = [compound]
         else:
@@ -160,14 +155,6 @@ def __init__(
         self.L = self._calculate_L()
         self.foyer_kwargs = foyer_kwargs
 
-    def _remove_hydrogen(self):
-        for subcompound in self.compound:
-            for atom in subcompound.particles():
-                if atom.name in ["_hc", "_ha", "_h1", "_h4"]:
-                    # NOTE: May not be a comprehensive list of
-                    # all hydrogen types.
-                    subcompound.remove(atom)
-
     def pack(self, box_expand_factor=5):
         """Pack compounds into a larger box in preparation for shrinking.
 
@@ -182,9 +169,6 @@ def pack(self, box_expand_factor=5):
         typed_system : ParmEd structure
             ParmEd structure of filled box
         """
-        if self.remove_hydrogen_atoms:
-            self._remove_hydrogen()
-
         L = self.L.value * box_expand_factor
         box = mb.Box([L, L, L])
         system = mb.packing.fill_box(
@@ -197,6 +181,10 @@ def pack(self, box_expand_factor=5):
         system.box = box
         pmd_system = system.to_parmed(residues=[self.residues])
         typed_system = self.ff.apply(pmd_system, **self.foyer_kwargs)
+        if self.remove_hydrogen_atoms:
+            typed_system.strip(
+                [a.atomic_number == 1 for a in typed_system.atoms]
+            )
         return typed_system
 
     def _calculate_L(self):

planckton/tests/test_hydrogen_removal.py

@@ -18,15 +18,8 @@ def test_hydrogen_removal():
         remove_hydrogen_atoms=True,
     )
 
-    packer._remove_hydrogen()
-
-    for atom in packer.compound[0].particles():
-        assert atom.name not in [
-            "_hc",
-            "_ha",
-            "_h1",
-            "_h4",
-        ], "Hydrogen found in system!"
+    system = packer.pack()
+    assert 1 not in [a.atomic_number for a in system.atoms]
 
 
 def test_hydrogen_removal_and_sim():
@@ -52,6 +45,21 @@ def test_hydrogen_removal_and_sim():
     my_sim.run()
 
 
+def test_hydrogen_remove_gaff():
+    p3ht = Compound("c1cscc1CCCCCC")
+    p3ht_Hs = [h for h in p3ht.particles_by_element("H")]
+    packer = Pack(
+        p3ht,
+        ff=FORCEFIELD["gaff"],
+        n_compounds=2,
+        density=0.01 * u.g / u.cm ** 3,
+        remove_hydrogen_atoms=True,
+    )
+    system = packer.pack()
+    assert "H" not in [a.name for a in system.atoms]
+    assert p3ht.n_particles * 2 - len(system.atoms) == len(p3ht_Hs) * 2
+
+
 if __name__ == "__main__":
     if path.isfile("restart.gsd"):
         remove("restart.gsd")

planckton/tests/test_sim.py

@@ -48,16 +48,27 @@ def test_smiles_gaff(self):
             target_length=packer.L,
         )
 
-    def test_gaff_noH_raises(self):
+    def test_gaff_noH(self):
         p3ht = Compound("c1cscc1CCCCCC")
-        with pytest.raises(NotImplementedError):
-            packer = Pack(
-                p3ht,
-                ff=FORCEFIELD["gaff"],
-                n_compounds=2,
-                density=0.01 * u.g / u.cm ** 3,
-                remove_hydrogen_atoms=True,
-            )
+        packer = Pack(
+            p3ht,
+            ff=FORCEFIELD["gaff"],
+            n_compounds=2,
+            density=0.01 * u.g / u.cm ** 3,
+            remove_hydrogen_atoms=True,
+        )
+        system = packer.pack()
+        my_sim = Simulation(
+            system,
+            kT=3.0,
+            gsd_write=1e2,
+            log_write=1e2,
+            e_factor=1,
+            n_steps=3e3,
+            mode="cpu",
+            shrink_steps=1e3,
+            target_length=packer.L,
+        )
 
     def test_smiles_opvgaff_raises(self):
         p3ht = Compound("c1cscc1CCCCCC")