Commit 2acba7ad authored by peastman's avatar peastman
Browse files

Merge pull request #1302 from rafwiewiora/pdb_extra_particles

Allow specification of extra particles in PDB files via element field and extraParticleIdentifier
parents 5eb3f8ad ef0ce25a
......@@ -123,7 +123,7 @@ class PdbStructure(object):
"""
def __init__(self, input_stream, load_all_models=False):
def __init__(self, input_stream, load_all_models=False, extraParticleIdentifier='EP'):
"""Create a PDB model from a PDB file stream.
Parameters
......@@ -135,9 +135,12 @@ class PdbStructure(object):
load_all_models : bool
Whether to load every model of an NMR structure or trajectory, or
just load the first model, to save memory.
extraParticleIdentifier : string='EP'
if this value appears in the element column for an ATOM record, the Atom's element will be set to 'EP' to mark it as an extra particle
"""
# initialize models
self.load_all_models = load_all_models
self.extraParticleIdentifier = extraParticleIdentifier
self.models = []
self._current_model = None
self.default_model = None
......@@ -157,7 +160,7 @@ class PdbStructure(object):
pdb_line = pdb_line.decode('utf-8')
# Look for atoms
if (pdb_line.find("ATOM ") == 0) or (pdb_line.find("HETATM") == 0):
self._add_atom(Atom(pdb_line, self))
self._add_atom(Atom(pdb_line, self, self.extraParticleIdentifier))
# Notice MODEL punctuation, for the next level of detail
# in the structure->model->chain->residue->atom->position hierarchy
elif (pdb_line.find("MODEL") == 0):
......@@ -682,7 +685,7 @@ class Residue(object):
class Atom(object):
"""Atom represents one atom in a PDB structure.
"""
def __init__(self, pdb_line, pdbstructure=None):
def __init__(self, pdb_line, pdbstructure=None, extraParticleIdentifier='EP'):
"""Create a new pdb.Atom from an ATOM or HETATM line.
Example line:
......@@ -795,6 +798,9 @@ class Atom(object):
try: self.formal_charge = int(pdb_line[78:80])
except ValueError: self.formal_charge = None
# figure out atom element
if self.element_symbol == extraParticleIdentifier:
self.element = 'EP'
else:
try:
# Try to find a sensible element symbol from columns 76-77
self.element = element.get_by_symbol(self.element_symbol)
......
......@@ -59,7 +59,7 @@ class PDBFile(object):
_residueNameReplacements = {}
_atomNameReplacements = {}
def __init__(self, file):
def __init__(self, file, extraParticleIdentifier='EP'):
"""Load a PDB file.
The atom positions and Topology can be retrieved by calling getPositions() and getTopology().
......@@ -68,6 +68,8 @@ class PDBFile(object):
----------
file : string
the name of the file to load
extraParticleIdentifier : string='EP'
if this value appears in the element column for an ATOM record, the Atom's element will be set to None to mark it as an extra particle
"""
metalElements = ['Al','As','Ba','Ca','Cd','Ce','Co','Cs','Cu','Dy','Fe','Gd','Hg','Ho','In','Ir','K','Li','Mg',
......@@ -91,7 +93,7 @@ class PDBFile(object):
if isinstance(file, str):
inputfile = open(file)
own_handle = True
pdb = PdbStructure(inputfile, load_all_models=True)
pdb = PdbStructure(inputfile, load_all_models=True, extraParticleIdentifier=extraParticleIdentifier)
if own_handle:
inputfile.close()
PDBFile._loadNameReplacementTables()
......@@ -116,7 +118,9 @@ class PDBFile(object):
atomName = atomReplacements[atomName]
atomName = atomName.strip()
element = atom.element
if element is None:
if element == 'EP':
element = None
elif element is None:
# Try to guess the element.
upper = atomName.upper()
......@@ -253,7 +257,7 @@ class PDBFile(object):
map[atom.attrib[id]] = name
@staticmethod
def writeFile(topology, positions, file=sys.stdout, keepIds=False):
def writeFile(topology, positions, file=sys.stdout, keepIds=False, extraParticleIdentifier=' '):
"""Write a PDB file containing a single model.
Parameters
......@@ -269,9 +273,11 @@ class PDBFile(object):
rather than generating new ones. Warning: It is up to the caller to
make sure these are valid IDs that satisfy the requirements of the
PDB format. Otherwise, the output file will be invalid.
extraParticleIdentifier : string=' '
String to write in the element column of the ATOM records for atoms whose element is None (extra particles)
"""
PDBFile.writeHeader(topology, file)
PDBFile.writeModel(topology, positions, file, keepIds=keepIds)
PDBFile.writeModel(topology, positions, file, keepIds=keepIds, extraParticleIdentifier=extraParticleIdentifier)
PDBFile.writeFooter(topology, file)
@staticmethod
......@@ -294,7 +300,7 @@ class PDBFile(object):
a*10, b*10, c*10, alpha*RAD_TO_DEG, beta*RAD_TO_DEG, gamma*RAD_TO_DEG), file=file)
@staticmethod
def writeModel(topology, positions, file=sys.stdout, modelIndex=None, keepIds=False):
def writeModel(topology, positions, file=sys.stdout, modelIndex=None, keepIds=False, extraParticleIdentifier=' '):
"""Write out a model to a PDB file.
Parameters
......@@ -313,6 +319,8 @@ class PDBFile(object):
rather than generating new ones. Warning: It is up to the caller to
make sure these are valid IDs that satisfy the requirements of the
PDB format. Otherwise, the output file will be invalid.
extraParticleIdentifier : string=' '
String to write in the element column of the ATOM records for atoms whose element is None (extra particles)
"""
if len(list(topology.atoms())) != len(positions):
......@@ -343,17 +351,17 @@ class PDBFile(object):
else:
resId = "%4d" % ((resIndex+1)%10000)
for atom in res.atoms():
if len(atom.name) < 4 and atom.name[:1].isalpha() and (atom.element is None or len(atom.element.symbol) < 2):
if atom.element is not None:
symbol = atom.element.symbol
else:
symbol = extraParticleIdentifier
if len(atom.name) < 4 and atom.name[:1].isalpha() and len(symbol) < 2:
atomName = ' '+atom.name
elif len(atom.name) > 4:
atomName = atom.name[:4]
else:
atomName = atom.name
coords = positions[posIndex]
if atom.element is not None:
symbol = atom.element.symbol
else:
symbol = ' '
line = "ATOM %5d %-4s %3s %s%4s %s%s%s 1.00 0.00 %2s " % (
atomIndex%100000, atomName, resName, chainName, resId, _format_83(coords[0]),
_format_83(coords[1]), _format_83(coords[2]), symbol)
......
......@@ -70,6 +70,20 @@ class TestPdbFile(unittest.TestCase):
pdb = PDBFile(open('systems/triclinic.pdb', 'rb'))
self.assertEqual(len(pdb.positions), 8)
def test_ExtraParticles(self):
"""Test reading, and writing and re-reading of a file containing extra particle atoms."""
pdb = PDBFile('systems/tip5p.pdb')
for atom in pdb.topology.atoms():
if atom.index > 2:
self.assertEqual(None, atom.element)
output = StringIO()
PDBFile.writeFile(pdb.topology, pdb.positions, output)
input = StringIO(output.getvalue())
pdb = PDBFile(input, extraParticleIdentifier = '')
for atom in pdb.topology.atoms():
if atom.index > 2:
self.assertEqual(None, atom.element)
def assertVecAlmostEqual(self, p1, p2, tol=1e-7):
unit = p1.unit
......
CRYST1 30.000 30.000 30.000 90.00 90.00 90.00 P 1 1
ATOM 1 O HOH A 1 20.612 20.351 16.218 1.00 0.00 O
ATOM 2 H1 HOH A 1 20.565 20.405 15.264 1.00 0.00 H
ATOM 3 H2 HOH A 1 21.273 19.679 16.387 1.00 0.00 H
ATOM 4 D1 HOH A 1 20.808 20.964 16.492 1.00 0.00 EP
ATOM 5 D2 HOH A 1 19.994 20.162 16.486 1.00 0.00 EP
TER 6 HOH A 1
END
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment