Merge 7798779 into 06306aa

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,81 @@ 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.
+
+        The same procedure as for the actual SlidingWindowMaxSum extractor is used, but
+        on the reference pulse_shape (that is also more finely binned)
+
+        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
+
+        # compute the number of slices to integrate in the pulse template
+        width_shape = int(
+            round(
+                (
+                    self.window_width.tel[telid]
+                    / readout.sampling_rate
+                    / readout.reference_pulse_sample_width
+                )
+                .to("")
+                .value
+            )
+        )
+
+        n_channels = len(readout.reference_pulse_shape)
+        correction = np.ones(n_channels, dtype=np.float)
+        for ichannel, pulse_shape in enumerate(readout.reference_pulse_shape):
+
+            # apply the same method as sliding window to find the highest sum
+            cwf = np.cumsum(pulse_shape)
+            # 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_shape:] - cwf[:-width_shape]
+            maxsum = np.max(sums)
+            correction[ichannel] = np.sum(pulse_shape) / maxsum
+
+        return correction
+
+    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)

ctapipe/image/tests/test_sliding_window_correction.py

@@ -0,0 +1,58 @@
+"""
+Test of sliding window extractor for LST camera pulse shape with
+the correction for the integration window completeness
+"""
+
+import numpy as np
+import astropy.units as u
+from numpy.testing import assert_allclose
+from traitlets.config.loader import Config
+
+from ctapipe.image.extractor import SlidingWindowMaxSum, ImageExtractor
+from ctapipe.image.toymodel import WaveformModel
+from ctapipe.instrument import SubarrayDescription, TelescopeDescription
+
+
+def test_sw_pulse_lst():
+    """
+    Test function of sliding window extractor for LST camera pulse shape with
+    the correction for the integration window completeness
+    """
+
+    # prepare array with 1 LST
+    subarray = SubarrayDescription(
+        "LST1",
+        tel_positions={1: np.zeros(3) * u.m},
+        tel_descriptions={
+            1: TelescopeDescription.from_name(optics_name="LST", camera_name="LSTCam")
+        },
+    )
+
+    telid = list(subarray.tel.keys())[0]
+
+    n_pixels = subarray.tel[telid].camera.geometry.n_pixels
+    n_samples = 40
+    readout = subarray.tel[telid].camera.readout
+
+    random = np.random.RandomState(1)
+    min_charge = 100
+    max_charge = 1000
+    charge_true = random.uniform(min_charge, max_charge, n_pixels)
+    time_true = random.uniform(
+        n_samples // 2 - 1, n_samples // 2 + 1, n_pixels
+    ) / readout.sampling_rate.to_value(u.GHz)
+
+    waveform_model = WaveformModel.from_camera_readout(readout)
+    waveform = waveform_model.get_waveform(charge_true, time_true, n_samples)
+    selected_gain_channel = np.zeros(charge_true.size, dtype=np.int)
+
+    # define extractor
+    config = Config({"SlidingWindowMaxSum": {"window_width": 8}})
+    extractor = SlidingWindowMaxSum(subarray=subarray)
+    extractor = ImageExtractor.from_name(
+        "SlidingWindowMaxSum", subarray=subarray, config=config
+    )
+
+    charge, _ = extractor(waveform, telid, selected_gain_channel)
+    print(charge / charge_true)
+    assert_allclose(charge, charge_true, rtol=0.02)