Merge pull request #2529 from cta-observatory/extractor_gain

Change R1- and DL0-waveforms shape to (n_channels, n_pixels, n_samples)

docs/changes/2529.datamodel.rst

@@ -0,0 +1,6 @@
+Change R1- and DL0-waveforms datamodel shape from (n_pixels, n_samples)
+to be always (n_channels, n_pixels, n_samples). ``HDF5EventSource`` was adjusted
+accordingly to support also older datamodel versions.
+
+Re-introduce also the possibility of running ``ImageExtractor``\s on data
+consisting of multiple gain channels.

examples/core/table_writer_reader.py

@@ -15,7 +15,6 @@
 
 """
 
-
 ######################################################################
 # Caveats to think about: \* vector columns in Containers *can* be
 # written, but some lilbraries like Pandas can not read those (so you must
@@ -228,6 +227,7 @@ def create_stream(n_event):
 # note that here we can still access the metadata
 #
 
+
 table.attrs
 
 ######################################################################

src/ctapipe/calib/camera/calibrator.py

@@ -2,6 +2,7 @@
 Definition of the `CameraCalibrator` class, providing all steps needed to apply
 calibration and image extraction, as well as supporting algorithms.
 """
+from functools import cache
 
 import astropy.units as u
 import numpy as np
@@ -21,17 +22,27 @@
 __all__ = ["CameraCalibrator"]
 
 
-def _get_invalid_pixels(n_pixels, pixel_status, selected_gain_channel):
-    broken_pixels = np.zeros(n_pixels, dtype=bool)
-    index = np.arange(n_pixels)
+@cache
+def _get_pixel_index(n_pixels):
+    """Cached version of ``np.arange(n_pixels)``"""
+    return np.arange(n_pixels)
+
+
+def _get_invalid_pixels(n_channels, n_pixels, pixel_status, selected_gain_channel):
+    broken_pixels = np.zeros((n_channels, n_pixels), dtype=bool)
+
+    index = _get_pixel_index(n_pixels)
     masks = (
         pixel_status.hardware_failing_pixels,
         pixel_status.pedestal_failing_pixels,
         pixel_status.flatfield_failing_pixels,
     )
     for mask in masks:
         if mask is not None:
-            broken_pixels |= mask[selected_gain_channel, index]
+            if selected_gain_channel is not None:
+                broken_pixels |= mask[selected_gain_channel, index]
+            else:
+                broken_pixels |= mask
 
     return broken_pixels
 
@@ -188,7 +199,7 @@ def _calibrate_dl0(self, event, tel_id):
         )
 
         dl0_waveform = r1.waveform.copy()
-        dl0_waveform[~signal_pixels] = 0
+        dl0_waveform[:, ~signal_pixels] = 0
 
         dl0_pixel_status = r1.pixel_status.copy()
         # set dvr pixel bit in pixel_status for pixels kept by DVR
@@ -211,24 +222,26 @@ def _calibrate_dl1(self, event, tel_id):
         if self._check_dl0_empty(waveforms):
             return
 
-        n_pixels, n_samples = waveforms.shape
+        n_channels, n_pixels, n_samples = waveforms.shape
 
         selected_gain_channel = event.dl0.tel[tel_id].selected_gain_channel
         broken_pixels = _get_invalid_pixels(
+            n_channels,
             n_pixels,
             event.mon.tel[tel_id].pixel_status,
             selected_gain_channel,
         )
+        pixel_index = _get_pixel_index(n_pixels)
 
         dl1_calib = event.calibration.tel[tel_id].dl1
-        time_shift = event.calibration.tel[tel_id].dl1.time_shift
         readout = self.subarray.tel[tel_id].camera.readout
 
         # subtract any remaining pedestal before extraction
         if dl1_calib.pedestal_offset is not None:
             # this copies intentionally, we don't want to modify the dl0 data
