Skip to content

Commit 1f7bbd8

Browse files
kosackkosack
committed
use a FeatureSetRegistry for FeatureSets
1 parent eba8c25 commit 1f7bbd8

3 files changed

Lines changed: 121 additions & 110 deletions

File tree

‎src/ctapipe/io/__init__.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
from .astropy_helpers import read_table, write_table # noqa: I001
99
from .datalevels import DataLevel
1010
from .dl2_tables_preprocessing import DL2EventPreprocessor, DL2EventLoader
11-
from .event_preprocessor import EventPreprocessor, PreprocessorFeatureSet
11+
from .event_preprocessor import EventPreprocessor
1212
from .eventsource import EventSource
1313
from .eventseeker import EventSeeker
1414
from .tableio import TableReader, TableWriter
@@ -48,5 +48,4 @@
4848
"DL2EventLoader",
4949
"get_hdf5_monitoring_types",
5050
"EventPreprocessor",
51-
"PreprocessorFeatureSet",
5251
]

‎src/ctapipe/io/event_preprocessor.py‎

Lines changed: 115 additions & 103 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,6 @@
11
"""Module containing classes related to event loading and preprocessing"""
22

3-
from enum import StrEnum, auto
4-
53
from astropy.coordinates import angular_separation
6-
from traitlets import default
74

85
from ..coordinates import altaz_to_nominal
96
from ..core import (
@@ -17,11 +14,113 @@
1714
__all__ = ["EventPreprocessor"]
1815

1916

20-
class PreprocessorFeatureSet(StrEnum):
21-
"""Pre-defined configurations for DL2EventPreprocessor for specific use cases."""
17+
from typing import Callable
18+
19+
20+
class FeatureSetRegistry:
21+
"""Registry for custom feature set configurations."""
2222

23-
custom = auto() #: use user-supplied configuration
24-
dl2_irf = auto() #: support IRF preprocessing use case
23+
_registry = {}
24+
25+
@classmethod
26+
def register(cls, name: str):
27+
"""Register a feature set configuration.
28+
29+
Examples
30+
--------
31+
>>> @FeatureSetRegistry.register("my_analysis")
32+
... def my_config(preprocessor):
33+
... return {
34+
... "features_to_generate": [("custom", "col_a / col_b")],
35+
... "quality_criteria": [("cut", "custom > 0.5")],
36+
... "output_features": ["event_id", "custom"]
37+
... }
38+
"""
39+
40+
def decorator(func: Callable):
41+
cls._registry[name] = func
42+
return func
43+
44+
return decorator
45+
46+
@classmethod
47+
def get(cls, name: str):
48+
"""Get a registered configuration function."""
49+
return cls._registry.get(name)
50+
51+
@classmethod
52+
def list_available(cls):
53+
"""List all registered feature set names."""
54+
return list(cls._registry.keys())
55+
56+
57+
@FeatureSetRegistry.register("dl2_irf")
58+
def _dl2_irf_config(preprocessor):
59+
"""Built-in configuration for DL2 IRF generation."""
60+
return {
61+
"features_to_generate": [
62+
("reco_energy", f"{preprocessor.energy_reconstructor}_energy"),
63+
("reco_alt", f"{preprocessor.geometry_reconstructor}_alt"),
64+
("reco_az", f"{preprocessor.geometry_reconstructor}_az"),
65+
("gh_score", f"{preprocessor.gammaness_reconstructor}_prediction"),
66+
("theta", "angular_separation(reco_az, reco_alt, true_az, true_alt)"),
67+
(
68+
"reco_fov_coord",
69+
"altaz_to_nominal(reco_az, reco_alt, subarray_pointing_lon, subarray_pointing_lat)",
70+
),
71+
(
72+
"reco_fov_lon",
73+
"reco_fov_coord[:,0]",
74+
), # note: GADF IRFs use the negative of this
75+
("reco_fov_lat", "reco_fov_coord[:,1]"),
76+
(
77+
"true_fov_coord",
78+
"altaz_to_nominal(true_az, true_alt, subarray_pointing_lon, subarray_pointing_lat)",
79+
),
80+
(
81+
"true_fov_lon",
82+
"true_fov_coord[:,0]",
83+
), # note: GADF IRFs use the negative of this
84+
("true_fov_lat", "true_fov_coord[:,1]"),
85+
(
86+
"true_fov_offset",
87+
"angular_separation(true_fov_lon, true_fov_lat, 0*u.deg, 0*u.deg)",
88+
),
89+
(
90+
"reco_fov_offset",
91+
"angular_separation(reco_fov_lon, reco_fov_lat, 0*u.deg, 0*u.deg)",
92+
),
93+
(
94+
"multiplicity",
95+
f"np.count_nonzero({preprocessor.gammaness_reconstructor}_telescopes,axis=1)",
96+
),
97+
],
98+
"quality_criteria": [
99+
("Valid geometry", f"{preprocessor.geometry_reconstructor}_is_valid"),
100+
("valid energy", f"{preprocessor.energy_reconstructor}_is_valid"),
101+
("valid gammaness", f"{preprocessor.gammaness_reconstructor}_is_valid"),
102+
("sufficient multiplicity", "multiplicity >= 4"),
103+
],
104+
"output_features": [
105+
"event_id",
106+
"obs_id",
107+
"reco_energy",
108+
"reco_alt",
109+
"reco_az",
110+
"gh_score",
111+
"true_energy",
112+
"true_alt",
113+
"true_az",
114+
"true_fov_offset",
115+
"reco_fov_offset",
116+
"theta",
117+
"reco_fov_lat",
118+
"true_fov_lat",
119+
"reco_fov_lon",
120+
"true_fov_lon",
121+
"multiplicity",
122+
],
123+
}
25124

26125

27126
class EventPreprocessor(Component):
@@ -56,9 +155,9 @@ class with the columns you to retain in the output table.
56155
help="Prefix of the classifier `_prediction` column",
57156
).tag(config=True)
58157

