Robert's refactoring of StateDataReporter to let subclasses add new reported values

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

@@ -6,9 +6,9 @@ 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) 2012 Stanford University and the Authors.
+Portions copyright (c) 2012-2013 Stanford University and the Authors.
 Authors: Peter Eastman
-Contributors:
+Contributors: Robert McGibbon
 
 Permission is hereby granted, free of charge, to any person obtaining a 
 copy of this software and associated documentation files (the "Software"),
@@ -38,15 +38,15 @@ import math
     
 class StateDataReporter(object):
     """StateDataReporter outputs information about a simulation, such as energy and temperature, to a file.
-    
+
     To use it, create a StateDataReporter, then add it to the Simulation's list of reporters.  The set of
     data to write is configurable using boolean flags passed to the constructor.  By default the data is
     written in comma-separated-value (CSV) format, but you can specify a different separator to use.
     """
-    
+
     def __init__(self, file, reportInterval, step=False, time=False, potentialEnergy=False, kineticEnergy=False, totalEnergy=False, temperature=False, volume=False, density=False, separator=',', systemMass=None):
         """Create a StateDataReporter.
-    
+
         Parameters:
          - file (string or file) The file to write to, specified as a file name or file object
          - reportInterval (int) The interval (in time steps) at which to write frames
@@ -79,13 +79,16 @@ class StateDataReporter(object):
         self._volume = volume
         self._density = density
         self._separator = separator
-        self._needEnergy = potentialEnergy or kineticEnergy or totalEnergy or temperature
         self._totalMass = systemMass
         self._hasInitialized = False
-    
+        self._needsPositions = False
+        self._needsVelocities = False
+        self._needsForces = False
+        self._needEnergy = potentialEnergy or kineticEnergy or totalEnergy or temperature
+
     def describeNextReport(self, simulation):
         """Get information about the next report this object will generate.
-        
+
         Parameters:
          - simulation (Simulation) The Simulation to generate a report for
         Returns: A five element tuple.  The first element is the number of steps until the
