Merge pull request #52 from jchodera/fix-mutate

Add support for specifying chain ids in chains to remove [ready to merge]

pdbfixer/pdbfixer.py

@@ -382,13 +382,15 @@ class PDBFixer(object):
                 templatePosition = template.positions[atom.index].value_in_unit(unit.nanometer)
                 newPositions.append(mm.Vec3(*np.dot(rotate, templatePosition))*unit.nanometer+translate)
     
-    def removeChains(self, chainIndices):
+    def removeChains(self, chainIndices=None, chainIds=None):
         """Remove a set of chains from the structure.
         
         Parameters
         ----------
-        chainIndices : list of int
-            List of the indices of the chains to remove.
+        chainIndices : list of int, optional, default=None
+            List of indices of chains to remove.
+        chainIds : list of str, optional, default=None
+            List of chain ids of chains to remove.
 
         Examples
         --------
@@ -396,15 +398,38 @@ class PDBFixer(object):
         Load a PDB file with two chains and eliminate the second chain.
 
         >>> fixer = PDBFixer(pdbid='4J7F')
-        >>> fixer.removeChains([1])
+        >>> fixer.removeChains(chainIndices=[1])
+
+        Load a PDB file with two chains and eliminate chains named 'B' and 'D'.
+
+        >>> fixer = PDBFixer(pdbid='4J7F')
+        >>> fixer.removeChains(chainIds=['B','D'])
 
         """
         modeller = app.Modeller(self.topology, self.positions)
         allChains = list(self.topology.chains())
+
+        if chainIndices == None:
+            chainIndices = list()
+        if chainIds != None:
+            # Add all chains that match the selection to the list.
+            chain_id_list = [c.chain_id for c in self.structure.models[0].chains]
+            for (chainNumber, chainId) in enumerate(chain_id_list):
+                if chainId in chainIds:
+                    chainIndices.append(chainNumber)
+            # Ensure only unique entries remain.
+            chainIndices = list(set(chainIndices))
+
+        # Do nothing if no chains will be deleted.
+        if len(chainIndices) == 0:
+            return
+
         modeller.delete(allChains[i] for i in chainIndices)
         self.topology = modeller.topology
         self.positions = modeller.positions
         self.structureChains = [self.structureChains[i] for i in range(len(self.structureChains)) if i not in chainIndices]
+
+        return
     
     def findMissingResidues(self):
         """Find residues that are missing from the structure.

tests/test_removechains.py

+from nose.tools import ok_, eq_, raises, assert_items_equal
+import simtk.openmm.app as app
+import pdbfixer
+import tempfile
+
+def remove_chain_ids_and_verify(pdbid, chain_ids_to_remove, expected_chain_ids_remaining):
+    # Create a PDBFixer instance for the given pdbid
+    fixer = pdbfixer.PDBFixer(pdbid=pdbid)
+    # Remove specified chains.
+    fixer.removeChains(chainIds=chain_ids_to_remove)
+    # Check to make sure asserted chains remain.
+    chain_ids_remaining = [c.chain_id for c in fixer.structureChains]
+    assert_items_equal(chain_ids_remaining, expected_chain_ids_remaining)
+
+def test_removechain_ids():
+    remove_chain_ids_and_verify('4JSV', [], ['B', 'D', 'A', 'C', 'B', 'A'])
+    remove_chain_ids_and_verify('4JSV', ['B', 'D'], ['A', 'C', 'A'])
+    remove_chain_ids_and_verify('4JSV', ['A', 'C'], ['B', 'D', 'B'])
+    remove_chain_ids_and_verify('4JSV', ['B', 'A'], ['D', 'C'])
+    remove_chain_ids_and_verify('4JSV', ['B', 'D', 'A', 'C'], [])
+
+def remove_chain_indices_and_verify(pdbid, chain_indices_to_remove, expected_chain_ids_remaining):
+    # Create a PDBFixer instance for the given pdbid
+    fixer = pdbfixer.PDBFixer(pdbid=pdbid)
+    # Remove specified chains.
+    fixer.removeChains(chainIndices=chain_indices_to_remove)
+    # Check to make sure asserted chains remain.
+    chain_ids_remaining = [c.chain_id for c in fixer.structureChains]
+    assert_items_equal(chain_ids_remaining, expected_chain_ids_remaining)
+
+def test_removechain_indices():
+    remove_chain_indices_and_verify('4JSV', [], ['B', 'D', 'A', 'C', 'B', 'A'])
+    remove_chain_indices_and_verify('4JSV', [0, 1], ['A', 'C', 'B', 'A'])
+    remove_chain_indices_and_verify('4JSV', [2, 3], ['B', 'D', 'B', 'A'])
+    remove_chain_indices_and_verify('4JSV', [0, 2], ['D', 'C', 'B', 'A'])
+    remove_chain_indices_and_verify('4JSV', [0, 1, 2, 3, 4, 5], [])