From af79392918097e8c42bb16474a494e0560dc7c0e Mon Sep 17 00:00:00 2001
From: Sebastien Bourdeauducq <sb@m-labs.hk>
Date: Wed, 4 Feb 2015 18:37:57 +0800
Subject: [PATCH] master: basic dumping of results to HDF5

---
 artiq/master/db.py          | 40 +++++++++++++++++++++++++++++++++++++
 artiq/master/worker_impl.py |  9 +++++++++
 setup.py                    |  2 +-
 3 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/artiq/master/db.py b/artiq/master/db.py
index f6283edf4..e8d4c9c7c 100644
--- a/artiq/master/db.py
+++ b/artiq/master/db.py
@@ -1,9 +1,45 @@
 from collections import OrderedDict
 import importlib
 
+import numpy
+import h5py
+
 from artiq.protocols.sync_struct import Notifier
 
 
+_type_to_hdf5 = {
+    int: h5py.h5t.STD_I64BE,
+    float: h5py.h5t.IEEE_F64BE
+}
+
+def _result_dict_to_hdf5(f, rd):
+    for name, data in rd.items():
+        if isinstance(data, list):
+            el_ty = type(data[0])
+            for d in data:
+                if type(d) != el_ty:
+                    raise TypeError("All list elements must have the same"
+                                    " type for HDF5 output")
+            try:
+                el_ty_h5 = _type_to_hdf5[el_ty]
+            except KeyError:
+                raise TypeError("List element type {} is not supported for"
+                                " HDF5 output".format(el_ty))
+            dataset = f.create_dataset(name, (len(data), ), el_ty_h5)
+            dataset[:] = data
+        elif isinstance(data, numpy.ndarray):
+            f.create_dataset(name, data=data)
+        else:
+            ty = type(data)
+            try:
+                ty_h5 = _type_to_hdf5[ty]
+            except KeyError:
+                raise TypeError("Type {} is not supported for HDF5 output"
+                                .format(ty))
+            dataset = f.create_dataset(name, (), ty_h5)
+            dataset[()] = data
+
+
 class ResultDB:
     def __init__(self, realtime_results):
         self.realtime_data = Notifier({x: [] for x in realtime_results})
@@ -30,6 +66,10 @@ class ResultDB:
         else:
             self.data[name] = value
 
+    def write_hdf5(self, f):
+        _result_dict_to_hdf5(f, self.realtime_data.read)
+        _result_dict_to_hdf5(f, self.data.read)
+
 
 def _create_device(desc, dbh):
     module = importlib.import_module(desc["module"])
diff --git a/artiq/master/worker_impl.py b/artiq/master/worker_impl.py
index a7cfe8aa3..8fda84c07 100644
--- a/artiq/master/worker_impl.py
+++ b/artiq/master/worker_impl.py
@@ -2,6 +2,8 @@ import sys
 from inspect import isclass
 import traceback
 
+import h5py
+
 from artiq.protocols import pyon
 from artiq.tools import file_import
 from artiq.language.db import AutoDB
@@ -102,6 +104,13 @@ def run(obj):
     finally:
         dbh.close()
 
+    filename = obj["file"] + ".h5"
+    f = h5py.File(filename, "w")
+    try:
+        rdb.write_hdf5(f)
+    finally:
+        f.close()
+
 
 def main():
     sys.stdout = sys.stderr
diff --git a/setup.py b/setup.py
index 2e70b4ef0..4838700d6 100755
--- a/setup.py
+++ b/setup.py
@@ -15,7 +15,7 @@ setup(
     license="BSD",
     install_requires=[
         "sphinx", "sphinx-argparse", "pyserial", "numpy", "scipy",
-        "python-dateutil", "prettytable"
+        "python-dateutil", "prettytable", "h5py"
     ],
     extras_require={},
     dependency_links=[],