Implement a shower processor

ctapipe/reco/__init__.py

@@ -1,6 +1,10 @@
 # Licensed under a 3-clause BSD style license - see LICENSE.rst
 from .HillasReconstructor import HillasReconstructor, Reconstructor
 from .ImPACT import ImPACTReconstructor
+from .shower_processor import ShowerProcessor
 
 
-__all__ = ["HillasReconstructor", "Reconstructor", "ImPACTReconstructor"]
+__all__ = ["ShowerProcessor",
+           "HillasReconstructor",
+           "Reconstructor",
+           "ImPACTReconstructor"]

ctapipe/reco/shower_processor.py

@@ -0,0 +1,193 @@
+"""
+High level processing of showers.
+
+This processor will be able to process a shower/event in 3 steps:
+- shower geometry
+- estimation of energy (optional, currently unavailable)
+- estimation of classification (optional, currently unavailable)
+
+"""
+import numpy as np
+
+from astropy.coordinates import SkyCoord, AltAz
+
+from ctapipe.core import Component, QualityQuery
+from ctapipe.core.traits import List
+from ctapipe.containers import ArrayEventContainer, ReconstructedShowerContainer
+from ctapipe.instrument import SubarrayDescription
+from ctapipe.reco import HillasReconstructor
+
+
+DEFAULT_SHOWER_PARAMETERS = ReconstructedShowerContainer()
+
+
+class ShowerQualityQuery(QualityQuery):
+    """Configuring shower-wise data checks."""
+
+    quality_criteria = List(
+        default_value=[
+            ("Positive charge", "lambda hillas: hillas.intensity > 0"),
+            ("Positive width", "lambda hillas: hillas.width.value > 0"),
+            ("Positive length", "lambda hillas: hillas.length.value > 0"),
+            ("Ellipticity > 0.1", "lambda hillas: hillas.width.value / hillas.length.value > 0.1"),
+            ("Ellipticity < 0.6", "lambda hillas: hillas.width.value / hillas.length.value < 0.6")
+        ],
+        help=QualityQuery.quality_criteria.help,
+    ).tag(config=True)
+
+
+class ShowerProcessor(Component):
+    """
+    Takes DL1/parameters data and estimates the shower geometry.
+    Should be run after ImageProcessor to produce all DL1b information.
+
+    For now the only supported reconstructor is HillasReconstructor.
+    """
+
+    def __init__(
+        self,
+        subarray: SubarrayDescription,
+        is_simulation,
+        reconstruct_energy,
+        classify,
+        config=None,
+        parent=None,
+        **kwargs,
+    ):
+        """
+        Parameters
+        ----------
+        subarray: SubarrayDescription
+            Description of the subarray. Provides information about the
+            camera which are useful in calibration. Also required for
+            configuring the TelescopeParameter traitlets.
+        is_simulation: bool
+            If true, also process simulated images if they exist
+        config: traitlets.loader.Config
+            Configuration specified by config file or cmdline arguments.
+            Used to set traitlet values.
+            This is mutually exclusive with passing a ``parent``.
+        parent: ctapipe.core.Component or ctapipe.core.Tool
+            Parent of this component in the configuration hierarchy,
+            this is mutually exclusive with passing ``config``
+        """
+
+        super().__init__(config=config, parent=parent, **kwargs)
+        self.subarray = subarray
+        self.check_shower = ShowerQualityQuery(parent=self)
+        self._is_simulation = is_simulation
+        self.reconstruct_energy = reconstruct_energy
+        self.classify = classify
+        self.reconstructor = HillasReconstructor()
+
+    def _reconstruct_shower(
+        self,
+        event,
+        default=DEFAULT_SHOWER_PARAMETERS,
+    ) -> ReconstructedShowerContainer:
+        """Perform shower reconstruction.
+
+        Parameters
+        ----------
+        tel_id: int
+            which telescope is being cleaned
+        image: np.ndarray
+            image to process
+        signal_pixels: np.ndarray[bool]
+            image mask
+        peak_time: np.ndarray
+            peak time image
+        Returns
+        -------
+        ReconstructedShowerContainer:
+            direction in the sky with uncertainty,
+            core position on the ground with uncertainty,
+            h_max with uncertainty,
+            is_valid boolean for successfull reconstruction,
+            average intensity of the intensities used for reconstruction,
+            measure of algorithm success (if fit),
+            list of tel_ids used if stereo, or None if Mono
+        """
+
+        # Read only valid HillasContainers (minimum condition to continue)
+        hillas_dict = {
+            tel_id: dl1.parameters.hillas
+            for tel_id, dl1 in event.dl1.tel.items()
+            if np.isfinite(event.dl1.tel[tel_id].parameters.hillas.intensity)
+        }
+
+        # On top of this check if the shower should be considered based
+        # on the user's configuration
+        shower_criteria = [self.check_shower(hillas_dict[tel_id]) for tel_id in hillas_dict]
+        self.log.debug(
+            "shower_criteria: %s",
+            list(zip(self.check_shower.criteria_names[1:], shower_criteria)),
+        )
+
+        # Reconstruct the shower only if all shower criteria are met
+        if np.count_nonzero(shower_criteria) > 2:
+
+            array_pointing = SkyCoord(
+                az=event.pointing.array_azimuth,
+                alt=event.pointing.array_altitude,
+                frame=AltAz(),
+            )
+
+            telescopes_pointings = {
+                tel_id: SkyCoord(
+                    alt=event.pointing.tel[tel_id].altitude,
+                    az=event.pointing.tel[tel_id].azimuth,
+                    frame=AltAz(),
+                )
+                for tel_id in hillas_dict
+            }
+
+            result = self.reconstructor.predict(hillas_dict,
+                                                self.subarray,
+                                                array_pointing,
+                                                telescopes_pointings)
+
+            return result
+
+        else:
+            return default
+
+    def _reconstruct_energy(self, event: ArrayEventContainer):
+        raise NotImplementedError("TO DO")
+
+    def _estimate_classification(self, event: ArrayEventContainer):
+        raise NotImplementedError("TO DO")
+
+    def _process_reconstructed_energy(self, event: ArrayEventContainer):
+        self._reconstruct_energy(event)
+
+    def _process_reconstructed_classification(self, event: ArrayEventContainer):
+        self._estimate_classification(event)
+
+    def _process_shower_geometry(self, event: ArrayEventContainer):
+        """Record the reconstructed shower geometry into the ArrayEventContainer."""
+
+        shower_geometry = self._reconstruct_shower(event)
+
+        self.log.debug("shower geometry: %s", shower_geometry.as_dict(recursive=True))
+
+        event.dl2.shower["HillasReconstructor"] = shower_geometry
+
+    def __call__(self, event: ArrayEventContainer):
+        """
+        Perform the full shower geometry reconstruction on the input event.
+
+        Parameters
+        ----------
+        event : container
+            A `ctapipe` event container
+        """
+
+        # This is always done when calling the ShowerProcessor
+        self._process_shower_geometry(event)
+
+        if self.reconstruct_energy:
+            self._process_reconstructed_classification(event)
+
+        if self.classify:
+            self._process_reconstructed_energy(event)

