Fixed some bugs due to surprising behavior of structure in which same chain ids…

Fixed some bugs due to surprising behavior of structure in which same chain ids can appear in multiple distinct chains.

pdbfixer/pdbfixer.py

@@ -412,12 +412,11 @@ class PDBFixer(object):
         if chainIndices == None:
             chainIndices = list()
         if chainIds != None:
-            # Convert chain ids to chain indices
-            which_model = 0
-            chain_id_to_chain_number = dict((c.chain_id, k) for k, c in enumerate(self.structure.models[which_model].chains))
-            for chainId in chainIds:
-                chainNumber = chain_id_to_chain_number[chainId]
-                chainIndices.append(chainNumber)
+            # 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))
 

tests/test_removechains.py

-from nose.tools import ok_, eq_, raises
+from nose.tools import ok_, eq_, raises, assert_items_equal
 import simtk.openmm.app as app
 import pdbfixer
 import tempfile
 
-def remove_chainids_and_verify(pdbid, chain_ids_to_remove, expected_chain_ids_remaining):
+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 = [ chain.id for chain in fixer.topology.chains() ]
     assertItemsEqual(chain_ids_remaining, expected_chain_ids_remaining)
 
-def test_removechainids():
-    remove_chainids_and_verify('4JSV', [], ['B', 'D', 'A', 'C'])
-    remove_chainids_and_verify('4JSV', ['B', 'D'], ['A', 'C'])
-    remove_chainids_and_verify('4JSV', ['A', 'C'], ['B', 'D'])
-    remove_chainids_and_verify('4JSV', ['B', 'A'], ['D', 'C'])
-    remove_chainids_and_verify('4JSV', ['B', 'D', 'A', 'C'], [])
+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], [])