-            # waveforms have shape (n_pixel, n_samples), pedestals (n_pixels, )
-            waveforms = waveforms - dl1_calib.pedestal_offset[:, np.newaxis]
+            # waveforms have shape (n_channels, n_pixel, n_samples), pedestals (n_pixels)
+            waveforms = waveforms.copy()
+            waveforms -= np.atleast_2d(dl1_calib.pedestal_offset)[..., np.newaxis]
 
         if n_samples == 1:
             # To handle ASTRI and dst
@@ -238,13 +251,18 @@ def _calibrate_dl1(self, event, tel_id):
             #   - Don't do anything if dl1 container already filled
             #   - Update on SST review decision
             dl1 = DL1CameraContainer(
-                image=waveforms[..., 0].astype(np.float32),
+                image=np.squeeze(waveforms).astype(np.float32),
                 peak_time=np.zeros(n_pixels, dtype=np.float32),
                 is_valid=True,
             )
         else:
             # shift waveforms if time_shift calibration is available
+            time_shift = dl1_calib.time_shift
+            remaining_shift = None
             if time_shift is not None:
+                if selected_gain_channel is not None:
+                    time_shift = time_shift[selected_gain_channel, pixel_index]
+
                 if self.apply_waveform_time_shift.tel[tel_id]:
                     sampling_rate = readout.sampling_rate.to_value(u.GHz)
                     time_shift_samples = time_shift * sampling_rate
@@ -264,11 +282,22 @@ def _calibrate_dl1(self, event, tel_id):
             )
 
             # correct non-integer remainder of the shift if given
-            if self.apply_peak_time_shift.tel[tel_id] and time_shift is not None:
+            if self.apply_peak_time_shift.tel[tel_id] and remaining_shift is not None:
                 dl1.peak_time -= remaining_shift
 
         # Calibrate extracted charge
-        dl1.image *= dl1_calib.relative_factor / dl1_calib.absolute_factor
+        if (
+            dl1_calib.relative_factor is not None
+            and dl1_calib.absolute_factor is not None
+        ):
+            if selected_gain_channel is None:
+                dl1.image *= dl1_calib.relative_factor / dl1_calib.absolute_factor
+            else:
+                corr = (
+                    dl1_calib.relative_factor[selected_gain_channel, pixel_index]
+                    / dl1_calib.absolute_factor[selected_gain_channel, pixel_index]
+                )
+                dl1.image *= corr
 
         # handle invalid pixels
         if self.invalid_pixel_handler is not None:
@@ -309,21 +338,21 @@ def shift_waveforms(waveforms, time_shift_samples):
 
     Parameters
     ----------
-    waveforms: ndarray of shape (n_pixels, n_samples)
+    waveforms: ndarray of shape (n_channels, n_pixels, n_samples)
         The waveforms to shift
-    time_shift_samples: ndarray of shape (n_pixels, )
+    time_shift_samples: ndarray
         The shift to apply in units of samples.
         Waveforms are shifted to the left by the smallest integer
         that minimizes inter-pixel differences.
 
     Returns
     -------
-    shifted_waveforms: ndarray of shape (n_pixels, n_samples)
+    shifted_waveforms: ndarray of shape (n_channels, n_pixels, n_samples)
         The shifted waveforms
-    remaining_shift: ndarray of shape (n_pixels, )
+    remaining_shift: ndarray
         The remaining shift after applying the integer shift to the waveforms.
     """
-    mean_shift = time_shift_samples.mean()
+    mean_shift = time_shift_samples.mean(axis=-1, keepdims=True)
     integer_shift = np.round(time_shift_samples - mean_shift).astype("int16")
     remaining_shift = time_shift_samples - integer_shift
     shifted_waveforms = _shift_waveforms_by_integer(waveforms, integer_shift)

src/ctapipe/calib/camera/flatfield.py

@@ -182,9 +182,11 @@ def _extract_charge(self, event) -> DL1CameraContainer:
         """
 
         waveforms = event.r1.tel[self.tel_id].waveform
+        n_channels, n_pixels, _ = waveforms.shape
         selected_gain_channel = event.r1.tel[self.tel_id].selected_gain_channel
         broken_pixels = _get_invalid_pixels(
-            n_pixels=waveforms.shape[-2],
+            n_channels=n_channels,
+            n_pixels=n_pixels,
             pixel_status=event.mon.tel[self.tel_id].pixel_status,
             selected_gain_channel=selected_gain_channel,
         )