ctapipe/reco/tests/test_shower_processor.py

@@ -0,0 +1,92 @@
+"""
+Tests for ShowerProcessor functionalities.
+"""
+import pytest
+from numpy import isfinite
+
+from traitlets.config.loader import Config
+
+from ctapipe.calib import CameraCalibrator
+from ctapipe.image import ImageProcessor
+from ctapipe.reco import ShowerProcessor
+
+
+def test_shower_processor_geometry(example_event, example_subarray):
+    """Ensure we get shower geometry when we input an event with parametrized images."""
+
+    calibrate = CameraCalibrator(subarray=example_subarray)
+
+    process_images = ImageProcessor(
+        subarray=example_subarray,
+        is_simulation=True,
+        image_cleaner_type="MARSImageCleaner",
+    )
+
+    process_shower = ShowerProcessor(
+        subarray=example_subarray,
+        is_simulation=True,
+        reconstruct_energy=False,
+        classify=False
+    )
+
+    calibrate(example_event)
+    process_images(example_event)
+
+    process_shower(example_event)
+    print(process_shower.check_shower.to_table())
+
+    DL2a = example_event.dl2.shower["HillasReconstructor"]
+    print(DL2a)
+    assert isfinite(DL2a.alt)
+    assert isfinite(DL2a.az)
+    assert isfinite(DL2a.core_x)
+    assert isfinite(DL2a.core_x)
+    assert isfinite(DL2a.core_y)
+    assert DL2a.is_valid
+    assert isfinite(DL2a.average_intensity)
+
+    # Increase some quality cuts and check that we get defaults
+    config = Config()
+    config.ShowerQualityQuery.quality_criteria = [("> 500 phes", "lambda hillas: hillas.intensity > 500")]
+
+    process_shower = ShowerProcessor(
+        config=config,
+        subarray=example_subarray,
+        is_simulation=True,
+        reconstruct_energy=False,
+        classify=False
+    )
+
+    process_shower(example_event)
+    print(process_shower.check_shower.to_table())
+
+    DL2a = example_event.dl2.shower["HillasReconstructor"]
+    assert not isfinite(DL2a.alt)
+    assert not isfinite(DL2a.az)
+    assert not isfinite(DL2a.core_x)
+    assert not isfinite(DL2a.core_x)
+    assert not isfinite(DL2a.core_y)
+    assert not DL2a.is_valid
+    assert not isfinite(DL2a.average_intensity)
+
+    # Now check that if energy reconstruction is enabled we get a TODO error
+    process_shower = ShowerProcessor(
+        config=config,
+        subarray=example_subarray,
+        is_simulation=True,
+        reconstruct_energy=True,
+        classify=False
+    )
+    with pytest.raises(NotImplementedError) as error_info:
+        process_shower(example_event)
+
+    # also for classification
+    process_shower = ShowerProcessor(
+        config=config,
+        subarray=example_subarray,
+        is_simulation=True,
+        reconstruct_energy=False,
+        classify=True
+    )
+    with pytest.raises(NotImplementedError) as error_info:
+        process_shower(example_event)