Add gain selection to calibration chain

ctapipe/calib/camera/calibrator.py

@@ -1,5 +1,6 @@
 import numpy as np
 from ctapipe.core import Component
+from ctapipe.calib.camera.gainselection import ManualGainSelector
 from ctapipe.image.reducer import NullDataVolumeReducer
 from ctapipe.image.extractor import NeighborPeakWindowSum
 import warnings
@@ -73,6 +74,7 @@ class CameraCalibrator(Component):
     the DL1 data level in the event container.
     """
     def __init__(self, config=None, parent=None,
+                 gain_selector=None,
                  data_volume_reducer=None,
                  image_extractor=None,
                  **kwargs):
@@ -87,6 +89,9 @@ def __init__(self, config=None, parent=None,
             Tool executable that is calling this component.
             Passes the correct logger to the component.
             Set to None if no Tool to pass.
+        gain_selector : ctapipe.calib.camera.gainselection.GainSelector
+            The GainSelector to use. If None, then ManualGainSelector will be
+            used, which by default selects the high/first gain channel.
         data_volume_reducer : ctapipe.image.reducer.DataVolumeReducer
             The DataVolumeReducer to use. If None, then
             NullDataVolumeReducer will be used by default, and waveforms
@@ -101,6 +106,10 @@ def __init__(self, config=None, parent=None,
         self._r1_empty_warn = False
         self._dl0_empty_warn = False
 
+        if gain_selector is None:
+            gain_selector = ManualGainSelector(parent=self)
+        self.gain_selector = gain_selector
+
         if data_volume_reducer is None:
             data_volume_reducer = NullDataVolumeReducer(parent=self)
         self.data_volume_reducer = data_volume_reducer
@@ -126,6 +135,7 @@ def _get_correction(self, event, telid):
         ndarray
         """
         try:
+            gain_selection = event.mon.tel[telid].gain_selection
             shift = self.image_extractor.window_shift
             width = self.image_extractor.window_width
             shape = event.mc.tel[telid].reference_pulse_shape
@@ -134,7 +144,8 @@ def _get_correction(self, event, telid):
             time_slice = event.mc.tel[telid].time_slice
             correction = integration_correction(n_chan, shape, step,
                                                 time_slice, width, shift)
-            return correction
+            pixel_correction = correction[gain_selection]
+            return pixel_correction
         except (AttributeError, KeyError):
             # Don't apply correction when window_shift or window_width
             # does not exist in extractor, or when container does not have
@@ -165,15 +176,34 @@ def _calibrate_dl0(self, event, telid):
         waveforms = event.r1.tel[telid].waveform
         if self._check_r1_empty(waveforms):
             return
-        # TODO: Add gain selection
-        reduced_waveforms = self.data_volume_reducer(waveforms)
+
+        # Perform gain selection. This is typically not the responsibility of
+        # ctapipe; DL0 (and R1) waveforms are aleady gain selected and
+        # therefore single channel. However, the waveforms read from
+        # simtelarray do not have the gain selection applied, and so must be
+        # done as part of the calibration step to ensure the correct
+        # waveform dimensions.
+        waveforms_gs, pixel_channel = self.gain_selector(waveforms)
+        if pixel_channel is not None:
+            event.mon.tel[telid].gain_selection = pixel_channel
+        else:
+            # If pixel_channel is None, then waveforms has already been
+            # pre-gainselected, and presumably the gain_selection container is
+            # filled by the EventSource
+            if event.mon.tel[telid].gain_selection is None:
+                raise ValueError(
+                    "EventSource is loading pre-gainselected waveforms "
+                    "without filling the gain_selection container"
+                )
+
+        reduced_waveforms = self.data_volume_reducer(waveforms_gs)
         event.dl0.tel[telid].waveform = reduced_waveforms
 
     def _calibrate_dl1(self, event, telid):
         waveforms = event.dl0.tel[telid].waveform
         if self._check_dl0_empty(waveforms):
             return
-        n_samples = waveforms.shape[2]
+        n_pixels, n_samples = waveforms.shape
         if n_samples == 1:
             # To handle ASTRI and dst
             # TODO: Improved handling of ASTRI and dst
@@ -182,7 +212,7 @@ def _calibrate_dl1(self, event, telid):
             #   - Don't do anything if dl1 container already filled
             #   - Update on SST review decision
             corrected_charge = waveforms[..., 0]
-            pulse_time = np.zeros(waveforms.shape[0:2])
+            pulse_time = np.zeros(n_pixels)
         else:
             # TODO: pass camera to ImageExtractor.__init__
             if self.image_extractor.requires_neighbors():
