Add possibility to compute multiplicities using FeatureGenerator

ctapipe/core/feature_generator.py

@@ -1,6 +1,8 @@
 """
 Generate Features.
 """
+from collections import ChainMap
+
 from .component import Component
 from .expression_engine import ExpressionEngine
 from .traits import List, Tuple, Unicode
@@ -51,7 +53,7 @@ def __init__(self, config=None, parent=None, **kwargs):
         self.engine = ExpressionEngine(expressions=self.features)
         self._feature_names = [name for name, _ in self.features]
 
-    def __call__(self, table):
+    def __call__(self, table, **kwargs):
         """
         Apply feature generation to the input table.
 
@@ -60,8 +62,9 @@ def __call__(self, table):
         however the new columns won't be visible in the input table.
         """
         table = _shallow_copy_table(table)
+        lookup = ChainMap(table, kwargs)
 
-        for result, name in zip(self.engine(table), self._feature_names):
+        for result, name in zip(self.engine(lookup), self._feature_names):
             if name in table.colnames:
                 raise FeatureGeneratorException(f"{name} is already a column of table.")
             try:

ctapipe/core/tests/test_feature_generator.py

@@ -67,3 +67,37 @@ def test_to_unit():
     table = generator(table)
     assert table["length_meter"] == 1000
     assert table["log_length_meter"] == 3
+
+
+def test_multiplicity(subarray_prod5_paranal):
+    """Test if generating features works."""
+
+    expressions = [
+        ("n_telescopes", "subarray.multiplicity(tels_with_trigger)"),
+        ("n_lsts", "subarray.multiplicity(tels_with_trigger, 'LST_LST_LSTCam')"),
+        ("n_msts", "subarray.multiplicity(tels_with_trigger, 'MST_MST_FlashCam')"),
+        ("n_ssts", "subarray.multiplicity(tels_with_trigger, 'SST_ASTRI_CHEC')"),
+    ]
+
+    subarray = subarray_prod5_paranal.select_subarray([1, 2, 20, 21, 80, 81])
+
+    masks = np.array(
+        [
+            [True, False, True, True, False, False],
+            [True, True, False, True, False, True],
+        ]
+    )
+
+    table = Table(
+        {
+            "tels_with_trigger": masks,
+        }
+    )
+
+    generator = FeatureGenerator(features=expressions)
+    table = generator(table, subarray=subarray)
+
+    np.testing.assert_equal(table["n_telescopes"], [3, 4])
+    np.testing.assert_equal(table["n_lsts"], [1, 2])
+    np.testing.assert_equal(table["n_msts"], [2, 1])
+    np.testing.assert_equal(table["n_ssts"], [0, 1])

ctapipe/instrument/subarray.py

