master: basic dumping of results to HDF5

This commit is contained in:
Sebastien Bourdeauducq 2015-02-04 18:37:57 +08:00
parent 2f4a83b97a
commit af79392918
3 changed files with 50 additions and 1 deletions

View File

@ -1,9 +1,45 @@
from collections import OrderedDict from collections import OrderedDict
import importlib import importlib
import numpy
import h5py
from artiq.protocols.sync_struct import Notifier 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: class ResultDB:
def __init__(self, realtime_results): def __init__(self, realtime_results):
self.realtime_data = Notifier({x: [] for x in realtime_results}) self.realtime_data = Notifier({x: [] for x in realtime_results})
@ -30,6 +66,10 @@ class ResultDB:
else: else:
self.data[name] = value 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): def _create_device(desc, dbh):
module = importlib.import_module(desc["module"]) module = importlib.import_module(desc["module"])

View File

@ -2,6 +2,8 @@ import sys
from inspect import isclass from inspect import isclass
import traceback import traceback
import h5py
from artiq.protocols import pyon from artiq.protocols import pyon
from artiq.tools import file_import from artiq.tools import file_import
from artiq.language.db import AutoDB from artiq.language.db import AutoDB
@ -102,6 +104,13 @@ def run(obj):
finally: finally:
dbh.close() dbh.close()
filename = obj["file"] + ".h5"
f = h5py.File(filename, "w")
try:
rdb.write_hdf5(f)
finally:
f.close()
def main(): def main():
sys.stdout = sys.stderr sys.stdout = sys.stderr

View File

@ -15,7 +15,7 @@ setup(
license="BSD", license="BSD",
install_requires=[ install_requires=[
"sphinx", "sphinx-argparse", "pyserial", "numpy", "scipy", "sphinx", "sphinx-argparse", "pyserial", "numpy", "scipy",
"python-dateutil", "prettytable" "python-dateutil", "prettytable", "h5py"
], ],
extras_require={}, extras_require={},
dependency_links=[], dependency_links=[],