Bugfixed chain ID assignment in addSolvent (openmm#287)

pdbfixer/pdbfixer.py

@@ -1094,15 +1094,31 @@ def addSolvent(self, boxSize=None, padding=None, boxVectors=None, positiveIon='N
 
         modeller = app.Modeller(self.topology, self.positions)
         forcefield = self._createForceField(self.topology, True)
+        original_chains = list(modeller.topology.chains())
         modeller.addSolvent(forcefield, padding=padding, boxSize=boxSize, boxVectors=boxVectors, boxShape=boxShape, positiveIon=positiveIon, negativeIon=negativeIon, ionicStrength=ionicStrength)
-        chains = list(modeller.topology.chains())
-        if len(chains) == 1:
-            chains[0].id = 'A'
-        else:
-            chains[-1].id = chr(ord(chains[-2].id)+1)
+        new_chains = list(modeller.topology.chains())
+        new_chain_indices = range(len(original_chains), len(new_chains))
+        assert 1 <= len(new_chain_indices) <= 2, "Expected a new solvent chain and optionally a new ion chain"
+
+        for new_chain_index in new_chain_indices:
+            new_chain = new_chains[new_chain_index]
+            assert new_chain.id.isnumeric(), "Expected numeric chain ID assigned by modeller.addSolvent"
+
+            # If the last chain ID is a letter between A and Y, assign the next letter(s) for the new chain ID(s).
+            if new_chain_index > 0:
+                prev_chain_id = new_chains[new_chain_index - 1].id
+                if len(prev_chain_id) == 1 and prev_chain_id.isalpha():
+                    abet_index = ord(prev_chain_id.upper()) - ord("A") + 1
+                    if abet_index < 26:
+                        new_chain.id = chr(ord("A") + abet_index)
+                    else:
+                        # Otherwise, leave the ID as a string representing its index.
+                        pass
+
         self.topology = modeller.topology
         self.positions = modeller.positions
 
+
     def addMembrane(self, lipidType='POPC', membraneCenterZ=0*unit.nanometer, minimumPadding=1*unit.nanometer, positiveIon='Na+', negativeIon='Cl-', ionicStrength=0*unit.molar):
         """Add a lipid membrane to the structure.