make list of hist.axis object configurable

Author: TjarkMiener

examples/tutorials/histograms_aggregation.py

@@ -8,11 +8,11 @@
 2. Configure and run HistogramsAggregator in chunks.
 3. Access counts, bin edges, and valid-event counts (n_events).
 4. Plot one pixel histogram from the selected chunks and both gain channels for both image and peak_time columns.
+5. Show how the plotted histograms follow the configured ``hist.axis`` definitions.
 """
 
 import matplotlib.pyplot as plt
 import numpy as np
-import hist
 from astropy.table import Table
 from astropy.time import Time
 from traitlets.config import Config
@@ -65,17 +65,31 @@
     {
         "HistogramsAggregator": {
             "chunking_type": "SizeChunking",
+            "hist_axis_list": [
+                {
+                    "axis_class_name": "Regular",
+                    "kwargs": {
+                        "bins": 50,
+                        "start": 40.0,
+                        "stop": 110.0,
+                        "name": "value",
+                    },
+                },
+                {
+                    "axis_class_name": "Regular",
+                    "kwargs": {
+                        "bins": 50,
+                        "start": 20.0,
+                        "stop": 90.0,
+                        "name": "value",
+                    },
+                },
+            ],
         },
         "SizeChunking": {"chunk_size": 1000},
     }
 )
-
-image_hist_templates = [
-    hist.Hist(hist.axis.Regular(50, 40.0, 110.0, name="value")),
-    hist.Hist(hist.axis.Regular(50, 20.0, 90.0, name="value")),
-]
-
-aggregator_image = HistogramsAggregator(image_hist_templates, config=config_image)
+aggregator_image = HistogramsAggregator(config=config_image)
 result = aggregator_image(
     table=table,
     col_name="image",
@@ -86,15 +100,23 @@
     {
         "HistogramsAggregator": {
             "chunking_type": "SizeChunking",
+            "hist_axis_list": [
+                {
+                    "axis_class_name": "Regular",
+                    "kwargs": {
+                        "bins": 50,
+                        "start": 2.0,
+                        "stop": 38.0,
+                        "name": "value",
+                    },
+                }
+            ],
         },
         "SizeChunking": {"chunk_size": 1000},
     }
 )
 
-aggregator_peak_time = HistogramsAggregator(
-    hist.axis.Regular(50, 2.0, 38.0, name="value"),
-    config=config_peak_time,
-)
+aggregator_peak_time = HistogramsAggregator(config=config_peak_time)
 result_peak_time = aggregator_peak_time(
     table=table,
     col_name="peak_time",
@@ -110,6 +132,7 @@
 # -------------------------------------------------------------------
 # Inspect one pixel histogram in both chunks and both gain channels
 # -------------------------------------------------------------------
+# The plotted curves use the axes configured above via ``hist_axis_list``.
 pixel_index = 10
 gain_label = {0: "High Gain", 1: "Low Gain"}
 

src/ctapipe/monitoring/aggregator.py

@@ -34,10 +34,11 @@
 from astropy.stats import sigma_clip
 from astropy.table import Table
 from hist import Hist
+from traitlets import TraitError
 
 from ..containers import ChunkStatisticsContainer
 from ..core import Component
-from ..core.traits import AstroQuantity, Bool, ComponentName, Enum, Int
+from ..core.traits import AstroQuantity, Bool, ComponentName, Dict, Enum, Int, List
 
 
 class BaseChunking(Component, metaclass=ABCMeta):
@@ -400,39 +401,87 @@ class HistogramsAggregator(BaseAggregator):
     Aggregation is performed along axis=0 (the event dimension) for any N-dimensional data.
     """
 
