Merge pull request #1081 from swails/pdbx

Add PDBxReporter class

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

 """
 OpenMM Application
 """
+from __future__ import absolute_import
+
 __docformat__ = "epytext en"
 
 __author__ = "Peter Eastman"
@@ -10,26 +12,26 @@ __license__ = "MIT"
 __maintainer__ = "Peter Eastman"
 __email__ = "peastman@stanford.edu"
 
-from topology import Topology, Chain, Residue, Atom
-from pdbfile import PDBFile
-from pdbxfile import PDBxFile
-from forcefield import ForceField
-from simulation import Simulation
-from pdbreporter import PDBReporter
-from amberprmtopfile import AmberPrmtopFile, HCT, OBC1, OBC2, GBn, GBn2
-from amberinpcrdfile import AmberInpcrdFile
-from dcdfile import DCDFile
-from gromacsgrofile import GromacsGroFile
-from gromacstopfile import GromacsTopFile
-from dcdreporter import DCDReporter
-from modeller import Modeller
-from statedatareporter import StateDataReporter
-from element import Element
-from desmonddmsfile import DesmondDMSFile
-from checkpointreporter import CheckpointReporter
-from charmmcrdfiles import CharmmCrdFile, CharmmRstFile
-from charmmparameterset import CharmmParameterSet
-from charmmpsffile import CharmmPsfFile, CharmmPSFWarning
+from .topology import Topology, Chain, Residue, Atom
+from .pdbfile import PDBFile
+from .pdbxfile import PDBxFile
+from .forcefield import ForceField
+from .simulation import Simulation
+from .pdbreporter import PDBReporter, PDBxReporter
+from .amberprmtopfile import AmberPrmtopFile, HCT, OBC1, OBC2, GBn, GBn2
+from .amberinpcrdfile import AmberInpcrdFile
+from .dcdfile import DCDFile
+from .gromacsgrofile import GromacsGroFile
+from .gromacstopfile import GromacsTopFile
+from .dcdreporter import DCDReporter
+from .modeller import Modeller
+from .statedatareporter import StateDataReporter
+from .element import Element
+from .desmonddmsfile import DesmondDMSFile
+from .checkpointreporter import CheckpointReporter
+from .charmmcrdfiles import CharmmCrdFile, CharmmRstFile
+from .charmmparameterset import CharmmParameterSet
+from .charmmpsffile import CharmmPsfFile, CharmmPSFWarning
 
 # Enumerated values
 

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

@@ -28,6 +28,7 @@ DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
 OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
 USE OR OTHER DEALINGS IN THE SOFTWARE.
 """
+from __future__ import print_function, division, absolute_import
 __author__ = "Peter Eastman"
 __version__ = "1.0"
 
@@ -39,12 +40,13 @@ from copy import copy
 from datetime import date
 from simtk.openmm import Vec3, Platform
 from simtk.openmm.app.internal.pdbstructure import PdbStructure
+from simtk.openmm.app.internal.unitcell import computeLengthsAndAngles
 from simtk.openmm.app import Topology
 from simtk.unit import nanometers, angstroms, is_quantity, norm, Quantity, dot
-import element as elem
+from . import element as elem
 try:
     import numpy
-except:
+except ImportError:
     pass
 
 class PDBFile(object):
@@ -252,17 +254,13 @@ class PDBFile(object):
          - topology (Topology) The Topology defining the molecular system being written
          - file (file=stdout) A file to write the file to
         """
-        print >>file, "REMARK   1 CREATED WITH OPENMM %s, %s" % (Platform.getOpenMMVersion(), str(date.today()))
+        print("REMARK   1 CREATED WITH OPENMM %s, %s" % (Platform.getOpenMMVersion(), str(date.today())), file=file)
         vectors = topology.getPeriodicBoxVectors()
         if vectors is not None:
