master: basic dumping of results to HDF5

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"])

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

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=[],