Add Histogram.to_stream() and .headings.

diff --git a/pysawsim/histogram.py b/pysawsim/histogram.py
index 880fd409ec4f8c1b487cf0b8eba2487f82cc6b8c..32a12cf2f2171ab8663d91b9b846f95d6d84b878 100644 (file)
--- a/pysawsim/histogram.py
+++ b/pysawsim/histogram.py
@@ -55,6 +55,7 @@ class Histogram (object):
 
         `data` should be a numpy array.
         """
+        self.headings = None
         self.bin_edges = bin_edges
         bin_width = self.bin_edges[1] - self.bin_edges[0]
 
@@ -85,11 +86,13 @@ class Histogram (object):
 
         >>> import StringIO
         >>> h = Histogram()
-        >>> h.from_stream(StringIO.StringIO('''# Force (N)\tUnfolding events
-        ... 150e-12\t10
-        ... 200e-12\t40
-        ... 250e-12\t5
+        >>> h.from_stream(StringIO.StringIO('''# Force (N)\\tUnfolding events
+        ... 150e-12\\t10
+        ... 200e-12\\t40
+        ... 250e-12\\t5
         ... '''))
+        >>> h.headings
+        ['Force (N)', 'Unfolding events']
         >>> h.total
         55.0
         >>> h.counts
@@ -99,11 +102,15 @@ class Histogram (object):
         >>> h.probabilities  # doctest: +ELLIPSIS
         [0.181..., 0.727..., 0.0909...]
         """
+        self.headings = None
         self.bin_edges = []
         self.counts = []
         for line in stream.readlines():
             line = line.strip()
-            if len(line) == 0 or line[0] == "#":
+            if len(line) == 0 or line.startswith('#'):
+                if self.headings == None and line.startswith('#'):
+                    line = line[len('#'):]
+                    self.headings = [x.strip() for x in line.split('\t')]
                 continue # ignore blank lines and comments
             bin_edge,count = line.split()
             self.bin_edges.append(float(bin_edge))
@@ -123,6 +130,26 @@ class Histogram (object):
         self.std_dev = numpy.sqrt(variance)
         self.normalize()
 
+    def to_stream(self, stream):
+        """Write to `stream` with the same format as `from_stream()`.
+
+        >>> import sys
+        >>> h = Histogram()
+        >>> h.headings = ['Force (N)', 'Unfolding events']
+        >>> h.bin_edges = [1.5e-10, 2e-10, 2.5e-10, 3e-10]
+        >>> h.counts = [10, 40, 5]
+        >>> h.to_stream(sys.stdout)
+        ... # doctest: +NORMALIZE_WHITESPACE, +REPORT_UDIFF
+        # Force (N)\tUnfolding events
+        1.5e-10\t10
+        2e-10\t40
+        2.5e-10\t5
+        """
+        if self.headings != None:
+            stream.write('# %s\n' % '\t'.join(self.headings))
+        for bin,count in zip(self.bin_edges, self.counts):
+            stream.write('%g\t%g\n' % (bin, count))
+
     def normalize(self):
         self.total = float(self.total)
         self.probabilities = [count/self.total for count in self.counts]