Major mutation code cleanup

pdbfixer/pdbfixer.py

@@ -742,41 +742,40 @@ class PDBFixer(object):
         index_to_old_name = dict((r.index, r.name) for r in self.topology.residues())
         index_to_old_name = dict((r.index, r.name) for r in self.topology.residues())
         index_to_new_residues = {}
         index_to_new_residues = {}
 
 
-        chain_numbers = list()
-        resSeq_to_index = dict()
+        # NOTE: Multiple chains may have the same chainid, but must have unique resSeq entries.
+        resSeq_to_residue = dict() # resSeq_to_residue[resid] is the residue in the requested chain corresponding to residue identifier 'resid'
         for (chain_number, chain) in enumerate(self.topology.chains()):
         for (chain_number, chain) in enumerate(self.topology.chains()):
             if chain.id == chain_id:
             if chain.id == chain_id:
-                chain_numbers.append(chain_id)
                 for (residue_number, residue) in enumerate(chain.residues()):
                 for (residue_number, residue) in enumerate(chain.residues()):
-                    resSeq_to_index[int(residue.id)] = residue_number
+                    resSeq_to_residue[int(residue.id)] = residue
 
 
+        residue_map = dict() # residue_map[residue] is the name of the new residue to mutate to, if a mutation is desired
         for mut_str in mutations:
         for mut_str in mutations:
             old_name, resSeq, new_name = mut_str.split("-")
             old_name, resSeq, new_name = mut_str.split("-")
             resSeq = int(resSeq)
             resSeq = int(resSeq)
-            index = resSeq_to_index[resSeq]
 
 
-            if index not in index_to_old_name:
-                raise(KeyError("Cannot find index %d in system!" % index))
+            if resSeq not in resSeq_to_residue:
+                raise(KeyError("Cannot find chain %s residue %d in system!" % (chain_id, resSeq)))
 
 
-            if index_to_old_name[index] != old_name:
-                raise(ValueError("You asked to mutate %s %d, but that residue is actually %s!" % (old_name, index, index_to_old_name[index])))
+            residue = resSeq_to_residue[resSeq] # retrieve the requested residue
+
+            if residue.name != old_name:
+                raise(ValueError("You asked to mutate chain %s residue %d name %s, but that residue is actually %s!" % (chain_id, resSeq, old_name, residue.name)))
 
 
             try:
             try:
                 template = self.templates[new_name]
                 template = self.templates[new_name]
             except KeyError:
             except KeyError:
                 raise(KeyError("Cannot find residue %s in template library!" % new_name))
                 raise(KeyError("Cannot find residue %s in template library!" % new_name))
 
 
-            index_to_new_residues[index] = new_name
-
-
-        residue_map = [(r, index_to_new_residues[r.index]) for r in self.topology.residues() if r.index in index_to_new_residues]
+            # Store mutation
+            residue_map[residue] = new_name
 
 
+        # If there are mutations to be made, make them.
         if len(residue_map) > 0:            
         if len(residue_map) > 0:            
-            deleteAtoms = []
+            deleteAtoms = [] # list of atoms to delete
 
 
             # Find atoms that should be deleted.
             # Find atoms that should be deleted.
-
-            for residue, replaceWith in residue_map:
+            for residue, replaceWith in residue_map.iteritems():
                 if residue.chain.index != chain_number:
                 if residue.chain.index != chain_number:
                     continue  # Only modify specified chain
                     continue  # Only modify specified chain
                 residue.name = replaceWith
                 residue.name = replaceWith
@@ -786,7 +785,7 @@ class PDBFixer(object):
                     if atom.element in (None, hydrogen) or atom.name not in standardAtoms:
                     if atom.element in (None, hydrogen) or atom.name not in standardAtoms:
                         deleteAtoms.append(atom)
                         deleteAtoms.append(atom)
 
 
-            # Delete them.
+            # Delete atoms queued to be deleted.
             modeller = app.Modeller(self.topology, self.positions)
             modeller = app.Modeller(self.topology, self.positions)
             modeller.delete(deleteAtoms)
             modeller.delete(deleteAtoms)
             self.topology = modeller.topology
             self.topology = modeller.topology