diff --git a/ctapipe/tools/process.py b/ctapipe/tools/process.py index e26a4466922..24846ac6d86 100644 --- a/ctapipe/tools/process.py +++ b/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) diff --git a/ctapipe/utils/__init__.py b/ctapipe/utils/__init__.py index dcdd2bed660..2e6eebd39aa 100644 --- a/ctapipe/utils/__init__.py +++ b/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", ] diff --git a/ctapipe/utils/event_type_filter.py b/ctapipe/utils/event_type_filter.py new file mode 100644 index 00000000000..aac751c5d36 --- /dev/null +++ b/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 diff --git a/ctapipe/utils/tests/test_event_filter.py b/ctapipe/utils/tests/test_event_filter.py new file mode 100644 index 00000000000..5ed567955a0 --- /dev/null +++ b/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)