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"),
@@ -79,9 +79,12 @@ 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.
@@ -93,7 +96,7 @@ 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.
@@ -103,6 +106,58 @@ class StateDataReporter(object):
          - state (State) The current state of the simulation
         """
         if not self._hasInitialized:
+            self._initializeConstants(simulation)
+            headers = self._constructHeaders()
+            print >>self._out, '#"%s"' % ('"'+self._separator+'"').join(headers)
+            self._hasInitialized = True
+
+        # Check for errors.
+        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]
+        if self._step:
+            values.append(simulation.currentStep)
+        if self._time:
+            values.append(state.getTime().value_in_unit(unit.picosecond))
+        if self._potentialEnergy:
+            values.append(state.getPotentialEnergy().value_in_unit(unit.kilojoules_per_mole))
+        if self._kineticEnergy:
+            values.append(state.getKineticEnergy().value_in_unit(unit.kilojoules_per_mole))
+        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*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))
+        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.
@@ -123,8 +178,11 @@ class StateDataReporter(object):
             elif not unit.is_quantity(self._totalMass):
                 self._totalMass = self._totalMass*unit.dalton
 
-            # Write the headers.
+    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')
@@ -142,11 +200,15 @@ class StateDataReporter(object):
             headers.append('Box Volume (nm^3)')
         if self._density:
             headers.append('Density (g/mL)')
-            print >>self._out, '#"%s"' % ('"'+self._separator+'"').join(headers)
-            self._hasInitialized = True
+        return headers
 
-        # Check for errors.
+    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):
@@ -154,29 +216,6 @@ class StateDataReporter(object):
             if math.isinf(energy):
                 raise ValueError('Energy is infinite')
 
-        # Write the values.
-        
-        values = []
-        box = state.getPeriodicBoxVectors()
-        volume = box[0][0]*box[1][1]*box[2][2]
-        if self._step:
-            values.append(simulation.currentStep)
-        if self._time:
-            values.append(state.getTime().value_in_unit(unit.picosecond))
-        if self._potentialEnergy:
-            values.append(state.getPotentialEnergy().value_in_unit(unit.kilojoules_per_mole))
-        if self._kineticEnergy:
-            values.append(state.getKineticEnergy().value_in_unit(unit.kilojoules_per_mole))
-        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))
-        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)
-        
     def __del__(self):
         if self._openedFile:
             self._out.close()