Repository: cta-observatory/ctapipe
Pull request #2027 (2022-07-18 17:54:26)

Return tuples from SubArray properties, add suport for single value in `get_tel_ids`
The properties of `SubarrayDescription` are immutable, so the proper type is a tuple, not a list.

This also enables `get_tel_ids(<single value>)`.

maxnoe: @kosack @LukasNickel had to resolve conflicts

Patch:
diff --git a/ctapipe/instrument/subarray.py b/ctapipe/instrument/subarray.py
index 5aa0bdd6..449c6e65 100644
--- a/ctapipe/instrument/subarray.py
+++ b/ctapipe/instrument/subarray.py
@@ -5,7 +5,7 @@ import warnings
 from contextlib import ExitStack
 from copy import copy
 from itertools import groupby
-from typing import Dict, List, Union
+from typing import Dict, Iterable, Tuple, Union
 
 import numpy as np
 import tables
@@ -392,40 +392,46 @@ class SubarrayDescription:
         plt.tight_layout()
 
     @lazyproperty
-    def telescope_types(self) -> List[TelescopeDescription]:
+    def telescope_types(self) -> Tuple[TelescopeDescription]:
         """list of telescope types in the array"""
-        return list({tel for tel in self.tel.values()})
+        return tuple({tel for tel in self.tel.values()})
 
     @lazyproperty
-    def camera_types(self) -> List[CameraDescription]:
+    def camera_types(self) -> Tuple[CameraDescription]:
         """list of camera types in the array"""
-        return list({tel.camera for tel in self.tel.values()})
+        return tuple({tel.camera for tel in self.tel.values()})
 
     @lazyproperty
-    def optics_types(self) -> List[OpticsDescription]:
+    def optics_types(self) -> Tuple[OpticsDescription]:
         """list of optics types in the array"""
-        return list({tel.optics for tel in self.tel.values()})
+        return tuple({tel.optics for tel in self.tel.values()})
 
-    def get_tel_ids_for_type(self, tel_type):
+    def get_tel_ids_for_type(self, tel_type) -> Tuple[int]:
         """
         return list of tel_ids that have the given tel_type
 
         Parameters
         ----------
         tel_type: str or TelescopeDescription
-           telescope type string (e.g. 'MST:NectarCam')
-
+           telescope type string (e.g. 'MST_MST_NectarCam')
         """
         if isinstance(tel_type, TelescopeDescription):
-            tel_str = str(tel_type)
+            if tel_type not in self.telescope_types:
+                raise ValueError(f"{tel_type} not in subarray: {self.telescope_types}")
+            return tuple(
+                tel_id for tel_id, descr in self.tels.items() if descr == tel_type
+            )
         else:
-            tel_str = tel_type
-
-        return [id for id, descr in self.tels.items() if str(descr) == tel_str]
+            valid = {str(tel) for tel in self.telescope_types}
+            if tel_type not in valid:
+                raise ValueError(f"{tel_type} not in subarray: {valid}")
+            return tuple(
+                tel_id for tel_id, descr in self.tels.items() if str(descr) == tel_type
+            )
 
     def get_tel_ids(
-        self, telescopes: List[Union[int, str, TelescopeDescription]]
-    ) -> List[int]:
+        self, telescopes: Iterable[Union[int, str, TelescopeDescription]]
+    ) -> Tuple[int]:
         """
         Convert a list of telescope ids and telescope descriptions to
         a list of unique telescope ids.
@@ -445,22 +451,21 @@ class SubarrayDescription:
         """
         ids = set()
 
-        valid_tel_types = {str(tel_type) for tel_type in self.telescope_types}
+        # support single telescope element
+        if isinstance(telescopes, (int, str, TelescopeDescription)):
+            telescopes = (telescopes,)
 
         for telescope in telescopes:
             if isinstance(telescope, (int, np.integer)):
-                if telescope not in self.tel_ids:
+                if telescope not in self.tel:
                     raise ValueError(
                         f"Telescope with tel_id={telescope} not in subarray."
                     )
                 ids.add(telescope)
+            else:
+                ids.update(self.get_tel_ids_for_type(telescope))
 
-            if isinstance(telescope, str) and telescope not in valid_tel_types:
-                raise ValueError("Invalid telescope type input.")
-
-            ids.update(self.get_tel_ids_for_type(telescope))
-
-        return sorted(ids)
+        return tuple(sorted(ids))
 
     def __eq__(self, other):
         if not isinstance(other, self.__class__):


Test patch:
diff --git a/ctapipe/instrument/tests/test_subarray.py b/ctapipe/instrument/tests/test_subarray.py
index b6907edf..e84584ba 100644
--- a/ctapipe/instrument/tests/test_subarray.py
+++ b/ctapipe/instrument/tests/test_subarray.py
@@ -16,7 +16,7 @@ from ctapipe.instrument import (
 )
 
 
-def example_subarray(tel_type, n_tels=10):
+def create_subarray(tel_type, n_tels=10):
     """generate a simple subarray for testing purposes"""
     rng = np.random.default_rng(0)
 
@@ -42,7 +42,7 @@ def example_subarray(tel_type, n_tels=10):
 def test_subarray_description(prod5_mst_nectarcam):
     """Test SubarrayDescription functionality"""
     n_tels = 10
-    sub = example_subarray(prod5_mst_nectarcam, n_tels)
+    sub = create_subarray(prod5_mst_nectarcam, n_tels)
     sub.peek()
 
     assert len(sub.telescope_types) == 1
@@ -119,6 +119,9 @@ def test_get_tel_ids_for_type(example_subarray):
         assert len(sub.get_tel_ids_for_type(teltype)) > 0
         assert len(sub.get_tel_ids_for_type(str(teltype))) > 0
 
+    with pytest.raises(ValueError):
+        sub.get_tel_ids_for_type("NON_EXISTENT_TYPE")
+
 
 def test_hdf(example_subarray, tmp_path):
     import tables
@@ -224,14 +227,28 @@ def test_get_tel_ids(example_subarray, prod3_astri):
     true_tel_ids = (
         subarray.get_tel_ids_for_type("MST_MST_FlashCam")
         + subarray.get_tel_ids_for_type(prod3_astri)
-        + [1, 2]
+        + (1, 2)
     )
 
     assert sorted(tel_ids) == sorted(true_tel_ids)
 
     # test invalid telescope type
     with pytest.raises(Exception):
-        tel_ids = subarray.get_tel_ids(["It's a-me, Mario!"])
+        subarray.get_tel_ids(["It's a-me, Mario!"])
+
+    # test single string
+    assert subarray.get_tel_ids("LST_LST_LSTCam") == (1, 2, 3, 4)
+
+    # test single id
+    assert subarray.get_tel_ids(1) == (1,)
+
+    # test invalid id
+    with pytest.raises(ValueError):
+        subarray.get_tel_ids(500)
+
+    # test invalid description
+    with pytest.raises(ValueError):
+        subarray.get_tel_ids("It's a-me, Mario!")
 
 
 def test_unknown_telescopes(example_subarray):