Allow lookup of telescope parameters by type

ctapipe/core/tests/test_traits.py

@@ -291,6 +291,33 @@ def test_telescope_parameter_lookup(mock_subarray):
         telparam_list2[None]
 
 
+def test_telescope_parameter_lookup_by_type(subarray_prod5_paranal):
+    subarray = subarray_prod5_paranal.select_subarray([1, 2, 3, 4, 100, 101])
+
+    lookup = TelescopeParameterLookup([("type", "*", 10), ("type", "LST*", 100)])
+    lookup.attach_subarray(subarray)
+
+    assert lookup["LST_LST_LSTCam"] == 100
+    assert lookup["MST_MST_NectarCam"] == 10
+    assert lookup[subarray.tel[1]] == 100
+    assert lookup[subarray.tel[100]] == 10
+
+    # no global default
+    lookup = TelescopeParameterLookup([("type", "LST*", 100)])
+    lookup.attach_subarray(subarray)
+    assert lookup["LST_LST_LSTCam"] == 100
+
+    with pytest.raises(KeyError, match="no parameter value"):
+        assert lookup["MST_MST_NectarCam"]
+
+    with pytest.raises(ValueError, match="Unknown telescope"):
+        assert lookup["Foo"]
+
+    with pytest.raises(ValueError, match="Unknown telescope"):
+        # sst
+        assert lookup[subarray_prod5_paranal.tel[30]]
+
+
 def test_telescope_parameter_patterns(mock_subarray):
     """Test validation of TelescopeParameters"""
 

ctapipe/core/traits.py

@@ -6,9 +6,10 @@
 import pathlib
 from collections import UserList
 from fnmatch import fnmatch
-from typing import Optional
+from typing import Optional, Union
 from urllib.parse import urlparse
 
+import numpy as np
 import traitlets
 import traitlets.config
 from astropy.time import Time
@@ -388,11 +389,12 @@ def __init__(self, telescope_parameter_list):
         telescope_parameter_list : list
             List of tuples in the form `[(command, argument, value), ...]`
         """
-        # self._telescope_parameter_list = copy.deepcopy(telescope_parameter_list)
         self._telescope_parameter_list = copy.deepcopy(telescope_parameter_list)
         self._value_for_tel_id = None
+        self._value_for_type = None
         self._subarray = None
-        self._subarray_global_value = None
+        self._subarray_global_value = Undefined
+        self._type_strs = None
         for param in telescope_parameter_list:
             if param[1] == "*":
                 self._subarray_global_value = param[2]
@@ -410,49 +412,78 @@ def attach_subarray(self, subarray):
         """
         self._subarray = subarray
         self._value_for_tel_id = {}
+        self._value_for_type = {}
+        self._type_strs = {str(tel) for tel in self._subarray.telescope_types}
         for command, arg, value in self._telescope_parameter_list:
             if command == "type":
                 matched_tel_types = [
                     str(t) for t in subarray.telescope_types if fnmatch(str(t), arg)
                 ]
                 logger.debug(f"argument '{arg}' matched: {matched_tel_types}")
+
                 if len(matched_tel_types) == 0:
                     logger.warning(
                         "TelescopeParameter type argument '%s' did not match "
                         "any known telescope types",
                         arg,
                     )
+
                 for tel_type in matched_tel_types:
+                    self._value_for_type[tel_type] = value
                     for tel_id in subarray.get_tel_ids_for_type(tel_type):
                         self._value_for_tel_id[tel_id] = value
             elif command == "id":
                 self._value_for_tel_id[int(arg)] = value
             else:
                 raise ValueError(f"Unrecognized command: {command}")
 
-    def __getitem__(self, tel_id: Optional[int]):
+    def __getitem__(self, tel: Optional[Union[int, str]]):
         """
         Returns the resolved parameter for the given telescope id
         """
-        if tel_id is None:
-            if self._subarray_global_value is not None:
+        if tel is None:
+            if self._subarray_global_value is not Undefined:
                 return self._subarray_global_value
-            else:
-                raise KeyError("No subarray global value set for TelescopeParameter")
+
+            raise KeyError("No subarray global value set for TelescopeParameter")
+
         if self._value_for_tel_id is None:
             raise ValueError(
                 "TelescopeParameterLookup: No subarray attached, call "
                 "`attach_subarray` first before trying to access a value by tel_id"
             )
-        try:
-            return self._value_for_tel_id[tel_id]
-        except KeyError:
-            raise KeyError(
-                f"TelescopeParameterLookup: no "
-                f"parameter value was set for telescope with tel_id="
-                f"{tel_id}. Please set it explicitly, "
-                f"or by telescope type or '*'."
-            )
+
+        if isinstance(tel, (int, np.integer)):
+            try:
+                return self._value_for_tel_id[tel]
+            except KeyError:
+                raise KeyError(
+                    f"TelescopeParameterLookup: no "
+                    f"parameter value was set for telescope with tel_id="
+                    f"{tel}. Please set it explicitly, "
+                    f"or by telescope type or '*'."
+                )
+
+        from ctapipe.instrument import TelescopeDescription
+
+        if isinstance(tel, TelescopeDescription):
+            tel = str(tel)
+
+        if isinstance(tel, str):
+            if tel not in self._type_strs:
+                raise ValueError(
+                    f"Unknown telescope type {tel}, known: {self._type_strs}"
+                )
+            try:
+                return self._value_for_type[tel]
+            except KeyError:
+                raise KeyError(
+                    "TelescopeParameterLookup: no "
+                    "parameter value was set for telescope type"
+                    f" '{tel}'. Please set explicitly or using '*'."
+                )
+
+        raise TypeError(f"Unsupported lookup type: {type(tel)}")
 
 
 class TelescopeParameter(List):