Change gain selection

protopipe/pipeline/event_preparer.py

@@ -7,6 +7,7 @@
 
 # CTAPIPE utilities
 from ctapipe.calib import CameraCalibrator
+from ctapipe.calib.camera.gainselection import GainSelector
 from ctapipe.image.extractor import LocalPeakWindowSum
 from ctapipe.image import hillas
 from ctapipe.utils.CutFlow import CutFlow
@@ -118,6 +119,37 @@ def largest_island(islands_labels):
     return islands_labels == np.argmax(np.bincount(islands_labels[islands_labels > 0]))
 
 
+class MyCameraCalibrator(CameraCalibrator):
+    """Create a child class of CameraCalibrator."""
+
+    def _calibrate_dl0(self, event, telid):
+        """Override class method to perform gain selection at R0 instead of R1.
+
+        Then select R1 waveforms using R0-selected gain channels.
+        """
+        waveforms_r0 = event.r0.tel[telid].waveform
+        _, selected_gain_channel = self.gain_selector(waveforms_r0)
+
+        waveforms_r1 = event.r1.tel[telid].waveform
+        if self._check_r1_empty(waveforms_r1):
+            return
+
+        # Use R0-selected gain channels to select R1 waveforms
+        _, n_pixels, _ = waveforms_r1.shape
+        waveforms_gs = waveforms_r1[selected_gain_channel, np.arange(n_pixels)]
+        if selected_gain_channel is not None:
+            event.r1.tel[telid].selected_gain_channel = selected_gain_channel
+        else:
+            if event.r1.tel[telid].selected_gain_channel is None:
+                raise ValueError(
+                    "EventSource is loading pre-gainselected waveforms "
+                    "without filling the selected_gain_channel container"
+                )
+
+        reduced_waveforms = self.data_volume_reducer(waveforms_gs)
+        event.dl0.tel[telid].waveform = reduced_waveforms
+
+
 # ==============================================================================
 
 
@@ -228,9 +260,13 @@ def __init__(self, config, mode, event_cutflow=None, image_cutflow=None):
         cfg = Config()
         cfg["ChargeExtractorFactory"]["window_width"] = 5
         cfg["ChargeExtractorFactory"]["window_shift"] = 2
+        cfg["ThresholdGainSelector"]["threshold"] = 4000.0  # 40 @ R1!
         extractor = LocalPeakWindowSum(config=cfg)
+        gain_selector = GainSelector.from_name("ThresholdGainSelector", config=cfg)
 
-        self.calib = CameraCalibrator(config=cfg, image_extractor=extractor)
+        self.calib = MyCameraCalibrator(
+            config=cfg, gain_selector=gain_selector, image_extractor=extractor
+        )
 
         # Reconstruction
         self.shower_reco = HillasReconstructor()

protopipe/scripts/write_dl1.py

@@ -101,6 +101,7 @@ class StoredImages(tb.IsDescription):
         tel_id = tb.Int16Col(dflt=1, pos=1)
         dl1_phe_image = tb.Float32Col(shape=(1855), pos=2)
         mc_phe_image = tb.Float32Col(shape=(1855), pos=3)
+        mc_energy = tb.Float32Col(dflt=1, pos=4)
 
     # Declaration of the column descriptor for the file containing DL1 data
     class EventFeatures(tb.IsDescription):
@@ -332,6 +333,7 @@ class EventFeatures(tb.IsDescription):
                     images_phe[cam_id]["tel_id"] = tel_id
                     images_phe[cam_id]["dl1_phe_image"] = dl1_phe_image
                     images_phe[cam_id]["mc_phe_image"] = mc_phe_image
+                    images_phe[cam_id]["mc_energy"] = event.mc.energy.value  # TeV
 
                     images_phe[cam_id].append()