-            (a, b, c) = vectors.value_in_unit(angstroms)
-            a_length = norm(a)
-            b_length = norm(b)
-            c_length = norm(c)
-            alpha = math.acos(dot(b, c)/(b_length*c_length))*180.0/math.pi
-            beta = math.acos(dot(c, a)/(c_length*a_length))*180.0/math.pi
-            gamma = math.acos(dot(a, b)/(a_length*b_length))*180.0/math.pi
-            print >>file, "CRYST1%9.3f%9.3f%9.3f%7.2f%7.2f%7.2f P 1           1 " % (a_length, b_length, c_length, alpha, beta, gamma)
+            a, b, c, alpha, beta, gamma = computeLengthsAndAngles(vectors)
+            RAD_TO_DEG = 180/math.pi
+            print("CRYST1%9.3f%9.3f%9.3f%7.2f%7.2f%7.2f P 1           1 " % (
+                    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):
@@ -288,7 +286,7 @@ class PDBFile(object):
         atomIndex = 1
         posIndex = 0
         if modelIndex is not None:
-            print >>file, "MODEL     %4d" % modelIndex
+            print("MODEL     %4d" % modelIndex, file=file)
         for (chainIndex, chain) in enumerate(topology.chains()):
             if keepIds:
                 chainName = chain.id
@@ -320,14 +318,14 @@ class PDBFile(object):
                         atomIndex%100000, atomName, resName, chainName, resId, _format_83(coords[0]),
                         _format_83(coords[1]), _format_83(coords[2]), symbol)
                     assert len(line) == 80, 'Fixed width overflow detected'
-                    print >>file, line
+                    print(line, file=file)
                     posIndex += 1
                     atomIndex += 1
                 if resIndex == len(residues)-1:
-                    print >>file, "TER   %5d      %3s %s%4s" % (atomIndex, resName, chainName, resId)
+                    print("TER   %5d      %3s %s%4s" % (atomIndex, resName, chainName, resId), file=file)
                     atomIndex += 1
         if modelIndex is not None:
-            print >>file, "ENDMDL"
+            print("ENDMDL", file=file)
 
     @staticmethod
     def writeFooter(topology, file=sys.stdout):
@@ -381,13 +379,13 @@ class PDBFile(object):
             for index1 in sorted(atomBonds):
                 bonded = atomBonds[index1]
                 while len(bonded) > 4:
-                    print >>file, "CONECT%5d%5d%5d%5d" % (index1, bonded[0], bonded[1], bonded[2])
+                    print("CONECT%5d%5d%5d%5d" % (index1, bonded[0], bonded[1], bonded[2]), file=file)
                     del bonded[:4]
                 line = "CONECT%5d" % index1
                 for index2 in bonded:
                     line = "%s%5d" % (line, index2)
-                print >>file, line
-        print >>file, "END"
+                print(line, file=file)
+        print("END", file=file)
 
 
 def _format_83(f):

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

@@ -32,7 +32,7 @@ __author__ = "Peter Eastman"
 __version__ = "1.0"
 
 import simtk.openmm as mm
-from simtk.openmm.app import PDBFile
+from simtk.openmm.app import PDBFile, PDBxFile
 
 class PDBReporter(object):
     """PDBReporter outputs a series of frames from a Simulation to a PDB file.
@@ -77,8 +77,34 @@ class PDBReporter(object):
             self._nextModel += 1
         PDBFile.writeModel(simulation.topology, state.getPositions(), self._out, self._nextModel)
         self._nextModel += 1
+        if hasattr(self._out, 'flush') and callable(self._out.flush):
+            self._out.flush()
 
     def __del__(self):
         if self._topology is not None:
             PDBFile.writeFooter(self._topology, self._out)
         self._out.close()
+
+class PDBxReporter(PDBReporter):
+    """PDBxReporter outputs a series of frames from a Simulation to a PDBx/mmCIF file.
+
+    To use it, create a PDBxReporter, then add it to the Simulation's list of reporters.
+    """
+
+    def report(self, simulation, state):
+        """Generate a report.
+
+        Parameters:
+         - simulation (Simulation) The Simulation to generate a report for
+         - state (State) The current state of the simulation
+        """
+        if self._nextModel == 0:
+            PDBxFile.writeHeader(simulation.topology, self._out)
+            self._nextModel += 1
+        PDBxFile.writeModel(simulation.topology, state.getPositions(), self._out, self._nextModel)
+        self._nextModel += 1
+        if hasattr(self._out, 'flush') and callable(self._out.flush):
+            self._out.flush()
+
+    def __del__(self):
+        self._out.close()

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

 """
