WIP: Add gain selection to r1 calibration

ctapipe/calib/camera/calibrator.py

@@ -32,8 +32,7 @@ class CameraCalibrator(Component):
         extractor_t0='ChargeExtractorFactory.t0',
         window_width='ChargeExtractorFactory.window_width',
         window_shift='ChargeExtractorFactory.window_shift',
-        sig_amp_cut_HG='ChargeExtractorFactory.sig_amp_cut_HG',
-        sig_amp_cut_LG='ChargeExtractorFactory.sig_amp_cut_LG',
+        peak_detection_threshold='ChargeExtractorFactory.peak_detection_threshold',
         lwt='ChargeExtractorFactory.lwt',
         clip_amplitude='CameraDL1Calibrator.clip_amplitude',
         radius='CameraDL1Calibrator.radius',

ctapipe/calib/camera/dl0.py

@@ -33,6 +33,7 @@ class CameraDL0Reducer(Component):
         will equal the r1 samples.
     kwargs
     """
+
     def __init__(self, config=None, tool=None, reductor=None, **kwargs):
         super().__init__(config=config, parent=tool, **kwargs)
         if reductor is None:
@@ -82,10 +83,16 @@ def reduce(self, event):
         """
         tels = event.r1.tels_with_data
         for telid in tels:
-            r1 = event.r1.tel[telid].waveform
             if self.check_r1_exists(event, telid):
+                dl0_tel = event.dl0.tel[telid]
+                r1_tel = event.r1.tel[telid]
+                r1_waveform = r1_tel.waveform
                 if self._reductor is None:
-                    event.dl0.tel[telid].waveform = r1
+                    dl0_tel.waveform = r1_waveform
                 else:
-                    reduction = self._reductor.reduce_waveforms(r1)
-                    event.dl0.tel[telid].waveform = reduction
+                    reduction = self._reductor.reduce_waveforms(r1_waveform)
+                    dl0_tel.waveform = reduction
+
+                dl0_tel.gain_channel = r1_tel.gain_channel
+                dl0_tel.trigger_time = r1_tel.trigger_time
+                dl0_tel.trigger_type = r1_tel.trigger_type

ctapipe/calib/camera/dl1.py

@@ -194,13 +194,14 @@ def calibrate(self, event):
         for telid in event.dl0.tels_with_data:
 
             if self.check_dl0_exists(event, telid):
+                dl0_tel = event.dl0.tel[telid]
                 waveforms = event.dl0.tel[telid].waveform
-                n_samples = waveforms.shape[2]
+                n_samples = waveforms.shape[1]
                 if n_samples == 1:
                     # To handle ASTRI and dst
-                    corrected = waveforms[..., 0]
+                    corrected = waveforms[:, 0]
                     window = np.ones(waveforms.shape)
-                    peakpos = np.zeros(waveforms.shape[0:2])
+                    peakpos = np.zeros(waveforms.shape[0:1])
                     cleaned = waveforms
                 else:
                     # Clean waveforms
@@ -215,7 +216,7 @@ def calibrate(self, event):
                     charge, peakpos, window = extract(cleaned)
 
                     # Apply integration correction
-                    correction = self.get_correction(event, telid)[:, None]
+                    correction = self.get_correction(event, telid)[0]
                     corrected = charge * correction
 
                 # Clip amplitude
@@ -227,4 +228,5 @@ def calibrate(self, event):
                 event.dl1.tel[telid].image = corrected
                 event.dl1.tel[telid].extracted_samples = window
                 event.dl1.tel[telid].peakpos = peakpos
-                event.dl1.tel[telid].cleaned = cleaned
+                event.dl1.tel[telid].waveform = cleaned
+                event.dl1.tel[telid].gain_channel = dl0_tel.gain_channel

ctapipe/calib/camera/gainselection.py

@@ -37,7 +37,6 @@ def pick_gain_channel(waveforms, threshold, select_by_sample=False):
 
     # if we have 2 channels:
     if waveforms.shape[0] == 2:
-        waveforms = np.squeeze(waveforms)
         new_waveforms = waveforms[0].copy()
 
         if select_by_sample:
@@ -51,12 +50,12 @@ def pick_gain_channel(waveforms, threshold, select_by_sample=False):
             new_waveforms[gain_mask] = waveforms[1][gain_mask]
 
     elif waveforms.shape[0] == 1:
