Merge pull request #978 from cta-observatory/allow_enums_in_table_wri…

…ter_and_reader

* allow enums in containers and support in tableio

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."""

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

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
+}

docs/examples/index.rst

@@ -10,18 +10,19 @@ the Tutorials section for more complete examples)
 .. toctree::
     :maxdepth: 1
     :caption: Algorithms
-    
+
     dilate_image
     nd_interpolation
     convert_hex_to_square
 
 .. toctree::
     :maxdepth: 1
     :caption: Core functionality
-    
+
     InstrumentDescription
     camera_display
     containers
     Tools
     provenance
     table_writer_reader
+    containers_with_enums_and_table_writer