Commit 28020b30 authored by Jason Swails's avatar Jason Swails
Browse files

Support gzip and bzip2 compression for StateDataReporter (based on filename

extension). Users were already able to pass gzip.GzipFile objects and
StateDataReporter would write compressed data. bz2.BZ2File objects also worked
after I removed the self._out.flush() calls.
parent 8eb840b6
......@@ -31,9 +31,10 @@ USE OR OTHER DEALINGS IN THE SOFTWARE.
__author__ = "Peter Eastman"
__version__ = "1.0"
import bz2
import gzip
import simtk.openmm as mm
import simtk.unit as unit
from simtk.openmm.app import PDBFile
import math
class StateDataReporter(object):
......@@ -67,6 +68,13 @@ class StateDataReporter(object):
self._reportInterval = reportInterval
self._openedFile = isinstance(file, str)
if self._openedFile:
# Detect the desired compression scheme from the filename extension
# and open all files unbuffered
if file.endswith('.gz'):
self._out = gzip.GzipFile(fileobj=open(file, 'wb', 0))
elif file.endswith('.bz2'):
self._out = bz2.BZ2File(file, 'w', 0)
else:
self._out = open(file, 'w', 0)
else:
self._out = file
......@@ -109,7 +117,6 @@ class StateDataReporter(object):
self._initializeConstants(simulation)
headers = self._constructHeaders()
print >>self._out, '#"%s"' % ('"'+self._separator+'"').join(headers)
self._out.flush()
self._hasInitialized = True
# Check for errors.
......@@ -120,7 +127,6 @@ class StateDataReporter(object):
# Write the values.
print >>self._out, self._separator.join(str(v) for v in values)
self._out.flush()
def _constructReportValues(self, simulation, state):
"""Query the simulation for the current state of our observables of interest.
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment