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