Add event type filter to ctapipe-process (#1841)

ctapipe/tools/process.py

@@ -13,6 +13,8 @@
 from ..io import DataLevel, DataWriter, EventSource, SimTelEventSource
 from ..io.datawriter import DATA_MODEL_VERSION
 from ..reco import ShowerProcessor
+from ..utils import EventTypeFilter
+
 
 COMPATIBLE_DATALEVELS = [
     DataLevel.R1,
@@ -48,10 +50,12 @@ class ProcessorTool(Tool):
     progress_bar = Bool(
         help="show progress bar during processing", default_value=False
     ).tag(config=True)
+
     force_recompute_dl1 = Bool(
         help="Enforce dl1 recomputation even if already present in the input file",
         default_value=False,
     ).tag(config=True)
+
     force_recompute_dl2 = Bool(
         help="Enforce dl2 recomputation even if already present in the input file",
         default_value=False,
@@ -136,6 +140,7 @@ class ProcessorTool(Tool):
         + classes_with_traits(ImageExtractor)
         + classes_with_traits(GainSelector)
         + classes_with_traits(QualityQuery)
+        + classes_with_traits(EventTypeFilter)
     )
 
     def setup(self):
@@ -162,6 +167,7 @@ def setup(self):
             subarray=self.event_source.subarray, parent=self
         )
         self.write = DataWriter(event_source=self.event_source, parent=self)
+        self.event_type_filter = EventTypeFilter(parent=self)
 
         # warn if max_events prevents writing the histograms
         if (
@@ -187,6 +193,7 @@ def should_compute_dl1(self):
         """returns true if we should compute DL1 info"""
         if self.force_recompute_dl1:
             return True
+
         if DataLevel.DL1_PARAMETERS in self.event_source.datalevels:
             return False
 
@@ -248,6 +255,9 @@ def start(self):
 
             self.log.debug("Processessing event_id=%s", event.index.event_id)
 
+            if not self.event_type_filter(event):
+                continue
+
             if self.should_calibrate:
                 self.calibrate(event)
 

ctapipe/utils/__init__.py

@@ -11,6 +11,7 @@
 from .astro import get_bright_stars
 from .cutflow import CutFlow, PureCountingCut, UndefinedCut
 from .index_finder import IndexFinder
+from .event_type_filter import EventTypeFilter
 
 
 __all__ = [
@@ -26,4 +27,5 @@
     "PureCountingCut",
     "UndefinedCut",
     "IndexFinder",
+    "EventTypeFilter",
 ]

ctapipe/utils/event_type_filter.py

@@ -0,0 +1,32 @@
+from ..core import Component
+from ..core.traits import Enum, Set
+from ..containers import EventType
+
+
+__all__ = ["EventTypeFilter"]
+
+
+class EventTypeFilter(Component):
+    """Check that an event has one of the allowed types"""
+
+    allowed_types = Set(
+        # add both the enum instance and the integer values to support
+        # giving the integers in config files.
+        trait=Enum(list(EventType) + [t.value for t in EventType]),
+        default_value=None,
+        allow_none=True,
+        help="The allowed types. Set to None to allow all types.",
+    ).tag(config=True)
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # convert ints to enum type
+        if self.allowed_types is not None:
+            self.allowed_types = {EventType(e) for e in self.allowed_types}
+
+    def __call__(self, event):
+        """Returns True if the event should be kept"""
+        if self.allowed_types is None:
+            return True
+
+        return event.trigger.event_type in self.allowed_types

ctapipe/utils/tests/test_event_filter.py

@@ -0,0 +1,32 @@
+from traitlets.config import Config
+from ctapipe.containers import EventType, ArrayEventContainer
+
+
+def test_event_filter():
+    from ctapipe.utils import EventTypeFilter
+
+    event_filter = EventTypeFilter(
+        allowed_types={EventType.SUBARRAY, EventType.FLATFIELD}
+    )
+
+    e = ArrayEventContainer()
+    e.trigger.event_type = EventType.SUBARRAY
+    assert event_filter(e)
+    e.trigger.event_type = EventType.FLATFIELD
+    assert event_filter(e)
+    e.trigger.event_type = EventType.DARK_PEDESTAL
+    assert not event_filter(e)
+
+
+def test_event_filter_config():
+    from ctapipe.utils import EventTypeFilter
+
+    config = Config({"EventTypeFilter": {"allowed_types": [EventType.SUBARRAY.value]}})
+    event_filter = EventTypeFilter(config=config)
+
+    e = ArrayEventContainer()
+    e.trigger.event_type = EventType.DARK_PEDESTAL
+    assert not event_filter(e)
+
+    e.trigger.event_type = EventType.SUBARRAY
+    assert event_filter(e)