-        new_waveforms = np.squeeze(waveforms)
+        new_waveforms = np.squeeze(waveforms, axis=0)
         gain_mask = np.zeros_like(new_waveforms).astype(bool)
 
     else:
-        raise ValueError("input waveforms has shape %s. not sure what to do "
-                         "with that.", waveforms.shape)
+        raise ValueError("input waveforms has shape {}. not sure what to do "
+                         "with that.".format(waveforms.shape))
 
     return new_waveforms, gain_mask
 
@@ -67,6 +66,9 @@ class GainSelector(Component, metaclass=ABCMeta):
     single waveform.
     """
 
+    def __init__(self, config=None, parent=None, **kwargs):
+        super().__init__(config=config, parent=parent, **kwargs)
+
     @abstractclassmethod
     def select_gains(self, cam_id, multi_gain_waveform):
         """

ctapipe/calib/camera/r1.py

@@ -15,6 +15,7 @@
 """
 from abc import abstractmethod
 
+from .gainselection import ThresholdGainSelector
 from ...core import Component, Factory
 from ...core.traits import Unicode
 from ...io import EventSource
@@ -158,23 +159,31 @@ class HESSIOR1Calibrator(CameraR1Calibrator):
         Tool executable that is calling this component.
         Passes the correct logger to the component.
         Set to None if no Tool to pass.
+
     kwargs
     """
 
+    # CALIB_SCALE is only relevant for MC calibration.
+    #
+    # CALIB_SCALE is the factor needed to transform from mean p.e. units to
+    # units of the single-p.e. peak: Depends on the collection efficiency,
+    # the asymmetry of the single p.e. amplitude  distribution and the
+    # electronic noise added to the signals. Default value is for GCT.
+    #
+    # To correctly calibrate to number of photoelectron, a fresh SPE calibration
+    # should be applied using a SPE sim_telarray run with an
+    # artificial light source.
+    #
+    # TODO: Handle calib_scale differently per simlated telescope
+
     calib_scale = 1.05
-    """
-    CALIB_SCALE is only relevant for MC calibration.
 
-    CALIB_SCALE is the factor needed to transform from mean p.e. units to 
-    units of the single-p.e. peak: Depends on the collection efficiency, 
-    the asymmetry of the single p.e. amplitude  distribution and the 
-    electronic noise added to the signals. Default value is for GCT.
+    def __init__(self, config=None, tool=None, gain_selector=None, **kwargs):
+        super().__init__(config=config, tool=tool, **kwargs)
 
-    To correctly calibrate to number of photoelectron, a fresh SPE calibration
-    should be applied using a SPE sim_telarray run with an 
-    artificial light source.
-    """
-    # TODO: Handle calib_scale differently per simlated telescope
+        self.gain_selector = gain_selector
+        if self.gain_selector is None:
+            self.gain_selector = ThresholdGainSelector(config, tool)
 
     def calibrate(self, event):
         if event.meta['origin'] != 'hessio':
@@ -188,11 +197,16 @@ def calibrate(self, event):
                 ped = event.mc.tel[telid].pedestal / n_samples
                 gain = event.mc.tel[telid].dc_to_pe * self.calib_scale
                 calibrated = (samples - ped[..., None]) * gain[..., None]
-                event.r1.tel[telid].waveform = calibrated
 
+                cam_id = event.inst.subarray.tel[telid].camera.cam_id
+                waveform, mask = self.gain_selector.select_gains(cam_id,
+                                                                 calibrated)
+                event.r1.tel[telid].waveform_full = calibrated
+                event.r1.tel[telid].waveform = waveform
+                event.r1.tel[telid].gain_channel = mask
 
-class TargetIOR1Calibrator(CameraR1Calibrator):
 
