Merge pull request #1302 from rafwiewiora/pdb_extra_particles

Allow specification of extra particles in PDB files via element field and extraParticleIdentifier

wrappers/python/simtk/openmm/app/internal/pdbstructure.py

@@ -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)

wrappers/python/simtk/openmm/app/pdbfile.py

@@ -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)

wrappers/python/tests/TestPdbFile.py

@@ -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

wrappers/python/tests/systems/tip5p.pdb

+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