Merge pull request #2943 from cta-observatory/sim_calib_gain_select

allow gain selection when skipping R1 calibration in SimTelEventSourve

Author: kosack

docs/changes/2943.feature.rst

@@ -0,0 +1 @@
+Allow to select gain for calibration events in SimTelEventSource, when skipping R1 calibration.

src/ctapipe/io/simteleventsource.py

@@ -3,6 +3,7 @@
 import warnings
 from contextlib import nullcontext
 from enum import Enum, IntFlag, auto, unique
+from functools import lru_cache
 from gzip import GzipFile
 from io import BufferedReader
 
@@ -119,6 +120,11 @@ class SimTelTriggerMask(IntFlag):
     RANDOM_MONO = auto()
 
 
+@lru_cache()
+def _get_pixel_index(n_pixels):
+    return np.arange(n_pixels)
+
+
 def _trigger_mask_to_event_type(trigger_mask):
     trigger_mask = SimTelTriggerMask(int(trigger_mask))
 
@@ -333,11 +339,10 @@ def _telescope_from_meta(telescope_meta, mirror_area):
 
 
 def apply_simtel_r1_calibration(
-    r0_waveforms, pedestal, factor, gain_selector, calib_scale=1.0, calib_shift=0.0
+    r0_waveforms, pedestal, factor, calib_scale=1.0, calib_shift=0.0
 ):
     """
     Perform the R1 calibration for R0 simtel waveforms. This includes:
-        - Gain selection
         - Pedestal subtraction
         - Conversion of samples into units proportional to photoelectrons
           (If the full signal in the waveform was integrated, then the resulting
@@ -356,7 +361,6 @@ def apply_simtel_r1_calibration(
         Conversion factor between R0 waveform samples and ~p.e., stored in the
         simtel file for each gain channel
         Shape: (n_channels, n_pixels)
-    gain_selector : ctapipe.calib.camera.gainselection.GainSelector
     calib_scale : float
         Extra global scale factor for calibration.
         Conversion factor to transform the integrated charges
@@ -369,27 +373,47 @@ def apply_simtel_r1_calibration(
     Returns
     -------
     r1_waveforms : ndarray
-        Calibrated waveforms
+        Calibrated waveforms not gain-selected, in units of photoelectrons (p.e.).
         Shape: (n_channels, n_pixels, n_samples)
-    selected_gain_channel : ndarray
-        The gain channel selected for each pixel
-        Shape: (n_pixels)
     """
-    n_pixels = r0_waveforms.shape[-2]
     ped = pedestal[..., np.newaxis]
     factor = factor[..., np.newaxis]
     gain = factor * calib_scale
 
     r1_waveforms = (r0_waveforms - ped) * gain + calib_shift
 
-    if gain_selector is not None:
-        selected_gain_channel = gain_selector(r0_waveforms)
-        r1_waveforms = r1_waveforms[
-            np.newaxis, selected_gain_channel, np.arange(n_pixels)
-        ]
-    else:
-        selected_gain_channel = None
+    return r1_waveforms
 
+
+def apply_gain_selection(r0_waveforms, r1_waveforms, gain_selector):
+    """
+    Apply gain selection to the R1 waveform.
+
+    Parameters
+    ----------
+    r0_waveforms : ndarray
+        Raw ADC waveforms from a simtel file. All gain channels available.
+        Shape: (n_channels, n_pixels, n_samples)
+    r1_waveforms : ndarray
+        Calibrated waveforms not gain-selected, in units of photoelectrons (p.e.).
+        Shape: (n_channels, n_pixels, n_samples)
+    gain_selector : ctapipe.calib.camera.gainselection.GainSelector
+        The GainSelector to use for selecting the gain channel.
+
+    Returns
+    -------
+    r1_waveforms : ndarray
+        Calibrated waveforms after gain selection, in units of photoelectrons (p.e.).
+        Shape: (1, n_pixels, n_samples)
+    selected_gain_channel : ndarray
+        The gain channel selected for each pixel. Shape: (n_pixels,)
+    """
+    selected_gain_channel = gain_selector(r0_waveforms)
+    r1_waveforms = r1_waveforms[
+        np.newaxis,
+        selected_gain_channel,
+        _get_pixel_index(r0_waveforms.shape[-2]),
+    ]
     return r1_waveforms, selected_gain_channel
 
 
@@ -1060,29 +1084,30 @@ def _generate_events(self):
                 factor = array_event["laser_calibrations"][tel_id]["calib"]
                 disabled_pixel_mask = self._get_disabled_pixel_mask(tel_id)
 