+class TargetIOR1Calibrator(CameraR1Calibrator):
     pedestal_path = Unicode(
         '',
         allow_none=True,

ctapipe/calib/camera/tests/test_calibrator.py

@@ -1,5 +1,7 @@
 from copy import deepcopy
+
 from numpy.testing import assert_allclose
+
 from ctapipe.calib.camera import (
     CameraCalibrator,
     HESSIOR1Calibrator,
@@ -11,14 +13,14 @@
 
 
 def test_camera_calibrator(test_event):
-    event = deepcopy(test_event) # so we don't modify the test event
+    event = deepcopy(test_event)  # so we don't modify the test event
     telid = 11
 
     calibrator = CameraCalibrator(r1_product="HESSIOR1Calibrator")
 
     calibrator.calibrate(event)
     image = event.dl1.tel[telid].image
-    assert_allclose(image[0, 0], -2.216, 1e-3)
+    assert_allclose(image[0], -2.216, 1e-3)
 
 
 def test_manual_r1():

ctapipe/calib/camera/tests/test_dl0.py

@@ -19,7 +19,7 @@ def test_camera_dl0_reducer(test_event):
     reducer = CameraDL0Reducer()
     reducer.reduce(event)
     waveforms = event.dl0.tel[telid].waveform
-    assert_almost_equal(waveforms[0, 0, 0], -0.091, 3)
+    assert_almost_equal(waveforms[0, 0], -0.091, 3)
 
 
 def test_check_r1_exists(test_event):

ctapipe/calib/camera/tests/test_dl1.py

@@ -41,14 +41,20 @@ def test_camera_dl1_calibrator(test_event):
 
     calibrator.calibrate(event)
     image = event.dl1.tel[telid].image
-    assert_allclose(image[0, 0], -2.216, 1e-3)
+    waveform = event.dl1.tel[telid].waveform
+    assert_allclose(image[0], -2.216, 1e-3)
+    assert image.ndim == 1
+    assert len(image) > 1  # make sure image has more than one pixel
+    assert waveform.ndim == 2
+    assert event.dl1.tel[telid].gain_channel.shape == waveform.shape
+    assert event.dl1.tel[telid]
 
 
 def test_check_dl0_exists(test_event):
     event = deepcopy(test_event)
     telid = 11
     previous_calibration(event)
     calibrator = CameraDL1Calibrator()
-    assert(calibrator.check_dl0_exists(event, telid) is True)
+    assert calibrator.check_dl0_exists(event, telid) is True
     event.dl0.tel[telid].waveform = None
-    assert(calibrator.check_dl0_exists(event, telid) is False)
+    assert calibrator.check_dl0_exists(event, telid) is False

ctapipe/calib/camera/tests/test_gainselection.py

@@ -1,8 +1,8 @@
 import numpy as np
 import pytest
 
-from ctapipe.calib.camera.gainselection import ThresholdGainSelector
 from ctapipe.calib.camera.gainselection import SimpleGainSelector
+from ctapipe.calib.camera.gainselection import ThresholdGainSelector
 from ctapipe.calib.camera.gainselection import pick_gain_channel
 
 
@@ -97,11 +97,23 @@ def test_threshold_gain_selector():
     with pytest.raises(ValueError):
         selector.select_gains("NectarCam", np.ones((3, 1000, 30)))
 
-    # 1-gain channel input:
-    wf0 = np.ones((1, 1000, 1))
+    # 1-gain channel input with many samples:
+    wf0 = np.ones((1, 1000, 30))
+    wf1, gm = selector.select_gains("ASTRICam", wf0)
+    assert wf1.shape == (1000, 30)
+    assert gm.shape == (1000, 30)
+
+    # 2-gain channel input with no samples:
+    wf0 = np.random.uniform(10, size=(2, 2368, 1))
+    wf1, gm = selector.select_gains("ASTRICam", wf0)
+    assert wf1.shape == (2368, 1)
+    assert gm.shape == (2368, 1)
+
+    # 1-gain channel input with no samples:
+    wf0 = np.random.uniform(10, size=(1, 2368, 1))
     wf1, gm = selector.select_gains("ASTRICam", wf0)
-    assert wf1.shape == (1000,)
-    assert gm.shape == (1000,)
+    assert wf1.shape == (2368, 1)
+    assert gm.shape == (2368, 1)
 
 
 def test_simple_gain_selector():

ctapipe/calib/camera/tests/test_r1.py

@@ -20,7 +20,7 @@ def test_hessio_r1_calibrator(test_event):
     calibrator = HESSIOR1Calibrator()
     calibrator.calibrate(event)
     r1 = event.r1.tel[telid].waveform
-    assert_almost_equal(r1[0, 0, 0], -0.091, 3)
+    assert_almost_equal(r1[0, 0], -0.091, 3)
 
 
 def test_null_r1_calibrator(test_event):