Keep channel dimension on waevforms, even if event is gain-selected

Author: jlenain

src/ctapipe_io_nectarcam/__init__.py

@@ -1158,9 +1158,17 @@ def fill_r0r1_camera_container(self, zfits_event):
         )
         expected_pixels = self.nectarcam_service.pixel_ids
 
+        has_low_gain = (zfits_event.pixel_status & PixelStatus.LOW_GAIN_STORED).astype(
+            bool
+        )
+        has_high_gain = (
+            zfits_event.pixel_status & PixelStatus.HIGH_GAIN_STORED
+        ).astype(bool)
+        not_broken = (has_low_gain | has_high_gain).astype(bool)
+
         # broken pixels have both false, so gain selected means checking
         # if there are any pixels where exactly one of high or low gain is stored
-        # gain_selected = np.any(has_low_gain != has_high_gain)
+        gain_selected = np.any(has_low_gain != has_high_gain)
 
         # fill value for broken pixels
 
@@ -1177,15 +1185,41 @@ def fill_r0r1_camera_container(self, zfits_event):
         # we assume that either all pixels are gain selected or none
         # only broken pixels are allowed to be missing completely
 
-        reshaped_waveform = zfits_event.waveform.reshape(N_GAINS, n_pixels, n_samples)
-        # re-order the waveform following the expected_pixels_id values
-        #  could also just do waveform = reshaped_waveform[np.argsort(expected_ids)]
-        reordered_waveform = np.full(
-            (N_GAINS, N_PIXELS, N_SAMPLES), fill, dtype=dtype
-        )  # VIM : use empty ?
-        reordered_waveform[:, expected_pixels, :] = reshaped_waveform
-        r0 = R0CameraContainer(waveform=reordered_waveform)
-        r1 = R1CameraContainer()
+        if gain_selected:
+            # print("GAIN SELECTED")
+            selected_gain = np.where(has_high_gain, 0, 1)
+            waveform = np.full(
+                (N_GAINS, n_pixels, n_samples), fill, dtype=dtype
+            )  # VIM : Replace full by empty ?
+            waveform[not_broken] = zfits_event.waveform.reshape((-1, n_samples))[
+                not_broken
+            ]
+
+            reordered_waveform = np.full(
+                (N_GAINS, N_PIXELS, N_SAMPLES), fill, dtype=dtype
+            )  # VIM : Replace full by empty ?
+            reordered_waveform[expected_pixels] = waveform
+
+            reordered_selected_gain = np.full(N_PIXELS, -1, dtype=np.int8)
+            reordered_selected_gain[expected_pixels] = selected_gain
+
+            r0 = R0CameraContainer()
+            r1 = R1CameraContainer(
+                waveform=reordered_waveform,
+                selected_gain_channel=reordered_selected_gain,
+            )
+        else:
+            reshaped_waveform = zfits_event.waveform.reshape(
+                N_GAINS, n_pixels, n_samples
+            )
+            # re-order the waveform following the expected_pixels_id values
+            #  could also just do waveform = reshaped_waveform[np.argsort(expected_ids)]
+            reordered_waveform = np.full(
+                (N_GAINS, N_PIXELS, N_SAMPLES), fill, dtype=dtype
+            )  # VIM : use empty ?
+            reordered_waveform[:, expected_pixels, :] = reshaped_waveform
+            r0 = R0CameraContainer(waveform=reordered_waveform)
+            r1 = R1CameraContainer()
 
         if not self.pre_v6_data:
             if r0.waveform is not None: