Add another cleaning method in cleaning.py

ctapipe/image/cleaning.py

@@ -18,6 +18,8 @@
     "mars_cleaning_1st_pass",
     "fact_image_cleaning",
     "apply_time_delta_cleaning",
+    "apply_time_average_cleaning",
+    "time_constrained_clean",
     "ImageCleaner",
     "TailcutsImageCleaner",
 ]
@@ -32,7 +34,7 @@
     IntTelescopeParameter,
     BoolTelescopeParameter,
 )
-
+from .morphology import number_of_islands, brightest_island
 
 def tailcuts_clean(
     geom,
@@ -238,6 +240,57 @@ def apply_time_delta_cleaning(
     return pixels_to_keep
 
 
+def apply_time_average_cleaning(
+    geom, mask, image, arrival_times, picture_thresh, time_limit
+):
+    """
+    Extract all pixels that arrived within a given timeframe 
+    with respect to the time average of the pixels on the main island.
+
+    In order to avoid removing signal pixels of large impact-parameter events, 
+    the time limit for bright pixels is doubled.
+
+    Parameters
+    ----------
+    geom: `ctapipe.instrument.CameraGeometry`
+        Camera geometry information
+    mask: array, boolean
+        boolean mask of *clean* pixels before time_delta_cleaning
+    image: array
+        pixel values
+    arrival_times: array
+        pixel timing information
+    picture_thresh: float
+        threshold above which time limit is extended twice its value
+    time_limit: int or float
+        arrival time limit w.r.t. the average time of the core pixels
+
+    Returns
+    -------
+
+    A boolean mask of *clean* pixels.  To get a zero-suppressed image and pixel
+    list, use `image[mask], geom.pix_id[mask]`, or to keep the same
+    image size and just set unclean pixels to 0 or similar, use
+    `image[~mask] = 0`
+
+    """
+    pixels_to_remove = []
+    mask = mask.copy()
+    if np.count_nonzero(mask) > 0:
+
+        # use main island (maximum charge) for time average calculation
+        num_islands, island_labels = number_of_islands(geom, mask)
+        mask_main = brightest_island(num_islands, island_labels, image)
+        time_ave = np.average(arrival_times[mask_main], weights=image[mask_main]**2)
+
+        time_diffs = np.abs(arrival_times[mask] - time_ave)
+        time_limit_pixwise = np.where(image < (2 * picture_thresh), time_limit, time_limit * 2)[mask]
+
+        mask[mask] &= time_diffs < time_limit_pixwise
+
+    return mask
+
+
 def fact_image_cleaning(
     geom,
     image,
@@ -323,6 +376,88 @@ def fact_image_cleaning(
     return pixels_to_keep
 
 
+def time_constrained_clean(
+    geom,
+    image,
+    arrival_times,
+    picture_thresh=7,
+    boundary_thresh=5,
+    time_limit_core=4.5,
+    time_limit_boundary=1.5,
+    min_number_picture_neighbors=1,
+):
+    """
+    time constrained cleaning by MAGIC
+
+    Cleaning contains the following steps:
+    - Find core pixels (containing more photons than a picture threshold)
+    - Remove pixels with less than N neighbors
+    - Remove pixels whose arrival times are within a time limit of the average time
+    - Find boundary pixels (containing more photons than a boundary threshold)
+    - Remove pixels with less than N neighbors arriving within a given timeframe
+
+    Parameters
+    ----------
+    geom: `ctapipe.instrument.CameraGeometry`
+        Camera geometry information
+    image: array
+        pixel values
+    arrival_times: array
+        pixel timing information
+    picture_threshold: float or array
+        threshold above which all pixels are retained
+    boundary_threshold: float or array
+        threshold above which pixels are retained if they have a neighbor
+        already above the picture_thresh
+    time_limit_core: int or float
+        arrival time limit of core pixels w.r.t the average time
+    time_limit_boundary: int or float
+        arrival time limit of boundary pixels w.r.t neighboring core pixels
+    min_number_neighbors: int
+        Threshold to determine if a pixel survives cleaning steps.
+        These steps include checks of neighbor arrival time and value
+
+    Returns
+    -------
+    A boolean mask of *clean* pixels.  To get a zero-suppressed image and pixel
+    list, use `image[mask], geom.pix_id[mask]`, or to keep the same
+    image size and just set unclean pixels to 0 or similar, use
+    `image[~mask] = 0`
+
+    """
+
+    # find core pixels that pass a picture threshold                                                                                                                                                      
+    pixels_above_picture = image >= picture_thresh
+
+    # require at least min_number_picture_neighbors                                                                                                                                                       
+    number_of_neighbors_above_picture = geom.neighbor_matrix_sparse.dot(
+        pixels_above_picture.view(np.byte)
+    )
+    pixels_in_picture = pixels_above_picture & (
+        number_of_neighbors_above_picture >= min_number_picture_neighbors
+    )
+
+    # keep core pixels whose arrival times are within a certain time limit of the average                                                                                                                 
+    mask_core = apply_time_average_cleaning(geom, pixels_in_picture, image, arrival_times, picture_thresh, time_limit_core)
+
+    # find boundary pixels that pass a boundary threshold                                                                                                                                                 
+    pixels_above_boundary = image >= boundary_thresh
+    pixels_with_picture_neighbors = geom.neighbor_matrix_sparse.dot(mask_core)
+    mask_boundary = ( pixels_above_boundary & pixels_with_picture_neighbors ) & np.invert(mask_core)
+
+    # keep boundary pixels whose arrival times are within a certain time limit of the neighboring core pixels                                                                                             
+    pixels_to_remove = []
+    mask_boundary = mask_boundary.copy()
+
+    time_diffs = np.abs(arrival_times[mask_boundary, None] - arrival_times)
+    valid_neighbors = (time_diffs < time_limit_boundary) & geom.neighbor_matrix[mask_boundary] & mask_core
+    enough_neighbors = np.count_nonzero(valid_neighbors, axis=1) >= min_number_picture_neighbors
+    mask_boundary[mask_boundary] &= enough_neighbors
+
+    return mask_core | mask_boundary
+
+
+
 class ImageCleaner(TelescopeComponent):
     """
     Abstract class for all configurable Image Cleaning algorithms.   Use
@@ -440,3 +575,34 @@ def __call__(
             min_number_neighbors=self.min_picture_neighbors.tel[tel_id],
             time_limit=self.time_limit_ns.tel[tel_id],
         )
+
+
+class TimeConstrainedImageCleaner(TailcutsImageCleaner):
+    """
+    MAGIC-like Image cleaner with timing information (See `ctapipe.image.time_constrained_clean`)
+    """
+
+    time_limit_core_ns = FloatTelescopeParameter(
+        default_value=4.5, help="arrival time limit for neighboring " "core pixels, in ns"
+    ).tag(config=True)
+    time_limit_boundary_ns = FloatTelescopeParameter(
+        default_value=1.5, help="arrival time limit for neighboring " "boundary pixels, in ns"
+    ).tag(config=True)
+
+    def __call__(
+        self, tel_id: int, image: np.ndarray, arrival_times=None
+    ) -> np.ndarray:
+        """
+        Apply MAGIC-like image cleaning with timing information. See `ImageCleaner.__call__()`
+        """
+
+        return mars_cleaning_1st_pass(
+            self.subarray.tel[tel_id].camera.geometry,
+            image,
+            arrival_times=arrival_times,
+            picture_thresh=self.picture_threshold_pe.tel[tel_id],
+            boundary_thresh=self.boundary_threshold_pe.tel[tel_id],
+            min_number_picture_neighbors=self.min_picture_neighbors.tel[tel_id],
+            time_limit_core=self.time_limit_core_ns.tel[tel_id],
+            time_limit_boundary=self.time_limit_boundary_ns.tel[tel_id]
+        )

ctapipe/image/tests/test_cleaning.py

@@ -1,3 +1,4 @@
+from tkinter import Image
 import numpy as np
 from numpy.testing import assert_allclose
 from ctapipe.image import cleaning
@@ -331,3 +332,78 @@ def test_apply_time_delta_cleaning():
     test_mask = mask.copy()
     test_mask[[41, 157]] = 0
     assert (test_mask == td_mask).all()
+
+
+def test_time_constrained_clean():
+    geom = CameraGeometry.from_name("LSTCam")
+    charge = np.zeros(geom.n_pixels, dtype=np.float64)
+    peak_time = np.zeros(geom.n_pixels, dtype=np.float64)
+
+    # define signal pixels and their charges/timings (1 core pixel + 6 neighboring core pixels + 12 neighboring boundary pixels)
+    core_pixel = 100
+    core_neighbors = geom.neighbors[core_pixel]
+    boundary_pixels = np.setdiff1d(np.array([geom.neighbors[core_neighbor] for core_neighbor in core_neighbors]), np.append(core_neighbors, core_pixel))
+    charge[core_pixel], charge[core_neighbors], charge[boundary_pixels] = 15, 10, 6
+    peak_time[core_pixel], peak_time[core_neighbors], peak_time[boundary_pixels]  = 18, 20, 21
+
+    # define initial cleaning parameters
+    picture_thresh, boundary_thresh = 8, 4
+    time_limit_core, time_limit_boundary = 4.5, 1.5
+    min_number_picture_neighbors = 1
+
+    mask_signal = charge > 0
+
+    # 1. basic test
+    mask_reco = cleaning.time_constrained_clean(
+        geom, charge, peak_time, picture_thresh=picture_thresh, boundary_thresh=boundary_thresh, time_limit_core=time_limit_core, time_limit_boundary=time_limit_boundary, min_number_picture_neighbors=min_number_picture_neighbors,
+    )
+    test_mask = mask_signal.copy()
+    assert (test_mask == mask_reco).all()
+
+    # 2. increased min_number_picture_neighbors test (here 3)
+    min_number_picture_neighbors = 3
+    mask_reco = cleaning.time_constrained_clean(
+        geom, charge, peak_time, picture_thresh=picture_thresh, boundary_thresh=boundary_thresh, time_limit_core=time_limit_core, time_limit_boundary=time_limit_boundary, min_number_picture_neighbors=min_number_picture_neighbors,
+    )
+    ## removed pixels : boundary pixels
+    test_mask = mask_signal.copy()
+    test_mask[boundary_pixels] = 0
+    assert (test_mask == mask_reco).all()
+
+    # 3. strict time_limit_boundary test (here 0.5)
+    min_number_picture_neighbors = 1
+    time_limit_boundary = 0.5
+    mask_reco = cleaning.time_constrained_clean(
+        geom, charge, peak_time, picture_thresh=picture_thresh, boundary_thresh=boundary_thresh, time_limit_core=time_limit_core, time_limit_boundary=time_limit_boundary, min_number_picture_neighbors=min_number_picture_neighbors,
+    )
+    ## removed pixels : boundary pixels
+    test_mask = mask_signal.copy()
+    test_mask[boundary_pixels] = 0
+    assert (test_mask == mask_reco).all()
+
+    # 4. time_limit_core test (one of core_neighbors have peak time >5 slice away from the average)
+    time_limit_boundary = 1.5
+    noise_core_neighbor = core_neighbors[0]
+    peak_time[noise_core_neighbor] = 25
+    mask_reco = cleaning.time_constrained_clean(
+        geom, charge, peak_time, picture_thresh=picture_thresh, boundary_thresh=boundary_thresh, time_limit_core=time_limit_core, time_limit_boundary=time_limit_boundary, min_number_picture_neighbors=min_number_picture_neighbors,
+    )
+    ## removed pixels : the noise core neighbor pixel + one neighboring boundary
+    test_mask = mask_signal.copy()
+    test_mask[noise_core_neighbor] = 0
+    noise_boundary = np.setdiff1d(geom.neighbors[noise_core_neighbor], np.array([geom.neighbors[core_neighbor] for core_neighbor in core_neighbors[1:]]))
+    test_mask[noise_boundary] = 0
+    assert (test_mask == mask_reco).all()
+
+    # 5. time_limit_core test for brighter pixels (one of core_neighbors have peak time >5 slice away from the average)
+    charge[core_pixel], charge[core_neighbors] = 30, 20
+    mask_reco = cleaning.time_constrained_clean(
+        geom, charge, peak_time, picture_thresh=picture_thresh, boundary_thresh=boundary_thresh, time_limit_core=time_limit_core, time_limit_boundary=time_limit_boundary, min_number_picture_neighbors=min_number_picture_neighbors,
+    )
+    ## removed pixels : one neighboring boundary to the noise core pixel
+    test_mask = mask_signal.copy()
+    noise_boundary = np.setdiff1d(geom.neighbors[noise_core_neighbor], np.array([geom.neighbors[core_neighbor] for core_neighbor in core_neighbors[1:]]))
+    test_mask[noise_boundary] = 0
+    assert (test_mask == mask_reco).all()
+
+