-    def __init__(self, hist_axis, config=None, parent=None, **kwargs):
+    hist_axis_list = List(
+        trait=Dict(),
+        allow_none=False,
+        help=(
+            "List of histogram axis definitions. Each entry must contain "
+            "``axis_class_name`` and ``kwargs`` and is used to construct a "
+            "``hist.axis.<axis_class_name>(**kwargs)`` instance. If a single "
+            "axis is provided, it is applied to all pixels/channels. If multiple "
+            "axes are provided, they are applied per gain channel or first data "
+            "dimension. In multi-axis mode, all axes must have the same number of "
+            "bins. E.g. ``[{'axis_class_name': 'Regular', 'kwargs': {'bins': 40, 'start': 20.0, 'stop': 80.0}}]``."
+        ),
+    ).tag(config=True)
+
+    def _axis_from_dict(self, axis_config, entry_index):
+        """Create a hist axis from one dict in ``hist_axis_list``."""
+        missing_keys = {"axis_class_name", "kwargs"} - axis_config.keys()
+        if missing_keys:
+            raise TraitError(
+                f"Entry '{entry_index}' in the ``hist_axis_list`` trait "
+                f"is missing required key(s): {', '.join(sorted(missing_keys))}"
+            )
+
+        axis_kwargs = axis_config["kwargs"]
+        if not isinstance(axis_kwargs, dict):
+            raise TraitError(
+                f"Entry '{entry_index}' in the ``hist_axis_list`` trait has "
+                "a non-dict 'kwargs' value."
+            )
+
+        axis_class_name = axis_config["axis_class_name"]
+        axis_class = getattr(hist.axis, axis_class_name, None)
+        if axis_class is None or not callable(axis_class):
+            raise TraitError(
+                f"Entry '{entry_index}' in the ``hist_axis_list`` trait has "
+                f"unknown axis_class_name '{axis_class_name}'."
+            )
+
+        try:
+            return axis_class(**axis_kwargs)
+        except TypeError as err:
+            raise TraitError(
+                f"Failed to initialize hist.axis.{axis_class_name} for entry "
+                f"'{entry_index}' with kwargs={axis_kwargs}: {err}"
+            ) from err
+
+    def __init__(self, config=None, parent=None, **kwargs):
         """
         Parameters
         ----------
-        hist_axis : hist.axis or hist.Hist or list[hist.Hist]
-            Histogram definition for aggregation.
-            If a `hist.axis` is passed, one histogram is used for all channels.
-            If a `hist.Hist` or list of `hist.Hist` is passed, the first axis of
-            each hist defines the value axis. A list with length > 1 must match
-            the first dimension in `data.shape[1:]` (e.g. gain channels).
+        hist_axis_list : list[dict]
+            List of axis definitions. Each entry must contain
+            ``axis_class_name`` and ``kwargs`` and is used to construct a
+            ``hist.axis`` via ``hist.axis.<axis_class_name>(**kwargs)``.
         config : traitlets.loader.Config
             Configuration specified by config file or cmdline arguments
         parent : ctapipe.core.Component or ctapipe.core.Tool
             Parent of this component in the configuration hierarchy
         """
         super().__init__(config=config, parent=parent, **kwargs)
 
-        self.hist_axis = None
-        self.hist_templates = None
+        axis_list = [
+            self._axis_from_dict(axis_config, index)
+            for index, axis_config in enumerate(self.hist_axis_list)
+        ]
+        if len(axis_list) == 0:
+            raise TraitError("``hist_axis_list`` must contain at least one axis.")
 
-        if isinstance(hist_axis, list):
-            if len(hist_axis) == 0:
-                raise ValueError("hist_axis list must not be empty")
-            if not all(isinstance(h, Hist) for h in hist_axis):
-                raise TypeError("All elements of hist_axis list must be hist.Hist")
-            self.hist_templates = hist_axis
-        elif isinstance(hist_axis, Hist):
-            self.hist_templates = [hist_axis]
-        else:
-            self.hist_axis = hist_axis
+        self.hist_axis = axis_list[0]
+        self.hist_templates = None
+        if len(axis_list) > 1:
+            self.hist_templates = [Hist(axis) for axis in axis_list]
 
     def _get_hist_templates_for_shape(self, spatial_shape):