@@ -93,69 +96,41 @@ class StateDataReporter(object):
         positions, velocities, forces, and energies respectively.
         """
         steps = self._reportInterval - simulation.currentStep%self._reportInterval
-        return (steps, False, False, False, self._needEnergy)
-    
+        return (steps, self._needsPositions, self._needsVelocities, self._needsForces, self._needEnergy)
+
     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 not self._hasInitialized:
-            system = simulation.system
-            if self._temperature:
-                # Compute the number of degrees of freedom.
-                dof = 0
-                for i in range(system.getNumParticles()):
-                    if system.getParticleMass(i) > 0*unit.dalton:
-                        dof += 3
-                dof -= system.getNumConstraints()
-                if any(type(system.getForce(i)) == mm.CMMotionRemover for i in range(system.getNumForces())):
-                    dof -= 3
-                self._dof = dof
-            if self._density:
-                if self._totalMass is None:
-                    # Compute the total system mass.
-                    self._totalMass = 0*unit.dalton
-                    for i in range(system.getNumParticles()):
-                        self._totalMass += system.getParticleMass(i)
-                elif not unit.is_quantity(self._totalMass):
-                    self._totalMass = self._totalMass*unit.dalton
-            
-            # Write the headers.
-            
-            headers = []
-            if self._step:
-                headers.append('Step')
-            if self._time:
-                headers.append('Time (ps)')
-            if self._potentialEnergy:
-                headers.append('Potential Energy (kJ/mole)')
-            if self._kineticEnergy:
-                headers.append('Kinetic Energy (kJ/mole)')
-            if self._totalEnergy:
-                headers.append('Total Energy (kJ/mole)')
-            if self._temperature:
-                headers.append('Temperature (K)')
-            if self._volume:
-                headers.append('Box Volume (nm^3)')
-            if self._density:
-                headers.append('Density (g/mL)')
+            self._initializeConstants(simulation)
+            headers = self._constructHeaders()
             print >>self._out, '#"%s"' % ('"'+self._separator+'"').join(headers)
             self._hasInitialized = True
 
         # Check for errors.
-        
-        if self._needEnergy:
-            energy = (state.getKineticEnergy()+state.getPotentialEnergy()).value_in_unit(unit.kilojoules_per_mole)
-            if math.isnan(energy):
-                raise ValueError('Energy is NaN')
-            if math.isinf(energy):
-                raise ValueError('Energy is infinite')
-                
+        self._checkForErrors(simulation, state)
+
+        # Query for the values
+        values = self._constructReportValues(simulation, state)
+
         # Write the values.
-        
+        print >>self._out, self._separator.join(str(v) for v in values)
+
+    def _constructReportValues(self, simulation, state):
+        """Query the simulation for the current state of our observables of interest.
+
+        Parameters:
+         - simulation (Simulation) The Simulation to generate a report for
+         - state (State) The current state of the simulation
+
+        Returns: A list of values summarizing the current state of
+        the simulation, to be printed or saved. Each element in the list
+        corresponds to one of the columns in the resulting CSV file.
+        """
         values = []
         box = state.getPeriodicBoxVectors()
         volume = box[0][0]*box[1][1]*box[2][2]
@@ -170,13 +145,77 @@ class StateDataReporter(object):
         if self._totalEnergy:
             values.append((state.getKineticEnergy()+state.getPotentialEnergy()).value_in_unit(unit.kilojoules_per_mole))
         if self._temperature:
-            values.append((2*state.getKineticEnergy()/(self._dof*0.00831451)).value_in_unit(unit.kilojoules_per_mole))
+            values.append((2*state.getKineticEnergy()/(self._dof*unit.MOLAR_GAS_CONSTANT_R)).value_in_unit(unit.kelvin))
         if self._volume:
             values.append(volume.value_in_unit(unit.nanometer**3))
         if self._density:
             values.append((self._totalMass/volume).value_in_unit(unit.gram/unit.item/unit.milliliter))
-        print >>self._out, self._separator.join(str(v) for v in values)
-        
+        return values
+
+    def _initializeConstants(self, simulation):
+        """Initialize a set of constants required for the reports
+
+        Parameters
+        - simulation (Simulation) The simulation to generate a report for
+        """
+        system = simulation.system
+        if self._temperature:
+            # Compute the number of degrees of freedom.
+            dof = 0
+            for i in range(system.getNumParticles()):
+                if system.getParticleMass(i) > 0*unit.dalton:
+                    dof += 3
+            dof -= system.getNumConstraints()
+            if any(type(system.getForce(i)) == mm.CMMotionRemover for i in range(system.getNumForces())):
+                dof -= 3
+            self._dof = dof
+        if self._density:
+            if self._totalMass is None:
+                # Compute the total system mass.
+                self._totalMass = 0*unit.dalton
+                for i in range(system.getNumParticles()):
+                    self._totalMass += system.getParticleMass(i)
+            elif not unit.is_quantity(self._totalMass):
+                self._totalMass = self._totalMass*unit.dalton
+
+    def _constructHeaders(self):
+        """Construct the headers for the CSV output
+
+        Returns: a list of strings giving the title of each observable being reported on.
+        """
+        headers = []
+        if self._step:
+            headers.append('Step')
+        if self._time:
+            headers.append('Time (ps)')
+        if self._potentialEnergy:
+            headers.append('Potential Energy (kJ/mole)')
+        if self._kineticEnergy:
+            headers.append('Kinetic Energy (kJ/mole)')
+        if self._totalEnergy:
+            headers.append('Total Energy (kJ/mole)')
+        if self._temperature:
+            headers.append('Temperature (K)')
+        if self._volume:
+            headers.append('Box Volume (nm^3)')
+        if self._density:
+            headers.append('Density (g/mL)')
+        return headers
+
+    def _checkForErrors(self, simulation, state):
+        """Check for errors in the current state of the simulation
+
+        Parameters
+         - simulation (Simulation) The Simulation to generate a report for
+         - state (State) The current state of the simulation
+        """
+        if self._needEnergy:
+            energy = (state.getKineticEnergy()+state.getPotentialEnergy()).value_in_unit(unit.kilojoules_per_mole)
+            if math.isnan(energy):
+                raise ValueError('Energy is NaN')
+            if math.isinf(energy):
+                raise ValueError('Energy is infinite')
+
     def __del__(self):
         if self._openedFile:
             self._out.close()