@@ -427,7 +427,7 @@ def optics_types(self) -> Tuple[OpticsDescription]:
 
     def get_tel_ids_for_type(self, tel_type) -> Tuple[int]:
         """
-        return list of tel_ids that have the given tel_type
+        return tuple of tel_ids that have the given tel_type
 
         Parameters
         ----------
@@ -448,6 +448,43 @@ def get_tel_ids_for_type(self, tel_type) -> Tuple[int]:
                 tel_id for tel_id, descr in self.tels.items() if str(descr) == tel_type
             )
 
+    def get_tel_indices_for_type(self, tel_type):
+        """
+        return tuple of telescope indices that have the given tel_type
+
+        Parameters
+        ----------
+        tel_type: str or TelescopeDescription
+           telescope type string (e.g. 'MST_MST_NectarCam')
+        """
+        return self.tel_ids_to_indices(self.get_tel_ids_for_type(tel_type))
+
+    def multiplicity(self, tel_mask, tel_type=None):
+        """
+        Compute multiplicity from a telescope mask
+
+        Parameters
+        ----------
+        tel_mask : np.ndarray
+            Boolean array with last dimension of size n_telescopes
+        tel_type : None, str or TelescopeDescription
+            If not given, compute multiplicity from all telescopes.
+            If given, the multiplicity of the given telescope type will
+            be computed.
+
+        Returns
+        -------
+        multiplicity : int or np.ndarray
+            Number of true values for given telescope mask and telescope type
+        """
+
+        if tel_type is None:
+            return np.count_nonzero(tel_mask, axis=-1)
+
+        return np.count_nonzero(
+            tel_mask[..., self.get_tel_indices_for_type(tel_type)], axis=-1
+        )
+
     def get_tel_ids(
         self, telescopes: Iterable[Union[int, str, TelescopeDescription]]
     ) -> Tuple[int]:

ctapipe/instrument/tests/test_subarray.py

@@ -246,3 +246,27 @@ def test_unknown_telescopes(example_subarray):
 
     with pytest.raises(UnknownTelescopeID):
         example_subarray.select_subarray([300, 201])
+
+
+def test_multiplicity(subarray_prod5_paranal):
+
+    subarray = subarray_prod5_paranal.select_subarray([1, 2, 20, 21, 80, 81])
+
+    mask = np.array([True, False, True, True, False, False])
+
+    assert subarray.multiplicity(mask) == 3
+    assert subarray.multiplicity(mask, "LST_LST_LSTCam") == 1
+    assert subarray.multiplicity(mask, "MST_MST_FlashCam") == 2
+    assert subarray.multiplicity(mask, "SST_ASTRI_CHEC") == 0
+
+    masks = np.array(
+        [
+            [True, False, True, True, False, False],
+            [True, True, False, True, False, True],
+        ]
+    )
+
+    np.testing.assert_equal(subarray.multiplicity(masks), [3, 4])
+    np.testing.assert_equal(subarray.multiplicity(masks, "LST_LST_LSTCam"), [1, 2])
+    np.testing.assert_equal(subarray.multiplicity(masks, "MST_MST_FlashCam"), [2, 1])
+    np.testing.assert_equal(subarray.multiplicity(masks, "SST_ASTRI_CHEC"), [0, 1])

ctapipe/reco/preprocessing.py

@@ -88,6 +88,8 @@ def collect_features(
     """
     features = {}
 
