From fcd29492f92405c1c184267d25a469e93582adc4 Mon Sep 17 00:00:00 2001 From: Sebastien Bourdeauducq Date: Sat, 24 Oct 2015 10:54:59 +0800 Subject: [PATCH] worker_db: support more types in HDF5 output. Closes #144. Closes #121 --- .gitignore | 1 + artiq/master/worker_db.py | 34 ++++++++++++++++++++++++++-------- artiq/test/h5types.py | 33 +++++++++++++++++++++++++++++++++ 3 files changed, 60 insertions(+), 8 deletions(-) create mode 100644 artiq/test/h5types.py diff --git a/.gitignore b/.gitignore index 50894b653..2cc1c622f 100644 --- a/.gitignore +++ b/.gitignore @@ -12,6 +12,7 @@ doc/manual/_build /*.egg-info /.coverage artiq/test/results +artiq/test/h5types.h5 examples/master/results examples/master/dataset_db.pyon examples/sim/dataset_db.pyon diff --git a/artiq/master/worker_db.py b/artiq/master/worker_db.py index 950ec147f..17dee5a1b 100644 --- a/artiq/master/worker_db.py +++ b/artiq/master/worker_db.py @@ -5,7 +5,7 @@ import os import time import re -import numpy +import numpy as np import h5py from artiq.protocols.sync_struct import Notifier @@ -119,7 +119,21 @@ def get_last_rid(): _type_to_hdf5 = { int: h5py.h5t.STD_I64BE, - float: h5py.h5t.IEEE_F64BE + float: h5py.h5t.IEEE_F64BE, + + np.int8: h5py.h5t.STD_I8BE, + np.int16: h5py.h5t.STD_I16BE, + np.int32: h5py.h5t.STD_I32BE, + np.int64: h5py.h5t.STD_I64BE, + + np.uint8: h5py.h5t.STD_U8BE, + np.uint16: h5py.h5t.STD_U16BE, + np.uint32: h5py.h5t.STD_U32BE, + np.uint64: h5py.h5t.STD_U64BE, + + np.float16: h5py.h5t.IEEE_F16BE, + np.float32: h5py.h5t.IEEE_F32BE, + np.float64: h5py.h5t.IEEE_F64BE } def result_dict_to_hdf5(f, rd): @@ -137,15 +151,19 @@ def result_dict_to_hdf5(f, rd): " HDF5 output".format(el_ty)) dataset = f.create_dataset(name, (len(data), ), el_ty_h5) dataset[:] = data - elif isinstance(data, numpy.ndarray): + elif isinstance(data, np.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)) + if ty is str: + ty_h5 = "S{}".format(len(data)) + data = data.encode() + else: + 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 diff --git a/artiq/test/h5types.py b/artiq/test/h5types.py new file mode 100644 index 000000000..dd204aa66 --- /dev/null +++ b/artiq/test/h5types.py @@ -0,0 +1,33 @@ +import unittest + +import h5py +import numpy as np + +from artiq.master.worker_db import result_dict_to_hdf5 + + +class TypesCase(unittest.TestCase): + def test_types(self): + d = { + "int": 42, + "float": 42.0, + "string": "abcdef", + + "intlist": [1, 2, 3], + "floatlist": [1.0, 2.0, 3.0] + } + + for size in 8, 16, 32, 64: + signed = getattr(np, "int" + str(size)) + unsigned = getattr(np, "uint" + str(size)) + d["i"+str(size)] = signed(42) + d["u"+str(size)] = unsigned(42) + d["i{}list".format(size)] = [signed(x) for x in range(3)] + d["u{}list".format(size)] = [unsigned(x) for x in range(3)] + for size in 16, 32, 64: + ty = getattr(np, "float" + str(size)) + d["f"+str(size)] = ty(42) + d["f{}list".format(size)] = [ty(x) for x in range(3)] + + with h5py.File("h5types.h5", "w") as f: + result_dict_to_hdf5(f, d)