Adapt InvalidPixelHandler to work with 2 gain channels

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,
         )

src/ctapipe/calib/camera/pedestals.py

@@ -211,9 +211,11 @@ 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,
         )

src/ctapipe/image/invalid_pixels.py

@@ -38,6 +38,7 @@ def __call__(
             Array of pixel peak_time values
         pixel_mask : np.ndarray
             Boolean mask of the pixels to be interpolated
+            Shape: (n_channels, n_pixels)
 
         Returns
         -------
@@ -66,6 +67,7 @@ def __call__(self, tel_id, image, peak_time, pixel_mask):
             Array of pixel peak_time values
         pixel_mask : np.ndarray
             Boolean mask of the pixels to be interpolated
+            Shape: (n_channels, n_pixels)
 
         Returns
         -------
@@ -76,31 +78,35 @@ def __call__(self, tel_id, image, peak_time, pixel_mask):
         """
         geometry = self.subarray.tel[tel_id].camera.geometry
 
-        n_interpolated = np.count_nonzero(pixel_mask)
-        if n_interpolated == 0:
+        n_interpolated = np.count_nonzero(pixel_mask, axis=-1, keepdims=True)
+        if (n_interpolated == 0).all():
             return image, peak_time
 
-        # exclude to-be-interpolated pixels from neighbors
-        neighbors = geometry.neighbor_matrix[pixel_mask] & ~pixel_mask
-
-        index, neighbor = np.nonzero(neighbors)
-        image_sum = np.zeros(n_interpolated, dtype=image.dtype)
-        count = np.zeros(n_interpolated, dtype=int)
-        peak_time_sum = np.zeros(n_interpolated, dtype=peak_time.dtype)
-
-        # calculate average of image and peak_time
-        np.add.at(count, index, 1)
-        np.add.at(image_sum, index, image[neighbor])
-        np.add.at(peak_time_sum, index, peak_time[neighbor])
-
-        valid = count > 0
-        np.divide(image_sum, count, out=image_sum, where=valid)
-        np.divide(peak_time_sum, count, out=peak_time_sum, where=valid)
-
-        peak_time = peak_time.copy()
-        peak_time[pixel_mask] = peak_time_sum
-
-        image = image.copy()
-        image[pixel_mask] = image_sum
-
-        return image, peak_time
+        image = np.atleast_2d(image)
+        peak_time = np.atleast_2d(peak_time)
+        for ichannel in range(image.shape[-2]):
+            if n_interpolated[ichannel] == 0:
+                continue
+            # exclude to-be-interpolated pixels from neighbors
+            neighbors = (
+                geometry.neighbor_matrix[pixel_mask[ichannel]] & ~pixel_mask[ichannel]
+            )
+
+            index, neighbor = np.nonzero(neighbors)
+            image_sum = np.zeros(n_interpolated[ichannel], dtype=image.dtype)
+            count = np.zeros(n_interpolated[ichannel], dtype=int)
+            peak_time_sum = np.zeros(n_interpolated[ichannel], dtype=peak_time.dtype)
+
+            # calculate average of image and peak_time
+            np.add.at(count, index, 1)
+            np.add.at(image_sum, index, image[ichannel, neighbor])
+            np.add.at(peak_time_sum, index, peak_time[ichannel, neighbor])
+
+            valid = count > 0
+            np.divide(image_sum, count, out=image_sum, where=valid)
+            np.divide(peak_time_sum, count, out=peak_time_sum, where=valid)
+
+            peak_time[ichannel, pixel_mask[ichannel]] = peak_time_sum
+            image[ichannel, pixel_mask[ichannel]] = image_sum
+
+        return np.squeeze(image), np.squeeze(peak_time)

src/ctapipe/image/reducer.py

@@ -183,7 +183,9 @@ def select_pixels(self, waveforms, tel_id=None, selected_gain_channel=None):
         # Pulse-integrate waveforms
         extractor = self.image_extractors[self.image_extractor_type.tel[tel_id]]
         # do not treat broken pixels in data volume reduction
-        broken_pixels = np.zeros(camera_geom.n_pixels, dtype=bool)
+        broken_pixels = np.zeros(
+            (waveforms.shape[-3], camera_geom.n_pixels), dtype=bool
+        )
         dl1: DL1CameraContainer = extractor(
             waveforms,
             tel_id=tel_id,

src/ctapipe/image/tests/test_invalid_pixels.py

@@ -12,16 +12,16 @@ def test_neighbor_average(prod5_gamma_simtel_path):
 
     neighbor_average = NeighborAverage(subarray)
 
-    broken_pixels = np.zeros(geometry.n_pixels, dtype=bool)
+    broken_pixels = np.zeros((2, geometry.n_pixels), dtype=bool)
 
     # one border pixel, one non-border pixel
     border = geometry.get_border_pixel_mask(1)
-    broken_pixels[np.nonzero(border)[0][0]] = True
-    broken_pixels[np.nonzero(~border)[0][0]] = True
+    broken_pixels[:, np.nonzero(border)[0][0]] = True
+    broken_pixels[:, np.nonzero(~border)[0][0]] = True
 
-    image = np.ones(geometry.n_pixels)
+    image = np.ones((2, geometry.n_pixels))
     image[broken_pixels] = 9999
-    peak_time = np.full(geometry.n_pixels, 20.0)
+    peak_time = np.full((2, geometry.n_pixels), 20.0)
     peak_time[broken_pixels] = -1
 
     interpolated_image, interpolated_peaktime = neighbor_average(
@@ -42,20 +42,20 @@ def test_negative_image(prod5_gamma_simtel_path):
 
     neighbor_average = NeighborAverage(subarray)
 
-    broken_pixels = np.zeros(geometry.n_pixels, dtype=bool)
+    broken_pixels = np.zeros((1, geometry.n_pixels), dtype=bool)
 
     # one border pixel, one non-border pixel
     border = geometry.get_border_pixel_mask(1)
-    broken_pixels[np.nonzero(border)[0][0]] = True
-    broken_pixels[np.nonzero(~border)[0][0]] = True
+    broken_pixels[:, np.nonzero(border)[0][0]] = True
+    broken_pixels[:, np.nonzero(~border)[0][0]] = True
 
     image = np.full(geometry.n_pixels, -10.0)
-    image[broken_pixels] = 9999
+    image[broken_pixels[0]] = 9999
     peak_time = np.full(geometry.n_pixels, 20.0)
-    peak_time[broken_pixels] = -1
+    peak_time[broken_pixels[0]] = -1
 
     interpolated_image, interpolated_peaktime = neighbor_average(
         tel_id=1, image=image, peak_time=peak_time, pixel_mask=broken_pixels
     )
-    assert np.allclose(interpolated_image[broken_pixels], -10)
-    assert np.allclose(interpolated_peaktime[broken_pixels], 20.0)
+    assert np.allclose(interpolated_image[broken_pixels[0]], -10)
+    assert np.allclose(interpolated_peaktime[broken_pixels[0]], 20.0)