diff --git a/ctapipe/io/hdf5tableio.py b/ctapipe/io/hdf5tableio.py index b4e0f0f4ef2..5c84fa0d599 100644 --- a/ctapipe/io/hdf5tableio.py +++ b/ctapipe/io/hdf5tableio.py @@ -1,3 +1,5 @@ +'''Implementations of TableWriter and -Reader for HDF5 files''' +import enum from functools import partial import numpy as np @@ -150,6 +152,14 @@ class Schema(tables.IsDescription): ) continue + if isinstance(value, enum.Enum): + def transform(enum_value): + '''transform enum instance into its (integer) value''' + return enum_value.value + meta[f'{col_name}_ENUM'] = value.__class__ + value = transform(value) + self.add_column_transform(table_name, col_name, transform) + if isinstance(value, Quantity): if self.add_prefix and container.prefix: key = col_name.replace(container.prefix + '_', '') @@ -185,6 +195,7 @@ class Schema(tables.IsDescription): coltype = PYTABLES_TYPE_MAP[typename] Schema.columns[col_name] = coltype() + self.log.debug("Table %s: added col: %s type: %s shape: %s", table_name, col_name, typename, shape) @@ -332,6 +343,21 @@ def _map_transforms_from_table_header(self, table_name): tr = partial(tr_add_unit, unitname=tab.attrs[attr]) self.add_column_transform(table_name, colname, tr) + for attr in tab.attrs._f_list(): + if attr.endswith("_ENUM"): + colname = attr[:-5] + + def transform_int_to_enum(int_val): + '''transform integer 'code' into enum instance''' + enum_class = tab.attrs[attr] + return enum_class(int_val) + + self.add_column_transform( + table_name, + colname, + transform_int_to_enum + ) + def _map_table_to_container(self, table_name, container): """ identifies which columns in the table to read into the container, by comparing their names.""" diff --git a/ctapipe/io/tests/test_hdf5.py b/ctapipe/io/tests/test_hdf5.py index 34b06a3dcb6..33a2b1128aa 100644 --- a/ctapipe/io/tests/test_hdf5.py +++ b/ctapipe/io/tests/test_hdf5.py @@ -1,5 +1,6 @@ import tempfile +import enum import numpy as np import pytest import tables @@ -285,6 +286,76 @@ class ContainerA(Container): assert a.a == 1 +class WithNormalEnum(Container): + class EventType(enum.Enum): + pedestal = 1 + physics = 2 + calibration = 3 + + event_type = Field( + EventType.calibration, + f'type of event, one of: {list(EventType.__members__.keys())}' + ) + + +def test_read_write_container_with_enum(tmp_path): + tmp_file = tmp_path / 'container_with_enum.hdf5' + + def create_stream(n_event): + data = WithNormalEnum() + for i in range(n_event): + data.event_type = data.EventType(i % 3 + 1) + yield data + + with HDF5TableWriter(tmp_file, group_name='data') as h5_table: + for data in create_stream(10): + h5_table.write('table', data) + + with HDF5TableReader(tmp_file, mode='r') as h5_table: + for group_name in ['data/']: + group_name = '/{}table'.format(group_name) + for data in h5_table.read(group_name, WithNormalEnum()): + assert isinstance( + data.event_type, + WithNormalEnum.EventType + ) + + +class WithIntEnum(Container): + class EventType(enum.IntEnum): + pedestal = 1 + physics = 2 + calibration = 3 + + event_type = Field( + EventType.calibration, + f'type of event, one of: {list(EventType.__members__.keys())}' + ) + + +def test_read_write_container_with_int_enum(tmp_path): + tmp_file = tmp_path / 'container_with_int_enum.hdf5' + + def create_stream(n_event): + data = WithIntEnum() + for i in range(n_event): + data.event_type = data.EventType(i % 3 + 1) + yield data + + with HDF5TableWriter(tmp_file, group_name='data') as h5_table: + for data in create_stream(10): + h5_table.write('table', data) + + with HDF5TableReader(tmp_file, mode='r') as h5_table: + for group_name in ['data/']: + group_name = '/{}table'.format(group_name) + for data in h5_table.read(group_name, WithIntEnum()): + assert isinstance( + data.event_type, + WithIntEnum.EventType + ) + + if __name__ == '__main__': import logging diff --git a/docs/examples/containers_with_enums_and_table_writer.ipynb b/docs/examples/containers_with_enums_and_table_writer.ipynb new file mode 100644 index 00000000000..5bd23e96eeb --- /dev/null +++ b/docs/examples/containers_with_enums_and_table_writer.ipynb @@ -0,0 +1,212 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# containers_with_enums_and_table_writer\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Create some example Containers" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import enum\n", + "from ctapipe.io import HDF5TableWriter\n", + "from ctapipe.core import Container, Field\n", + "from astropy import units as u\n", + "import numpy as np" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "class WithEnum(Container):\n", + " \n", + " # this class could also be defined in global namespace \n", + " # outside this container, but this looks a bit tidier.\n", + " # both variants work however\n", + " class EventType(enum.Enum):\n", + " pedestal = 1\n", + " physics = 2\n", + " calibration = 3\n", + " \n", + " event_type = Field(\n", + " EventType.calibration, \n", + " f'type of event, one of: {list(EventType.__members__.keys())}'\n", + " )" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "let's also make a dummy stream (generator) that will create a series of these containers" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def create_stream(n_event):\n", + " data = WithEnum()\n", + " for i in range(n_event):\n", + " data.event_type = WithEnum.EventType(i % 3 + 1)\n", + " yield data" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "for data in create_stream(3):\n", + " for key, val in data.items():\n", + " print('{}: {}, type : {}'.format(key, val, type(val)))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Writing the Data" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "with HDF5TableWriter('container.h5', group_name='data') as h5_table:\n", + " for data in create_stream(10):\n", + " h5_table.write('table', data)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "!ls container.h5" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Reading the Data" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import pandas as pd\n", + "\n", + "data = pd.read_hdf('container.h5', key='/data/table')\n", + "data.head()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Reading with PyTables" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import tables\n", + "h5 = tables.open_file('container.h5')\n", + "table = h5.root['data']['table']\n", + "table" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "table.attrs" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from ctapipe.io import HDF5TableReader\n", + "\n", + "def read(mode):\n", + " \n", + " print('reading mode {}'.format(mode))\n", + "\n", + " with HDF5TableReader('container.h5', mode=mode) as h5_table:\n", + "\n", + " for group_name in ['data/']:\n", + "\n", + " group_name = '/{}table'.format(group_name)\n", + " print(group_name)\n", + "\n", + " for data in h5_table.read(group_name, WithEnum()):\n", + "\n", + " print(data.as_dict())" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "read('r')" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.7.2" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/docs/examples/index.rst b/docs/examples/index.rst index 92ef4ad655f..fdae4095d3b 100644 --- a/docs/examples/index.rst +++ b/docs/examples/index.rst @@ -10,7 +10,7 @@ the Tutorials section for more complete examples) .. toctree:: :maxdepth: 1 :caption: Algorithms - + dilate_image nd_interpolation convert_hex_to_square @@ -18,10 +18,11 @@ the Tutorials section for more complete examples) .. toctree:: :maxdepth: 1 :caption: Core functionality - + InstrumentDescription camera_display containers Tools provenance table_writer_reader + containers_with_enums_and_table_writer