diff --git a/pdbfixer/pdbfixer.py b/pdbfixer/pdbfixer.py index 431b021..20c5a30 100644 --- a/pdbfixer/pdbfixer.py +++ b/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.