PdbFile and AmberInpcrdFile can return values as Numpy arrays

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

@@ -5,6 +5,11 @@ __author__ = "Peter Eastman"
 __version__ = "1.0"
 
 from simtk.openmm.app.internal import amber_file_parser
+from simtk.unit import nanometers, picoseconds
+try:
+    import numpy
+except:
+    pass
 
 class AmberInpcrdFile(object):
     """AmberInpcrdFile parses an AMBER inpcrd file and loads the data stored in it."""
@@ -35,3 +40,48 @@ class AmberInpcrdFile(object):
             self.boxVectors = results[1]
         else:
             self.positions = results
+        self._numpyPositions = None
+        if loadVelocities:
+            self._numpyVelocities = None
+        if loadBoxVectors:
+            self._numpyBoxVectors = None
+
+    def getPositions(self, asNumpy=False):
+        """Get the atomic positions.
+        
+        Parameters:
+         - asNumpy (boolean=False) if true, the values are returned as a numpy array instead of a list of Vec3s
+         """
+        if asNumpy:
+            if self._numpyPositions is None:
+                self._numpyPositions = numpy.array(self.positions.value_in_unit(nanometers))*nanometers
+            return self._numpyPositions
+        return self.positions
+    
+    def getVelocities(self, asNumpy=False):
+        """Get the atomic velocities.
+        
+        Parameters:
+         - asNumpy (boolean=False) if true, the vectors are returned as numpy arrays instead of Vec3s
+         """
+        if asNumpy:
+            if self._numpyVelocities is None:
+                self._numpyVelocities = numpy.array(self.velocities.value_in_unit(nanometers/picoseconds))*nanometers/picoseconds
+            return self._numpyVelocities
+        return self.velocities
+    
+    def getBoxVectors(self, asNumpy=False):
+        """Get the periodic box vectors.
+        
+        Parameters:
+         - asNumpy (boolean=False) if true, the values are returned as a numpy array instead of a list of Vec3s
+         """
+        if asNumpy:
+            if self._numpyBoxVectors is None:
+                self._numpyBoxVectors = []
+                self._numpyBoxVectors.append(numpy.array(self.boxVectors[0].value_in_unit(nanometers))*nanometers)
+                self._numpyBoxVectors.append(numpy.array(self.boxVectors[1].value_in_unit(nanometers))*nanometers)
+                self._numpyBoxVectors.append(numpy.array(self.boxVectors[2].value_in_unit(nanometers))*nanometers)
+            return self._numpyBoxVectors
+        return self.boxVectors
+        
\ No newline at end of file

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

@@ -780,7 +780,7 @@ def readAmberCoordinates(filename, read_box=False, read_velocities=False, verbos
             line = infile.readline().strip()
             elements = line.split()
             while (len(elements) > 0):
-                velocities.append(mm.Vec3(float(elements.pop(0)), float(elements.pop(0)), float(elements.pop(0))))
+                velocities.append(20.455*mm.Vec3(float(elements.pop(0)), float(elements.pop(0)), float(elements.pop(0))))
             natoms_read += 1
         if asNumpy:
             newvel = numpy.zeros([natoms,3], numpy.float32)
@@ -789,7 +789,7 @@ def readAmberCoordinates(filename, read_box=False, read_velocities=False, verbos
                     newvel[i,j] = velocities[i][j]
             velocities = newvel
         # Assign units.
-        velocities = units.Quantity(velocities, units.angstroms)
+        velocities = units.Quantity(velocities, units.angstroms/units.picoseconds)
             
     # Read box size if present
     box_vectors = None

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

@@ -13,7 +13,11 @@ from simtk.openmm.app.internal.pdbstructure import PdbStructure
 from simtk.openmm.app import Topology
 from simtk.unit import nanometers, angstroms, is_quantity
 import element as elem
-    
+try:
+    import numpy
+except:
+    pass
+
 class PDBFile(object):
     """PDBFile parses a Protein Data Bank (PDB) file and constructs a Topology and a set of atom positions from it.
     
@@ -27,7 +31,7 @@ class PDBFile(object):
     def __init__(self, file):
         """Load a PDB file.
         
-        The atom positions and Topology are stored into this object's "positions" and "topology" fields.
+        The atom positions and Topology can be retrieved by calling getPositions() and getTopology().
         
         Parameters:
          - file (string) the name of the file to load
@@ -90,6 +94,23 @@ class PDBFile(object):
         self.topology.setUnitCellDimensions(pdb.get_unit_cell_dimensions())
         self.topology.createStandardBonds()
         self.topology.createDisulfideBonds(self.positions)
+        self._numpyPositions = None
+
+    def getTopology(self):
+        """Get the Topology of the model."""
+        return self.topology
+        
+    def getPositions(self, asNumpy=False):
+        """Get the atomic positions.
+        
+        Parameters:
+         - asNumpy (boolean=False) if true, the values are returned as a numpy array instead of a list of Vec3s
+         """
+        if asNumpy:
+            if self._numpyPositions is None:
+                self._numpyPositions = numpy.array(self.positions.value_in_unit(nanometers))*nanometers
+            return self._numpyPositions
+        return self.positions
 
     @staticmethod
     def _loadNameReplacementTables():