59-
feature_set = traits.UseEnum(
60-
PreprocessorFeatureSet,
61-
default_value=PreprocessorFeatureSet.dl2_irf,
158+
feature_set = traits.CaselessStrEnum(
159+
["custom"] + FeatureSetRegistry.list_available(),
160+
default_value="custom",
62161
help=(
63162
"Set up the FeatureGenerator.features, output features, and quality criteria "
64163
"based on standard use cases."
@@ -79,16 +178,18 @@ class with the columns you to retain in the output table.
79178

80179
def __init__(self, config=None, parent=None, **kwargs):
81180
super().__init__(config=config, parent=parent, **kwargs)
82-
if self.feature_set == PreprocessorFeatureSet.custom:
181+
if self.feature_set == "custom":
83182
self.feature_generator = FeatureGenerator(parent=self)
84183
self.quality_query = QualityQuery(parent=self)
85-
else:
184+
else: # use a pre-registered feature set
185+
feature_set = FeatureSetRegistry.get(self.feature_set)(self)
86186
self.feature_generator = FeatureGenerator(
87-
parent=self, features=self._get_predefined_features_to_generate()
187+
parent=self, features=feature_set["features_to_generate"]
88188
)
89189
self.quality_query = QualityQuery(
90-
parent=self, quality_criteria=self._get_predefined_quality_criteria()
190+
parent=self, quality_criteria=feature_set["quality_criteria"]
91191
)
192+
self.features = feature_set["output_features"]
92193
# sanity checks:
93194
if len(self.features) == 0:
94195
raise ToolConfigurationError(
@@ -114,92 +215,3 @@ def __call__(self, table):
114215
# return only the columns specified in `self.features`, and rows in
115216
# `selected_mask`
116217
return generated[self.features][selected_mask]
117-
118-
def _get_predefined_features_to_generate(self) -> list[tuple]:
119-
"""Return a default list of FeatureGenerator features."""
120-
if self.feature_set == PreprocessorFeatureSet.dl2_irf:
121-
# Default features for DL2/Subarray events
122-
return [
123-
("reco_energy", f"{self.energy_reconstructor}_energy"),
124-
("reco_alt", f"{self.geometry_reconstructor}_alt"),
125-
("reco_az", f"{self.geometry_reconstructor}_az"),
126-
("gh_score", f"{self.gammaness_reconstructor}_prediction"),
127-
("theta", "angular_separation(reco_az, reco_alt, true_az, true_alt)"),
128-
(
129-
"reco_fov_coord",
130-
"altaz_to_nominal(reco_az, reco_alt, subarray_pointing_lon, subarray_pointing_lat)",
131-
),
132-
(
133-
"reco_fov_lon",
134-
"reco_fov_coord[:,0]",
135-
), # note: GADF IRFs use the negative of this
136-
("reco_fov_lat", "reco_fov_coord[:,1]"),
137-
(
138-
"true_fov_coord",
139-
"altaz_to_nominal(true_az, true_alt, subarray_pointing_lon, subarray_pointing_lat)",
140-
),
141-
(
142-
"true_fov_lon",
143-
"true_fov_coord[:,0]",
144-
), # note: GADF IRFs use the negative of this
145-
("true_fov_lat", "true_fov_coord[:,1]"),
146-
(
147-
"true_fov_offset",
148-
"angular_separation(true_fov_lon, true_fov_lat, 0*u.deg, 0*u.deg)",
149-
),
150-
(
151-
"reco_fov_offset",
152-
"angular_separation(reco_fov_lon, reco_fov_lat, 0*u.deg, 0*u.deg)",
153-
),
154-
(
155-
"multiplicity",
156-
f"np.count_nonzero({self.gammaness_reconstructor}_telescopes,axis=1)",
157-
),
158-
]
159-
else:
160-
raise NotImplementedError(f"unsupported feature_set: {self.feature_set}")
161-
162-
def _get_predefined_quality_criteria(self) -> list[tuple]:
163-
"""
164-
Set the quality criteria for a DL2FeatureSet.
165-
166-
Here you can use any columns in the input table, or any that are
167-
specified in the FeatureGenerator.
168-
"""
169-
if self.feature_set == PreprocessorFeatureSet.dl2_irf:
170-
return [
171-
("Valid geometry", f"{self.geometry_reconstructor}_is_valid"),
172-
("valid energy", f"{self.energy_reconstructor}_is_valid"),
173-
("valid gammaness", f"{self.gammaness_reconstructor}_is_valid"),
174-
("sufficient multiplicity", "multiplicity >= 4"),
175-
]
176-
else:
177-
raise NotImplementedError(f"unsupported feature_set: {self.feature_set}")
178-
179-
@default("features")
180-
def _features(self):
181-
"""Set the columns to output, for a given FeatureSet."""
182-
if self.feature_set == PreprocessorFeatureSet.dl2_irf:
183-
return [
184-
"event_id",
185-
"obs_id",
186-
"reco_energy",
187-
"reco_alt",
188-
"reco_az",
189-
"gh_score",
190-
"true_energy",
191-
"true_alt",
192-
"true_az",
193-
"true_fov_offset",
194-
"reco_fov_offset",
195-
"theta",
196-
"reco_fov_lat",
197-
"true_fov_lat",
198-
"reco_fov_lon",
199-
"true_fov_lon",
200-
"multiplicity",
201-
]
202-
elif self.feature_set == PreprocessorFeatureSet.custom:
203-
return []
204-
else:
205-
raise NotImplementedError(f"unsupported feature_set: {self.feature_set}")

‎src/ctapipe/io/tests/test_event_preprocessor.py‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
from astropy import units as u
66
from astropy.table import QTable
77

8-
from ctapipe.io import PreprocessorFeatureSet
8+
from ctapipe.io.event_preprocessor import FeatureSetRegistry
99

1010

1111
@pytest.fixture(scope="function")
@@ -39,7 +39,7 @@ def minimal_dl2_table():
3939
)
4040

4141

42-
@pytest.mark.parametrize("feature_set", list(PreprocessorFeatureSet))
42+
@pytest.mark.parametrize("feature_set", FeatureSetRegistry.list_available())
4343
def test_event_preprocessing(feature_set, minimal_dl2_table):
4444
from traitlets.config import Config
4545

@@ -65,10 +65,10 @@ def test_event_preprocessing(feature_set, minimal_dl2_table):
6565
def test_no_output():
6666
"""Check error is raised if no columns are specified for output."""
6767
from ctapipe.core import ToolConfigurationError
68-
from ctapipe.io import EventPreprocessor, PreprocessorFeatureSet
68+
from ctapipe.io import EventPreprocessor
6969

7070
with pytest.raises(ToolConfigurationError):
71-
EventPreprocessor(feature_set=PreprocessorFeatureSet.custom)
71+
EventPreprocessor(feature_set="custom")
7272

7373

7474
def test_nondefault_reconstructors(minimal_dl2_table):
@@ -94,7 +94,7 @@ def test_nondefault_reconstructors(minimal_dl2_table):
9494
table[f"{gammaness}_telescopes"] = table["RandomForestClassifier_telescopes"]
9595

9696
preprocess = EventPreprocessor(
97-
feature_set=PreprocessorFeatureSet.dl2_irf,
97+
feature_set="dl2_irf",
9898
geometry_reconstructor=geom,
9999
energy_reconstructor=energy,
100100
gammaness_reconstructor=gammaness,

0 commit comments

Comments
 (0)