Add arg_n_largest function

 - substitute it with the argsort term in the GlobalPeakWindowSum

 Co-authored-by: Maximilian Linhoff <[email protected]>

src/ctapipe/image/extractor.py

@@ -47,6 +47,7 @@
 from .hillas import camera_to_shower_coordinates, hillas_parameters
 from .invalid_pixels import InvalidPixelHandler
 from .morphology import brightest_island, number_of_islands
+from .statistics import arg_n_largest
 from .timing import timing_parameters
 
 
@@ -585,11 +586,13 @@ def __call__(
             ).argmax(axis=-1)
         else:
             n_pixels = int(self.pixel_fraction.tel[tel_id] * waveforms.shape[-2])
-            brightest = np.argsort(
-                waveforms.max(
-                    axis=-1, where=~broken_pixels[..., np.newaxis], initial=-np.inf
-                )
-            )[..., -n_pixels:]
+            n_channels = waveforms.shape[0]
+            brightest = np.empty((n_channels, n_pixels), dtype=np.int64)
+            waveforms_max = waveforms.max(
+                axis=-1, where=~broken_pixels[..., np.newaxis], initial=-np.inf
+            )
+            for channel in range(n_channels):
+                brightest[channel] = arg_n_largest(n_pixels, waveforms_max[channel])
 
             # average over brightest pixels then argmax over samples
             peak_index = (

src/ctapipe/image/statistics.py

@@ -100,3 +100,16 @@ def descriptive_statistics(
 def n_largest(n, array):
     """return the n largest values of an array"""
     return nlargest(n, array)
+
+
+@njit
+def arg_n_largest(n, array):
+    """return the indices of the n largest values of an array"""
+    values = []
+    for i in range(len(array)):
+        values.append((array[i], i))
+    largest = n_largest(n, values)
+    idx = np.empty(n, dtype=np.int64)
+    for i in range(n):
+        idx[i] = largest[i][1]
+    return idx

src/ctapipe/image/tests/test_statistics.py

@@ -62,11 +62,23 @@ def test_return_type():
     assert isinstance(stats, PeakTimeStatisticsContainer)
 
 
-def test_nlargest():
+def test_n_largest():
     from ctapipe.image.statistics import n_largest
 
-    image = np.random.rand(1855)
+    rng = np.random.default_rng(0)
+    image = rng.random(1855)
     image[-3:] = 10
 
     largest_3 = n_largest(3, image)
     assert largest_3 == [10, 10, 10]
+
+
+def test_arg_n_largest():
+    from ctapipe.image.statistics import arg_n_largest
+
+    rng = np.random.default_rng(0)
+    image = rng.random(1855)
+    image[-3:] = 10
+
+    largest_3 = arg_n_largest(3, image)
+    assert (largest_3 == [1854, 1853, 1852]).all()