@@ -192,7 +222,7 @@ def _calibrate_dl1(self, event, telid):
 
             # Apply integration correction
             # TODO: Remove integration correction
-            correction = self._get_correction(event, telid)[:, None]
+            correction = self._get_correction(event, telid)
             corrected_charge = charge * correction
 
         event.dl1.tel[telid].image = corrected_charge

ctapipe/calib/camera/gainselection.py

@@ -45,14 +45,17 @@ def __call__(self, waveforms):
             Shape: (n_pix, n_samples)
         """
         if waveforms.ndim == 2:  # Return if already gain selected
-            return waveforms
+            pixel_channel = None  # Provided by EventSource
+            return waveforms, pixel_channel
         elif waveforms.ndim == 3:
             n_channels, n_pixels, _ = waveforms.shape
             if n_channels == 1:  # Reduce if already single channel
-                return waveforms[0]
+                pixel_channel = np.zeros(n_pixels, dtype=int)
+                return waveforms[0], pixel_channel
             else:
                 pixel_channel = self.select_channel(waveforms)
-                return waveforms[pixel_channel, np.arange(n_pixels)]
+                gain_selected = waveforms[pixel_channel, np.arange(n_pixels)]
+                return gain_selected, pixel_channel
         else:
             raise ValueError(
                 f"Cannot handle waveform array of shape: {waveforms.ndim}"
@@ -92,7 +95,8 @@ class ManualGainSelector(GainSelector):
     ).tag(config=True)
 
     def select_channel(self, waveforms):
-        return GainChannel[self.channel]
+        n_pixels = waveforms.shape[1]
+        return np.full(n_pixels, GainChannel[self.channel])
 
 
 class ThresholdGainSelector(GainSelector):

ctapipe/calib/camera/tests/test_calibrator.py

@@ -2,17 +2,48 @@
     CameraCalibrator,
     integration_correction,
 )
+from ctapipe.io.containers import DataContainer
 from ctapipe.image.extractor import LocalPeakWindowSum
 from traitlets.config.configurable import Config
 import pytest
+import numpy as np
 
 
 def test_camera_calibrator(example_event):
     telid = list(example_event.r0.tel)[0]
     calibrator = CameraCalibrator()
     calibrator(example_event)
     image = example_event.dl1.tel[telid].image
+    pulse_time = example_event.dl1.tel[telid].pulse_time
     assert image is not None
+    assert pulse_time is not None
+    assert image.shape == (1764,)
+    assert pulse_time.shape == (1764,)
+
+
+def test_select_gain():
+    n_channels = 2
+    n_pixels = 2048
+    n_samples = 128
+    telid = 0
+
+    calibrator = CameraCalibrator()
+
+    event = DataContainer()
+    event.r1.tel[telid].waveform = np.ones((n_channels, n_pixels, n_samples))
+    calibrator._calibrate_dl0(event, telid)
+    assert event.dl0.tel[telid].waveform.shape == (n_pixels, n_samples)
+
+    event = DataContainer()
+    event.r1.tel[telid].waveform = np.ones((n_pixels, n_samples))
+    with pytest.raises(ValueError):
+        calibrator._calibrate_dl0(event, telid)
+
+    event = DataContainer()
+    event.r1.tel[telid].waveform = np.ones((n_pixels, n_samples))
+    event.mon.tel[telid].gain_selection = np.zeros(n_pixels)
+    calibrator._calibrate_dl0(event, telid)
+    assert event.dl0.tel[telid].waveform.shape == (n_pixels, n_samples)
 
 
 def test_manual_extractor():
@@ -55,7 +86,7 @@ def test_integration_correction_no_ref_pulse(example_event):
     calibrator = CameraCalibrator()
     calibrator._calibrate_dl0(example_event, telid)
     correction = calibrator._get_correction(example_event, telid)
-    assert correction[0] == 1
+    assert (correction == 1).all()
 
 
 def test_check_r1_empty(example_event):

ctapipe/calib/camera/tests/test_flatfield.py

@@ -1,9 +1,21 @@
 import numpy as np
 from ctapipe.calib.camera.flatfield import *
-from ctapipe.io.containers import EventAndMonDataContainer
+from ctapipe.io.containers import DataContainer
 from traitlets.config.loader import Config
 
 
+# class FixedWindowSum2Chan(FixedWindowSum):
+#
+#     def __call__(self, waveforms):
+#         start = self.window_start
+#         end = self.window_start + self.window_width
+#         charge = waveforms[..., start:end].sum(2)
+#         pulse_time = extract_pulse_time_around_peak(
+#             waveforms, self.window_start, self.window_width, 0
+#         )
+#         return charge, pulse_time
+
+
 def test_flasherflatfieldcalculator():
     """test of flasherFlatFieldCalculator"""
     tel_id = 0
@@ -22,7 +34,7 @@ def test_flasherflatfieldcalculator():
                                                sample_size=n_events,
                                                tel_id=tel_id, config=config)
     # create one event
-    data = EventAndMonDataContainer()
+    data = DataContainer()
     data.meta['origin'] = 'test'
 
     # initialize mon and r1 data

ctapipe/calib/camera/tests/test_gainselection.py

@@ -14,12 +14,29 @@ def test_gain_selector():
     waveforms[1] *= 2
 
     gain_selector = TestGainSelector()
-    waveforms_gs = gain_selector(waveforms)
+    waveforms_gs, pixel_channel = gain_selector(waveforms)
     np.testing.assert_equal(waveforms[GainChannel.HIGH], waveforms_gs)
-    waveforms_gs = gain_selector(waveforms[0])
-    np.testing.assert_equal(waveforms_gs, waveforms[0])
-    waveforms_gs = gain_selector(waveforms[[0]])
-    np.testing.assert_equal(waveforms_gs, waveforms[0])
+    np.testing.assert_equal(pixel_channel, 0)
+
+
+def test_pre_selected():
+    shape = (2048, 128)
+    waveforms = np.zeros(shape)
+
+    gain_selector = TestGainSelector()
+    waveforms_gs, pixel_channel = gain_selector(waveforms)
+    assert waveforms.shape == waveforms_gs.shape
+    assert pixel_channel is None
+
+
+def test_single_channel():
+    shape = (1, 2048, 128)
+    waveforms = np.zeros(shape)
+
+    gain_selector = TestGainSelector()
+    waveforms_gs, pixel_channel = gain_selector(waveforms)
+    assert waveforms_gs.shape == (2048, 128)
+    assert (pixel_channel == 0).all()
 
 
 def test_manual_gain_selector():
@@ -28,12 +45,14 @@ def test_manual_gain_selector():
     waveforms[1] *= 2
 
     gs_high = ManualGainSelector(channel="HIGH")
-    waveforms_gs_high = gs_high(waveforms)
-    np.testing.assert_equal(waveforms[GainChannel.HIGH], waveforms_gs_high)
+    waveforms_gs, pixel_channel = gs_high(waveforms)
+    np.testing.assert_equal(waveforms[GainChannel.HIGH], waveforms_gs)
+    np.testing.assert_equal(pixel_channel, 0)
 
-    gs_high = ManualGainSelector(channel="LOW")
-    waveforms_gs_low = gs_high(waveforms)
-    np.testing.assert_equal(waveforms[GainChannel.LOW], waveforms_gs_low)
+    gs_low = ManualGainSelector(channel="LOW")
+    waveforms_gs, pixel_channel = gs_low(waveforms)
+    np.testing.assert_equal(waveforms[GainChannel.LOW], waveforms_gs)
+    np.testing.assert_equal(pixel_channel, 1)
 
 
 def test_threshold_gain_selector():
@@ -43,6 +62,8 @@ def test_threshold_gain_selector():
     waveforms[0, 0] = 100
 
     gain_selector = ThresholdGainSelector(threshold=50)
-    waveforms_gs = gain_selector(waveforms)
+    waveforms_gs, pixel_channel = gain_selector(waveforms)
     assert (waveforms_gs[0] == 1).all()
     assert (waveforms_gs[np.arange(1, 2048)] == 0).all()
+    assert pixel_channel[0] == 1
+    assert (pixel_channel[np.arange(1, 2048)] == 0).all()

ctapipe/calib/camera/tests/test_pedestals.py

@@ -1,6 +1,6 @@
 import numpy as np
 from ctapipe.calib.camera.pedestals import *
-from ctapipe.io.containers import EventAndMonDataContainer
+from ctapipe.io.containers import DataContainer
 
 
 def test_pedestal_calculator():
@@ -16,7 +16,7 @@ def test_pedestal_calculator():
                                         sample_size=n_events,
                                         tel_id=tel_id)
     # create one event
-    data = EventAndMonDataContainer()
+    data = DataContainer()
     data.meta['origin'] = 'test'
 
     # fill the values necessary for the pedestal calculation