-        """Return one hist template per channel and validate compatibility."""
+        """
+        Return one histogram template per gain channel or first spatial dimension.
+
+        A single configured axis is reused for all channels. When multiple axes
+        are configured, the number of axes must match the first data dimension
+        and all axes must have the same bin count.
+        """
         if len(spatial_shape) == 0:
             n_channels = 1
         else:
@@ -490,6 +539,22 @@ def _build_data_mask(self, data, masked_elements_of_sample):
         invalid = ~np.isfinite(data)
         return mask | invalid
 
+    def _iter_channel_views(self, data, mask, n_events, spatial_shape):
+        """Yield per-channel data and mask views for histogram filling."""
+        if len(spatial_shape) == 0:
+            yield data, mask
+            return
+
+        for channel in range(spatial_shape[0]):
+            yield data[:, channel, ...], mask[:, channel, ...]
+
+    def _combine_edges(self, edges_per_channel):
+        """Return either a single edge array or a stacked per-channel array."""
+        if len(edges_per_channel) == 1:
+            return edges_per_channel[0]
+
+        return np.stack(edges_per_channel, axis=0)
+
     def _compute_single_histos(self, data, mask, n_events, spatial_shape):
         """Compute histograms using one value axis for all pixels/channels."""
         n_pixels = int(np.prod(spatial_shape))
@@ -519,7 +584,6 @@ def _compute_single_histos(self, data, mask, n_events, spatial_shape):
     def _compute_multi_histos(self, data, mask, n_events, spatial_shape):
         """Compute histograms with one template per channel."""
         templates = self._get_hist_templates_for_shape(spatial_shape)
-        n_channels = 1 if len(spatial_shape) == 0 else spatial_shape[0]
         channel_shape = spatial_shape[1:] if len(spatial_shape) > 1 else ()
         n_pixels_per_channel = (
             int(np.prod(channel_shape)) if len(channel_shape) > 0 else 1
@@ -530,9 +594,10 @@ def _compute_multi_histos(self, data, mask, n_events, spatial_shape):
         n_events_valid = np.zeros(spatial_shape, dtype=int)
         edges_per_channel = []
         hist_objects = []
-        is_scalar = len(spatial_shape) == 0
 
-        for channel in range(n_channels):
+        for channel, (channel_data, channel_mask) in enumerate(
+            self._iter_channel_views(data, mask, n_events, spatial_shape)
+        ):
             template = templates[channel]
             channel_hist = Hist(
                 template.axes[0],
@@ -542,13 +607,6 @@ def _compute_multi_histos(self, data, mask, n_events, spatial_shape):
             hist_objects.append(channel_hist)
             edges_per_channel.append(channel_hist.axes[0].edges)
 
-            if is_scalar:
-                channel_data = data
-                channel_mask = mask
-            else:
-                channel_data = data[:, channel, ...]
-                channel_mask = mask[:, channel, ...]
-
             flat_channel_data = channel_data.reshape(n_events, n_pixels_per_channel)
             flat_channel_mask = channel_mask.reshape(n_events, n_pixels_per_channel)
 
@@ -563,18 +621,10 @@ def _compute_multi_histos(self, data, mask, n_events, spatial_shape):
             channel_counts = channel_hist.values().reshape((n_bins,) + channel_shape)
             valid_events = np.sum(~flat_channel_mask, axis=0).reshape(channel_shape)
 
-            if is_scalar:
-                hist_counts[...] = channel_counts
-                n_events_valid[...] = valid_events
-            else:
-                hist_counts[:, channel, ...] = channel_counts
-                n_events_valid[channel, ...] = valid_events
+            hist_counts[:, channel, ...] = channel_counts
+            n_events_valid[channel, ...] = valid_events
 
-        edges = (
-            np.stack(edges_per_channel, axis=0)
-            if len(edges_per_channel) > 1
-            else edges_per_channel[0]
-        )
+        edges = self._combine_edges(edges_per_channel)
         return hist_objects, hist_counts, edges, n_events_valid
 
     def compute_histos(