Adapt get_combinations

Author: Hckjs

src/ctapipe/reco/stereo_combination.py

@@ -613,7 +613,7 @@ def _combine_altaz(self, event):
                 ids.append(tel_id)
 
         if len(fov_lon_values) > 0:
-            index_tel_combs = get_combinations(range(len(ids)), self.n_tel_combinations)
+            index_tel_combs = get_combinations(len(ids), self.n_tel_combinations)
             fov_lons, fov_lats, comb_weights = calc_combs_min_distances_event(
                 index_tel_combs,
                 np.array(fov_lon_values),

src/ctapipe/reco/telescope_event_handling.py

@@ -175,26 +175,26 @@ def weighted_mean_std_ufunc(
 
 
 @lru_cache(maxsize=4096)
-def get_combinations(iterable, size):
+def get_combinations(array_length, comb_size):
     """
-    Generate all possible combinations of elements of a given `size` from
-    the given `iterable`.
+    Generate all possible index combinations of elements of a given `comb_size`
+    from an array with a given `array_length`.
 
     Uses ``itertools.combinations`` and caching to speed up repeated calls.
 
     Parameters
     ----------
-    iterable: iterable
-        Input iterable from which to generate combinations.
+    array_length: int
+        Length of the list or array to generate index combinations from.
     size : int
         The size of each combination.
 
     Returns
     -------
     np.ndarray
-        Array of combinations of the specified size.
+        Array of indey combinations of the specified size.
     """
-    return np.array(list(combinations(iterable, size)))
+    return np.array(list(combinations(range(array_length), comb_size)))
 
 
 @njit
@@ -470,9 +470,9 @@ def create_combs_array(max_multiplicity, k):
         - An array of all k-combinations for different multiplicities.
         - An array mapping each combination to its respective multiplicity.
     """
-    combs_array = get_combinations(range(k), k)
+    combs_array = get_combinations(k, k)
     for i in range(k + 1, max_multiplicity + 1):
-        combs = get_combinations(range(i), k)
+        combs = get_combinations(i, k)
         combs_array = np.concatenate([combs_array, combs])
 
     n_combs = _calc_n_combs(np.arange(k, max_multiplicity + 1), k)

src/ctapipe/reco/tests/test_telescope_event_handling.py

@@ -56,12 +56,13 @@ def test_get_combinations():
     from ctapipe.reco.telescope_event_handling import get_combinations
 
     tel_ids = [1, 2, 3]
+    comb_size = 2
     get_combinations.cache_clear()
-    combinations = get_combinations(tel_ids)
+    index_combinations = get_combinations(len(tel_ids), comb_size)
 
-    expected_combinations = np.array([[1, 2], [1, 3], [2, 3]])
+    expected_combinations = np.array([[0, 1], [0, 2], [1, 2]])
 
-    assert np.allclose(combinations, expected_combinations)
+    assert np.allclose(index_combinations, expected_combinations)
 
 
 def test_calc_combs_min_distances():