-pdbfile.py: Used for loading PDB files.
+pdbxfile.py: Used for loading PDBx/mmCIF files.
 
 This is part of the OpenMM molecular simulation toolkit originating from
 Simbios, the NIH National Center for Physics-Based Simulation of
 Biological Structures at Stanford, funded under the NIH Roadmap for
 Medical Research, grant U54 GM072970. See https://simtk.org.
 
-Portions copyright (c) 2014-2015 Stanford University and the Authors.
+Portions copyright (c) 2015 Stanford University and the Authors.
 Authors: Peter Eastman
-Contributors:
+Contributors: Jason Swails
 
 Permission is hereby granted, free of charge, to any person obtaining a
 copy of this software and associated documentation files (the "Software"),
@@ -28,18 +28,20 @@ DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
 OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
 USE OR OTHER DEALINGS IN THE SOFTWARE.
 """
+from __future__ import division, absolute_import, print_function
+
 __author__ = "Peter Eastman"
-__version__ = "1.0"
+__version__ = "2.0"
 
-import os
 import sys
 import math
-from simtk.openmm import Vec3
+from simtk.openmm import Vec3, Platform
+from datetime import date
 from simtk.openmm.app.internal.pdbx.reader.PdbxReader import PdbxReader
-from simtk.openmm.app.internal.unitcell import computePeriodicBoxVectors
+from simtk.openmm.app.internal.unitcell import computePeriodicBoxVectors, computeLengthsAndAngles
 from simtk.openmm.app import Topology
-from simtk.unit import nanometers, angstroms, is_quantity, norm, Quantity
-import element as elem
+from simtk.unit import nanometers, angstroms, is_quantity, norm, Quantity, dot
+from . import element as elem
 try:
     import numpy
 except:
@@ -202,3 +204,120 @@ class PDBxFile(object):
                 self._numpyPositions[frame] = Quantity(numpy.array(self._positions[frame].value_in_unit(nanometers)), nanometers)
             return self._numpyPositions[frame]
         return self._positions[frame]
+
+    @staticmethod
+    def writeFile(topology, positions, file=sys.stdout, keepIds=False,
+                  entry=None):
+        """Write a PDBx/mmCIF file containing a single model.
+
+        Parameters:
+         - topology (Topology) The Topology defining the model to write
+         - positions (list) The list of atomic positions to write
+         - file (file=stdout) A file to write to
+         - keepIds (bool=False) If True, keep the residue and chain IDs specified in the Topology 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 PDBx/mmCIF format.  Otherwise, the output file will be invalid.
+         - entry (str=None) The entry ID to assign to the CIF file
+        """
+        PDBxFile.writeHeader(topology, file, entry)
+        PDBxFile.writeModel(topology, positions, file, keepIds=keepIds)
+
+    @staticmethod
+    def writeHeader(topology, file=sys.stdout, entry=None):
+        """Write out the header for a PDBx/mmCIF file.
+
+        Parameters:
+         - topology (Topology) The Topology defining the molecular system being written
+         - file (file=stdout) A file to write the file to
+         - entry (str=None) The entry ID to assign to the CIF file
+        """
+        if entry is not None:
+            print('data_%s' % entry, file=file)
+        else:
+            print('data_cell', file=file)
+        print("# Created with OpenMM %s, %s" % (Platform.getOpenMMVersion(), str(date.today())), file=file)
+        print('#', file=file)
+        vectors = topology.getPeriodicBoxVectors()
+        if vectors is not None:
+            a, b, c, alpha, beta, gamma = computeLengthsAndAngles(vectors)
+            RAD_TO_DEG = 180/math.pi
+            print('_cell.length_a     %10.4f' % (a*10), file=file)
+            print('_cell.length_b     %10.4f' % (b*10), file=file)
+            print('_cell.length_c     %10.4f' % (c*10), file=file)
+            print('_cell.angle_alpha  %10.4f' % (alpha*RAD_TO_DEG), file=file)
+            print('_cell.angle_beta   %10.4f' % (beta*RAD_TO_DEG), file=file)
+            print('_cell.angle_gamma  %10.4f' % (gamma*RAD_TO_DEG), file=file)
+            print('##', file=file)
+        print('loop_', file=file)
+        print('_atom_site.group_PDB', file=file)
+        print('_atom_site.id', file=file)
+        print('_atom_site.type_symbol', file=file)
+        print('_atom_site.label_atom_id', file=file)
+        print('_atom_site.label_alt_id', file=file)
+        print('_atom_site.label_comp_id', file=file)
+        print('_atom_site.label_asym_id', file=file)
+        print('_atom_site.label_entity_id', file=file)
+        print('_atom_site.label_seq_id', file=file)
+        print('_atom_site.pdbx_PDB_ins_code', file=file)
+        print('_atom_site.Cartn_x', file=file)
+        print('_atom_site.Cartn_y', file=file)
+        print('_atom_site.Cartn_z', file=file)
+        print('_atom_site.occupancy', file=file)
+        print('_atom_site.B_iso_or_equiv', file=file)
+        print('_atom_site.Cartn_x_esd', file=file)
+        print('_atom_site.Cartn_y_esd', file=file)
+        print('_atom_site.Cartn_z_esd', file=file)
+        print('_atom_site.occupancy_esd', file=file)
+        print('_atom_site.B_iso_or_equiv_esd', file=file)
+        print('_atom_site.pdbx_formal_charge', file=file)
+        print('_atom_site.auth_seq_id', file=file)
+        print('_atom_site.auth_comp_id', file=file)
+        print('_atom_site.auth_asym_id', file=file)
+        print('_atom_site.auth_atom_id', file=file)
+        print('_atom_site.pdbx_PDB_model_num', file=file)
+
+    @staticmethod
+    def writeModel(topology, positions, file=sys.stdout, modelIndex=1, keepIds=False):
+        """Write out a model to a PDBx/mmCIF file.
+
+        Parameters:
+         - topology (Topology) The Topology defining the model to write
+         - positions (list) The list of atomic positions to write
+         - file (file=stdout) A file to write the model to
+         - modelIndex (int=1) The model number of this frame
+         - keepIds (bool=False) If True, keep the residue and chain IDs specified in the Topology 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 PDBx/mmCIF format.  Otherwise, the output file will be invalid.
+        """
+        if len(list(topology.atoms())) != len(positions):
+            raise ValueError('The number of positions must match the number of atoms')
+        if is_quantity(positions):
+            positions = positions.value_in_unit(angstroms)
+        if any(math.isnan(norm(pos)) for pos in positions):
+            raise ValueError('Particle position is NaN')
+        if any(math.isinf(norm(pos)) for pos in positions):
+            raise ValueError('Particle position is infinite')
+        atomIndex = 1
+        posIndex = 0
+        for (chainIndex, chain) in enumerate(topology.chains()):
+            if keepIds:
+                chainName = chain.id
+            else:
+                chainName = chr(ord('A')+chainIndex%26)
+            residues = list(chain.residues())
+            for (resIndex, res) in enumerate(residues):
+                if keepIds:
+                    resId = res.id
+                else:
+                    resId = resIndex + 1
+                for atom in res.atoms():
+                    coords = positions[posIndex]
+                    if atom.element is not None:
+                        symbol = atom.element.symbol
+                    else:
+                        symbol = '?'
+                    line = "ATOM  %5d %-3s %-4s . %-4s %s ? %5s . %10.4f %10.4f %10.4f  0.0  0.0  ?  ?  ?  ?  ?  .  %5s %4s %s %4s %5d"
+                    print(line % (atomIndex, symbol, atom.name, res.name, chainName, resId, coords[0], coords[1], coords[2],
+                                  resId, res.name, chainName, atom.name, modelIndex), file=file)
+                    posIndex += 1
+                    atomIndex += 1

wrappers/python/tests/TestPdbxFile.py

@@ -3,6 +3,7 @@ from simtk.openmm.app import *
 from simtk.openmm import *
 from simtk.unit import *
 import simtk.openmm.app.element as elem
+import os
 
 class TestPdbxFile(unittest.TestCase):
     """Test the PDBx/mmCIF file parser"""
@@ -48,6 +49,69 @@ class TestPdbxFile(unittest.TestCase):
             diff = abs(p1[i]-p2[i])/scale
             self.assertTrue(diff < tol)
 
+    def testReporterImplicit(self):
+        """ Tests the PDBxReporter without PBC """
+        parm = AmberPrmtopFile('systems/alanine-dipeptide-implicit.prmtop')
+        system = parm.createSystem()
+        sim = Simulation(parm.topology, system, VerletIntegrator(1*femtoseconds),
+                         Platform.getPlatformByName('Reference'))
+        sim.context.setPositions(PDBFile('systems/alanine-dipeptide-implicit.pdb').getPositions())
+        sim.reporters.append(PDBxReporter('test.cif', 1))
+        sim.step(10)
+        pdb = PDBxFile('test.cif')
+        self.assertEqual(len(list(pdb.topology.atoms())), len(list(parm.topology.atoms())))
+        self.assertEqual(len(list(pdb.topology.residues())), len(list(parm.topology.residues())))
+        for res1, res2 in zip(pdb.topology.residues(), parm.topology.residues()):
+            self.assertEqual(res1.name, res2.name)
+            for atom1, atom2 in zip(res1.atoms(), res2.atoms()):
+                self.assertEqual(atom1.name, atom2.name)
+        positions = pdb.getPositions(frame=9)
+        self.assertFalse(all(x1 == x2 for x1, x2 in zip(positions, pdb.getPositions(frame=0))))
+        # There should only be 10 frames (0 through 9)
+        self.assertRaises(IndexError, lambda: pdb.getPositions(frame=10))
+        self.assertIs(pdb.topology.getPeriodicBoxVectors(), None)
+        os.unlink('test.cif')
+
+    def assertAlmostEqualVec(self, vec1, vec2, *args, **kwargs):
+        if is_quantity(vec1):
+            vec1 = vec1.value_in_unit_system(md_unit_system)
+        if is_quantity(vec2):
+            vec2 = vec2.value_in_unit_system(md_unit_system)
+        for x, y in zip(vec1, vec2):
+            self.assertAlmostEqual(x, y, *args, **kwargs)
+
+    def testReporterExplicit(self):
+        """ Tests the PDBxReporter with PBC """
+        parm = AmberPrmtopFile('systems/alanine-dipeptide-explicit.prmtop')
+        system = parm.createSystem(nonbondedCutoff=1.0, nonbondedMethod=PME)
+        sim = Simulation(parm.topology, system, VerletIntegrator(1*femtoseconds),
+                         Platform.getPlatformByName('Reference'))
+        orig_pdb = PDBFile('systems/alanine-dipeptide-explicit.pdb')
+        sim.context.setPositions(orig_pdb.getPositions())
+        sim.context.setPeriodicBoxVectors(*parm.topology.getPeriodicBoxVectors())
+        sim.reporters.append(PDBxReporter('test.cif', 1))
+        sim.step(10)
+        pdb = PDBxFile('test.cif')
+        self.assertEqual(len(list(pdb.topology.atoms())), len(list(parm.topology.atoms())))
+        self.assertEqual(len(list(pdb.topology.residues())), len(list(parm.topology.residues())))
+        for res1, res2 in zip(pdb.topology.residues(), parm.topology.residues()):
+            self.assertEqual(res1.name, res2.name)
+            for atom1, atom2 in zip(res1.atoms(), res2.atoms()):
+                self.assertEqual(atom1.name, atom2.name)
+        positions = pdb.getPositions(frame=9)
+        self.assertFalse(all(x1 == x2 for x1, x2 in zip(positions, pdb.getPositions(frame=0))))
+        # There should only be 10 frames (0 through 9)
+        self.assertRaises(IndexError, lambda: pdb.getPositions(frame=10))
+        self.assertAlmostEqualVec(parm.topology.getPeriodicBoxVectors()[0],
+                                  pdb.topology.getPeriodicBoxVectors()[0],
+                                  places=5)
+        self.assertAlmostEqualVec(parm.topology.getPeriodicBoxVectors()[1],
+                                  pdb.topology.getPeriodicBoxVectors()[1],
+                                  places=5)
+        self.assertAlmostEqualVec(parm.topology.getPeriodicBoxVectors()[2],
+                                  pdb.topology.getPeriodicBoxVectors()[2],
+                                  places=5)
+        os.unlink('test.cif')
 
 if __name__ == '__main__':
     unittest.main()