+    features.update(event.trigger.as_dict(recursive=False, flatten=True))
+
     features.update(
         event.dl1.tel[tel_id].parameters.as_dict(
             add_prefix=True,

ctapipe/reco/sklearn.py

@@ -370,7 +370,7 @@ def __call__(self, event: ArrayEventContainer) -> None:
         """
         for tel_id in event.trigger.tels_with_trigger:
             table = collect_features(event, tel_id, self.instrument_table)
-            table = self.feature_generator(table)
+            table = self.feature_generator(table, subarray=self.subarray)
 
             # get_table_mask returns a table with a single row
             passes_quality_checks = self.quality_query.get_table_mask(table)[0]
@@ -397,7 +397,7 @@ def __call__(self, event: ArrayEventContainer) -> None:
 
     def predict_table(self, key, table: Table) -> Dict[ReconstructionProperty, Table]:
         """Predict on a table of events"""
-        table = self.feature_generator(table)
+        table = self.feature_generator(table, subarray=self.subarray)
 
         n_rows = len(table)
         energy = u.Quantity(np.full(n_rows, np.nan), self.unit, copy=False)
@@ -439,7 +439,7 @@ class ParticleClassifier(SKLearnClassificationReconstructor):
     def __call__(self, event: ArrayEventContainer) -> None:
         for tel_id in event.trigger.tels_with_trigger:
             table = collect_features(event, tel_id, self.instrument_table)
-            table = self.feature_generator(table)
+            table = self.feature_generator(table, subarray=self.subarray)
             passes_quality_checks = self.quality_query.get_table_mask(table)[0]
 
             if passes_quality_checks:
@@ -464,7 +464,7 @@ def __call__(self, event: ArrayEventContainer) -> None:
 
     def predict_table(self, key, table: Table) -> Dict[ReconstructionProperty, Table]:
         """Predict on a table of events"""
-        table = self.feature_generator(table)
+        table = self.feature_generator(table, subarray=self.subarray)
 
         n_rows = len(table)
         score = np.full(n_rows, np.nan)
@@ -681,7 +681,7 @@ def __call__(self, event: ArrayEventContainer) -> None:
         """
         for tel_id in event.trigger.tels_with_trigger:
             table = collect_features(event, tel_id, self.instrument_table)
-            table = self.feature_generator(table)
+            table = self.feature_generator(table, subarray=self.subarray)
 
             passes_quality_checks = self.quality_query.get_table_mask(table)[0]
 
@@ -751,7 +751,7 @@ def predict_table(self, key, table: Table) -> Dict[ReconstructionProperty, Table
         altaz_table : `~astropy.table.Table`
             Table with resulting predictions of horizontal coordinates
         """
-        table = self.feature_generator(table)
+        table = self.feature_generator(table, subarray=self.subarray)
 
         n_rows = len(table)
         disp = u.Quantity(np.full(n_rows, np.nan), self.unit, copy=False)

ctapipe/resources/train_energy_regressor.yaml

@@ -50,6 +50,8 @@ TrainEnergyRegressor:
       - intensity_skewness
       - intensity_kurtosis
       - area
+      - n_telescopes_triggered
+      - n_telescopes_hillas_reconstructor
 
     stereo_combiner_cls: "StereoMeanCombiner"
     StereoMeanCombiner:
@@ -66,6 +68,8 @@ TrainEnergyRegressor:
     FeatureGenerator:
       features:
         - ["area", "hillas_width * hillas_length"]
+        - ["n_telescopes_triggered", "subarray.multiplicity(tels_with_trigger)"]
+        - ["n_telescopes_hillas_reconstructor", "subarray.multiplicity(HillasReconstructor_telescopes)"]
 
 
 

ctapipe/tools/train_disp_reconstructor.py

@@ -127,7 +127,7 @@ def _read_table(self, telescope_type):
                 f"No events after quality query for telescope type {telescope_type}"
             )
 
-        table = self.models.feature_generator(table)
+        table = self.models.feature_generator(table, subarray=self.loader.subarray)
 
         table[self.models.target] = self._get_true_disp(table)
 

ctapipe/tools/train_energy_regressor.py

@@ -128,7 +128,7 @@ def _read_table(self, telescope_type):
                 f"No events after quality query for telescope type {telescope_type}"
             )
 
-        table = self.regressor.feature_generator(table)
+        table = self.regressor.feature_generator(table, subarray=self.loader.subarray)
 
         feature_names = self.regressor.features + [self.regressor.target]
         table = table[feature_names]

ctapipe/tools/train_particle_classifier.py

@@ -130,13 +130,11 @@ def setup(self):
         if self.signal_loader.subarray != self.background_loader.subarray:
             raise ValueError("Signal and background subarrays do not match")
 
-        self.n_signal.attach_subarray(self.signal_loader.subarray)
-        self.n_background.attach_subarray(self.signal_loader.subarray)
+        self.subarray = self.signal_loader.subarray
+        self.n_signal.attach_subarray(self.subarray)
+        self.n_background.attach_subarray(self.subarray)
 
-        self.classifier = ParticleClassifier(
-            subarray=self.signal_loader.subarray,
-            parent=self,
-        )
+        self.classifier = ParticleClassifier(subarray=self.subarray, parent=self)
         self.rng = np.random.default_rng(self.random_seed)
         self.cross_validate = CrossValidator(
             parent=self, model_component=self.classifier
@@ -179,7 +177,7 @@ def _read_table(self, telescope_type, loader, n_events=None):
                 f"No events after quality query for telescope type {telescope_type}"
             )
 
-        table = self.classifier.feature_generator(table)
+        table = self.classifier.feature_generator(table, subarray=self.subarray)
 
         # Add true energy for energy-dependent performance plots
         columns = self.classifier.features + [self.classifier.target, "true_energy"]

docs/changes/2308.feature.rst

@@ -0,0 +1,7 @@
+Add ``SubarrayDescription.mulitplicity`` method that can compute
+telescope multiplicity for a given telescope boolean mask, either for
+all telescope or a given telescope type.
+
+Enable adding additional keyword arguments to ``FeatureGenerator``.
+
+Pass the ``SubarrayDescription`` to ``FeatureGenerator`` in sklearn classes.