-                select_gain = self.select_gain is True or (
-                    self.select_gain is None
-                    and trigger.event_type is EventType.SUBARRAY
-                )
-                if select_gain:
-                    gain_selector = self.gain_selector
-                else:
-                    gain_selector = None
-
                 if self.skip_r1_calibration:
                     # Skip the simtel R1 calibration
                     r1_waveform = adc_samples.astype(np.float32)
-                    selected_gain_channel = None
                 else:
-                    # Apply the simtel R1 calibration and the gain selector if selected
-                    r1_waveform, selected_gain_channel = apply_simtel_r1_calibration(
+                    # Apply the simtel R1 calibration to get waveforms in units of photoelectrons (p.e.)
+                    r1_waveform = apply_simtel_r1_calibration(
                         adc_samples,
                         pedestal,
                         factor,
-                        gain_selector,
                         self.calib_scale,
                         self.calib_shift,
                     )
+                # Perform the gain selection if requested.
+                # By default, the cosmic events will be gain-selected, not for calibration events.
+                select_gain = self.select_gain is True or (
+                    self.select_gain is None
+                    and trigger.event_type is EventType.SUBARRAY
+                )
+                if select_gain:
+                    r1_waveform, selected_gain_channel = apply_gain_selection(
+                        adc_samples, r1_waveform, self.gain_selector
+                    )
+                else:
+                    selected_gain_channel = None
 
                 pixel_status = self._get_r1_pixel_status(
                     tel_id=tel_id,

src/ctapipe/io/tests/test_simteleventsource.py

@@ -19,6 +19,7 @@
 from ctapipe.io.simteleventsource import (
     AtmosphereProfileKind,
     SimTelEventSource,
+    apply_gain_selection,
     apply_simtel_r1_calibration,
     read_atmosphere_profile_from_simtel,
 )
@@ -270,6 +271,37 @@ def test_skip_r1_calibration():
     assert n_processed == n_expected
 
 
+def test_skip_r1_calibration_with_gain_selection():
+    n_expected = 1
+    with SimTelEventSource(
+        input_url=calib_events_path,
+        max_events=n_expected,
+        skip_calibration_events=False,
+        select_gain=True,
+        focal_length_choice="EQUIVALENT",
+    ) as reader:
+        n_processed = 0
+        for event in reader:
+            n_processed += 1
+            # R0 should have 2 channels (high gain and low gain)
+            assert event.r0.tel[1].waveform.ndim == 3
+            assert event.r0.tel[1].waveform.shape[0] == 2, (
+                f"R0 waveforms should have 2 channels, got {event.r0.tel[1].waveform.shape[0]}"
+            )
+            # R1 should have 1 channel after gain selection
+            assert event.r1.tel[1].waveform.ndim == 3
+            assert event.r1.tel[1].waveform.shape[0] == 1, (
+                f"R1 waveforms should have 1 channel after gain selection, got {event.r1.tel[1].waveform.shape[0]}"
+            )
+            # Check that selected_gain_channel exists and has correct shape
+            assert event.r1.tel[1].selected_gain_channel is not None
+            assert (
+                event.r1.tel[1].selected_gain_channel.shape[0]
+                == event.r1.tel[1].waveform.shape[1]
+            ), "selected_gain_channel should have one entry per pixel"
+    assert n_processed == n_expected
+
+
 def test_time_shift():
     source = SimTelEventSource(
         input_url=calib_events_path,
@@ -326,8 +358,9 @@ def test_apply_simtel_r1_calibration_1_channel():
     dc_to_pe = np.full((n_channels, n_pixels), 0.5)
 
     gain_selector = ThresholdGainSelector(threshold=90)
-    r1_waveforms, selected_gain_channel = apply_simtel_r1_calibration(
-        r0_waveforms, pedestal, dc_to_pe, gain_selector
+    r1_waveforms = apply_simtel_r1_calibration(r0_waveforms, pedestal, dc_to_pe)
+    r1_waveforms, selected_gain_channel = apply_gain_selection(
+        r0_waveforms, r1_waveforms, gain_selector
     )
 
     assert (selected_gain_channel == 0).all()
@@ -357,8 +390,9 @@ def test_apply_simtel_r1_calibration_2_channel():
     dc_to_pe[1] = 0.1
 
     gain_selector = ThresholdGainSelector(threshold=90)
-    r1_waveforms, selected_gain_channel = apply_simtel_r1_calibration(
-        r0_waveforms, pedestal, dc_to_pe, gain_selector
+    r1_waveforms = apply_simtel_r1_calibration(r0_waveforms, pedestal, dc_to_pe)
+    r1_waveforms, selected_gain_channel = apply_gain_selection(
+        r0_waveforms, r1_waveforms, gain_selector
     )
 
     assert selected_gain_channel[0] == 1