signal extractor SlidingWindowMaxSum

ctapipe/image/extractor.py

@@ -8,10 +8,12 @@
     "FixedWindowSum",
     "GlobalPeakWindowSum",
     "LocalPeakWindowSum",
+    "SlidingWindowMaxSum",
     "NeighborPeakWindowSum",
     "BaselineSubtractedNeighborPeakWindowSum",
     "TwoPassWindowSum",
     "extract_around_peak",
+    "extract_sliding_window",
     "neighbor_average_waveform",
     "subtract_baseline",
     "integration_correction",
@@ -119,6 +121,77 @@ def extract_around_peak(
     sum_[0] = i_sum
 
 
+@guvectorize(
+    [
+        (float64[:], int64, float64, float32[:], float32[:]),
+        (float32[:], int64, float64, float32[:], float32[:]),
+    ],
+    "(s),(),()->(),()",
+    nopython=True,
+)
+def extract_sliding_window(waveforms, width, sampling_rate_ghz, sum_, peak_time):
+    """
+    This function performs the following operations:
+
+    - Find the largest sum of width consecutive slices
+    - Obtain the pulse time within a window defined by a peak finding
+    algorithm, using the weighted average of the samples.
+
+    This function is a numpy universal function which defines the operation
+    applied on the waveform for every channel and pixel. Therefore in the
+    code body of this function:
+        - waveforms is a 1D array of size n_samples.
+        - width is integer
+
+    The ret argument is required by numpy to create the numpy array which is
+    returned. It can be ignored when calling this function.
+
+    Parameters
+    ----------
+    waveforms : ndarray
+        Waveforms stored in a numpy array.
+        Shape: (n_pix, n_samples)
+    width : ndarray or int
+        Window size of integration window for each pixel.
+    sampling_rate_ghz : float
+        Sampling rate of the camera, in units of GHz
+        Astropy units should have to_value('GHz') applied before being passed
+    sum_ : ndarray
+        Return argument for ufunc (ignore)
+        Returns the sum of the waveform samples
+    peak_time : ndarray
+        Return argument for ufunc (ignore)
+        Returns the peak_time in units "ns"
+
+    Returns
+    -------
+    charge : ndarray
+        Extracted charge.
+        Shape: (n_pix)
+
+    """
+
+    # first find the cumulative waveform
+    cwf = np.cumsum(waveforms)
+    # add zero at the begining so it is easier to substract the two arrays later
+    cwf = np.concatenate((np.zeros(1), cwf))
+    sums = cwf[width:] - cwf[:-width]
+    maxpos = np.argmax(sums)  # start of the window with largest sum
+    sum_[0] = sums[maxpos]
+
+    time_num = float64(0.0)
+    time_den = float64(0.0)
+    # now compute the timing as the average of non negative slices
+    for isample in prange(maxpos, maxpos + width):
+        if waveforms[isample] > 0:
+            time_num += waveforms[isample] * isample
+            time_den += waveforms[isample]
+
+    peak_time[0] = time_num / time_den if time_den > 0 else maxpos + 0.5 * width
+    # Convert to units of ns
+    peak_time[0] /= sampling_rate_ghz
+
+
 @njit(parallel=True, cache=True)
 def neighbor_average_waveform(waveforms, neighbors_indices, neighbors_indptr, lwt):
     """
@@ -524,6 +597,66 @@ def __call__(self, waveforms, telid, selected_gain_channel):
         return charge, peak_time
 
 
+class SlidingWindowMaxSum(ImageExtractor):
+    """
+    Sliding window extractor that maximizes the signal in window_width consecutive slices.
+    """
+
+    window_width = IntTelescopeParameter(
+        default_value=7, help="Define the width of the integration window"
+    ).tag(config=True)
+
+    apply_integration_correction = BoolTelescopeParameter(
+        default_value=True, help="Apply the integration window correction"
+    ).tag(config=True)
+
+    @lru_cache(maxsize=128)
+    def _calculate_correction(self, telid):
+        """
+        Calculate the correction for the extracted charge such that the value
+        returned would equal 1 for a noise-less unit pulse.
+
+        This method is decorated with @lru_cache to ensure it is only
+        calculated once per telescope.
+
+        WARNING: TO BE DONE properly, the current code reuses the function of
+        LocalPeakWindowSum assuming that the pulse is symmetric (which normally is not
+        true). The proper approach would be to apply a similar sliding window over the
+        reference pulse shape and compute the maximum fraction of the pulse contained
+        in window of a given size. Thus the correction for the leakage of signal outside
+        of the integration window is slightly overestimated.
+
+
+        Parameters
+        ----------
+        telid : int
+
+        Returns
+        -------
+        correction : ndarray
+        The correction to apply to an extracted charge using this ImageExtractor
+        Has size n_channels, as a different correction value might be required
+        for different gain channels.
+        """
+
+        readout = self.subarray.tel[telid].camera.readout
+        return integration_correction(
+            readout.reference_pulse_shape,
+            readout.reference_pulse_sample_width.to_value("ns"),
+            (1 / readout.sampling_rate).to_value("ns"),
+            self.window_width.tel[telid],
+            self.window_width.tel[telid] // 2,  # assuming that shift = 0.5 * window
+        )
+
+    def __call__(self, waveforms, telid, selected_gain_channel):
+        charge, peak_time = extract_sliding_window(
+            waveforms, self.window_width.tel[telid], self.sampling_rate[telid]
+        )
+        if self.apply_integration_correction.tel[telid]:
+            charge *= self._calculate_correction(telid=telid)[selected_gain_channel]
+        return charge, peak_time
+
+
 class NeighborPeakWindowSum(ImageExtractor):
     """
     Extractor which sums in a window about the

ctapipe/image/tests/test_extractor.py

@@ -8,6 +8,7 @@
 from ctapipe.core import non_abstract_children
 from ctapipe.image.extractor import (
     extract_around_peak,
+    extract_sliding_window,
     neighbor_average_waveform,
     subtract_baseline,
     integration_correction,
@@ -16,6 +17,7 @@
     NeighborPeakWindowSum,
     TwoPassWindowSum,
     FullWaveformSum,
+    SlidingWindowMaxSum,
 )
 from ctapipe.image.toymodel import WaveformModel
 from ctapipe.instrument import SubarrayDescription, TelescopeDescription
@@ -103,6 +105,26 @@ def test_extract_around_peak(toymodel):
     assert charge.dtype == np.float32
 
 
+def test_extract_sliding_window(toymodel):
+    waveforms, _, _, _, _, _ = toymodel
+    n_pixels, n_samples = waveforms.shape
+    charge, peak_time = extract_sliding_window(waveforms, 7, 1)
+    assert (charge >= 0).all()
+    assert (peak_time >= 0).all() and (peak_time <= n_samples).all()
+
+    x = np.arange(100)
+    y = norm.pdf(x, 41.2, 6)
+    charge, peak_time = extract_sliding_window(y[np.newaxis, :], x.size, 1)
+    assert_allclose(charge[0], 1.0, rtol=1e-3)
+    assert_allclose(peak_time[0], 41.2, rtol=1e-3)
+
+    # Test negative amplitude
+    y_offset = y - y.max() / 2
+    charge, _ = extract_sliding_window(y_offset[np.newaxis, :], x.size, 1)
+    assert_allclose(charge, y_offset.sum(), rtol=1e-3)
+    assert charge.dtype == np.float32
+
+
 def test_extract_around_peak_charge_expected(toymodel):
     waveforms = np.ones((2048, 96))
     n_samples = waveforms.shape[-1]
@@ -289,6 +311,15 @@ def test_fixed_window_sum(toymodel):
     assert_allclose(peak_time, true_time, rtol=0.1)
 
 
+def test_sliding_window_max_sum(toymodel):
+    waveforms, subarray, telid, selected_gain_channel, true_charge, true_time = toymodel
+    extractor = SlidingWindowMaxSum(subarray=subarray)
+    charge, peak_time = extractor(waveforms, telid, selected_gain_channel)
+    print(true_charge, charge, true_time, peak_time)
+    assert_allclose(charge, true_charge, rtol=0.1)
+    assert_allclose(peak_time, true_time, rtol=0.1)
+
+
 def test_neighbor_peak_window_sum_lwt(toymodel):
     waveforms, subarray, telid, selected_gain_channel, true_charge, true_time = toymodel
     extractor = NeighborPeakWindowSum(subarray=subarray, lwt=4)