@@ -215,20 +217,15 @@ def calculate_relative_gain(self, event):
         if self.n_events_seen == self.sample_size:
             self.n_events_seen = 0
 
-        # real data
         trigger_time = event.trigger.time
-        if event.meta["origin"] != "hessio":
-            hardware_or_pedestal_mask = np.logical_or(
-                event.mon.tel[self.tel_id].pixel_status.hardware_failing_pixels,
-                event.mon.tel[self.tel_id].pixel_status.pedestal_failing_pixels,
-            )
-            pixel_mask = np.logical_or(
-                hardware_or_pedestal_mask,
-                event.mon.tel[self.tel_id].pixel_status.flatfield_failing_pixels,
-            )
-
-        else:  # patches for MC data
-            pixel_mask = np.zeros(waveform.shape[1], dtype=bool)
+        hardware_or_pedestal_mask = np.logical_or(
+            event.mon.tel[self.tel_id].pixel_status.hardware_failing_pixels,
+            event.mon.tel[self.tel_id].pixel_status.pedestal_failing_pixels,
+        )
+        pixel_mask = np.logical_or(
+            hardware_or_pedestal_mask,
+            event.mon.tel[self.tel_id].pixel_status.flatfield_failing_pixels,
+        )
 
         if self.n_events_seen == 0:
             self.time_start = trigger_time

src/ctapipe/calib/camera/gainselection.py

@@ -1,6 +1,7 @@
 """
 Algorithms to select correct gain channel
 """
+
 from abc import abstractmethod
 from enum import IntEnum
 
@@ -47,16 +48,15 @@ def __call__(self, waveforms):
             Shape: n_pix
             Dtype: int8
         """
-        if waveforms.ndim == 2:  # Return None if already gain selected
-            return None
-        elif waveforms.ndim == 3:
-            n_channels, n_pixels, _ = waveforms.shape
-            if n_channels == 1:  # Must be first channel if only one channel
-                return np.zeros(n_pixels, dtype=np.int8)
-            else:
-                return self.select_channel(waveforms)
-        else:
-            raise ValueError(f"Cannot handle waveform array of shape: {waveforms.ndim}")
+        if waveforms.ndim != 3:
+            raise ValueError(
+                f"Cannot handle waveform array of shape: {waveforms.shape}"
+            )
+
+        n_channels, n_pixels, _ = waveforms.shape
+        if n_channels == 1:
+            return np.zeros(n_pixels, dtype=np.int8)
+        return self.select_channel(waveforms)
 
     @abstractmethod
     def select_channel(self, waveforms):

src/ctapipe/calib/camera/pedestals.py

@@ -211,13 +211,14 @@ def _extract_charge(self, event) -> DL1CameraContainer:
         DL1CameraContainer
         """
         waveforms = event.r1.tel[self.tel_id].waveform
+        n_channels, n_pixels, _ = waveforms.shape
         selected_gain_channel = event.r1.tel[self.tel_id].selected_gain_channel
         broken_pixels = _get_invalid_pixels(
-            n_pixels=waveforms.shape[-2],
+            n_channels=n_channels,
+            n_pixels=n_pixels,
             pixel_status=event.mon.tel[self.tel_id].pixel_status,
             selected_gain_channel=selected_gain_channel,
         )
-
         # Extract charge and time
         if self.extractor:
             return self.extractor(
@@ -247,10 +248,7 @@ def calculate_pedestals(self, event):
 
         # real data
         trigger_time = event.trigger.time
-        if event.meta["origin"] != "hessio":
-            pixel_mask = event.mon.tel[self.tel_id].pixel_status.hardware_failing_pixels
-        else:  # patches for MC data
-            pixel_mask = np.zeros(waveform.shape[1], dtype=bool)
+        pixel_mask = event.mon.tel[self.tel_id].pixel_status.hardware_failing_pixels
 
         if self.n_events_seen == 0:
             self.time_start = trigger_time