Add a component to reduce waveform readout window

Author: maxnoe

src/ctapipe/image/reducer.py

@@ -3,6 +3,7 @@
 """
 
 from abc import abstractmethod
+from copy import deepcopy
 
 import numpy as np
 
@@ -214,3 +215,39 @@ def select_pixels(self, waveforms, tel_id=None, selected_gain_channel=None):
             mask = dilate(camera_geom, mask)
 
         return mask
+
+
+class ReadoutWindowReducer(TelescopeComponent):
+    """
+    Reduce the readout window size of telescope using a fixed window.
+    """
+
+    window_start = IntTelescopeParameter(
+        default_value=None,
+        allow_none=True,
+        help="Start sample of readout window",
+    ).tag(config=True)
+
+    window_end = IntTelescopeParameter(
+        default_value=None,
+        allow_none=True,
+        help="Last sample of readout window (non-inclusive)",
+    ).tag(config=True)
+
+    def __init__(self, subarray, **kwargs):
+        # we mutate the subarray, make sure we do not modify it for someone else
+        super().__init__(deepcopy(subarray), **kwargs)
+
+        for tel_id, tel in self.subarray.tel.items():
+            start = self.window_start.tel[tel_id]
+            end = self.window_end.tel[tel_id]
+            n_samples = len(np.arange(tel.camera.readout.n_samples)[start:end])
+            tel.camera.readout.n_samples = n_samples
+
+    def __call__(self, event):
+        for container in (event.r0, event.r1, event.dl0):
+            for tel_id, tel_container in container.tel.items():
+                if tel_container.waveform is not None:
+                    start = self.window_start.tel[tel_id]
+                    end = self.window_end.tel[tel_id]
+                    tel_container.waveform = tel_container.waveform[..., start:end]

src/ctapipe/image/tests/test_reducer.py

@@ -1,3 +1,5 @@
+from copy import deepcopy
+
 import astropy.units as u
 import numpy as np
 import pytest
@@ -6,6 +8,7 @@
 
 from ctapipe.image.reducer import NullDataVolumeReducer, TailCutsDataVolumeReducer
 from ctapipe.instrument import SubarrayDescription
+from ctapipe.io import EventSource
 
 
 @pytest.fixture(scope="module")
@@ -92,3 +95,72 @@ def test_tailcuts_data_volume_reducer(subarray_lst):
 
     assert (reduced_waveforms != 0).sum() == (1 + 4 + 14) * n_samples
     assert_array_equal(expected_waveforms, reduced_waveforms)
+
+
+def test_readout_window_reducer(prod5_gamma_simtel_path):
+    from ctapipe.image.reducer import ReadoutWindowReducer
+
+    def check(subarray, event, event_reduced, event_no_op):
+        checked = 0
+        for dl in ("r0", "r1", "dl0"):
+            for tel_id in getattr(event, dl).tel.keys():
+                if str(subarray.tel[tel_id]).startswith("LST"):
+                    start = 10
+                    end = 20
+                    n_samples = end - start
+                elif "Nectar" in str(subarray.tel[tel_id]):
+                    start = 15
+                    end = 40
+                    n_samples = end - start
+                else:
+                    start = None
+                    end = None
+                    n_samples = subarray.tel[tel_id].camera.readout.n_samples
+
+                original_waveform = getattr(event, dl).tel[tel_id].waveform
+                reduced_waveform = getattr(event_reduced, dl).tel[tel_id].waveform
+                no_op_waveform = getattr(event_no_op, dl).tel[tel_id].waveform
+
+                assert reduced_waveform.ndim == 3
+                assert reduced_waveform.shape[-1] == n_samples
+                np.testing.assert_array_equal(original_waveform, no_op_waveform)
+                np.testing.assert_array_equal(
+                    original_waveform[..., start:end], reduced_waveform
+                )
+
+                checked += 1
+        return checked
+
+    with EventSource(prod5_gamma_simtel_path) as source:
+        reducer_no_op = ReadoutWindowReducer(source.subarray)
+        reducer = ReadoutWindowReducer(
+            source.subarray,
+            window_start=[
+                ("type", "*", None),
+                ("type", "LST*", 10),
+                ("type", "*Nectar*", 15),
+            ],
+            window_end=[
+                ("type", "*", None),
+                ("type", "LST*", 20),
+                ("type", "*Nectar*", 40),
+            ],
+        )
+
+        # make sure we didn't modify the original subarray
+        assert source.subarray.tel[1].camera.readout.n_samples == 40
+        assert reducer_no_op.subarray.tel[1].camera.readout.n_samples == 40
+        # new subarray should have reduced window
+        assert reducer.subarray.tel[1].camera.readout.n_samples == 10
+
+        n_checked = 0
+        for event in source:
+            event_no_op = deepcopy(event)
+            event_reduced = deepcopy(event)
+
+            reducer(event_reduced)
+            reducer_no_op(event_no_op)
+
+            n_checked += check(source.subarray, event, event_reduced, event_no_op)
+
+        assert n_checked > 0