From 318b707dead758860ada32facaae6ea49cdd782b Mon Sep 17 00:00:00 2001 From: Severin Magel Date: Tue, 14 Jul 2026 10:10:27 -0400 Subject: [PATCH 1/4] PrometheusReader: raise on missing extractor tables instead of skipping An extractor whose table is absent from the input file (e.g. a non-default Prometheus photon_field_name, or a typo'd table_name) was silently skipped, producing truth-only databases with no diagnostic. Validate all configured tables once per file and raise a ValueError naming the missing and available tables. Part of #909. Co-Authored-By: Claude Fable 5 --- .../data/readers/prometheus_reader.py | 24 ++++++++--- tests/data/test_prometheus_reader.py | 40 +++++++++++++++++++ 2 files changed, 59 insertions(+), 5 deletions(-) create mode 100644 tests/data/test_prometheus_reader.py diff --git a/src/graphnet/data/readers/prometheus_reader.py b/src/graphnet/data/readers/prometheus_reader.py index 3f97c4468..3bf2cc0e0 100644 --- a/src/graphnet/data/readers/prometheus_reader.py +++ b/src/graphnet/data/readers/prometheus_reader.py @@ -22,17 +22,31 @@ def __call__(self, file_path: str) -> List[OrderedDict]: Returns: Extracted data. + + Raises: + ValueError: If a table required by one of the configured + extractors does not exist in the file. """ # Open file outputs = [] file = pd.read_parquet(file_path) + extractors: List[PrometheusExtractor] = [] + missing_tables = [] + for extractor in self._extractors: + assert isinstance(extractor, PrometheusExtractor) + extractors.append(extractor) + if extractor._table not in file.columns: + missing_tables.append(extractor._table) + if missing_tables: + raise ValueError( + f"Table(s) {missing_tables} not found in {file_path}. " + f"Available tables: {list(file.columns)}." + ) for k in range(len(file)): # Loop over events in file extracted_event = OrderedDict() - for extractor in self._extractors: - assert isinstance(extractor, PrometheusExtractor) - if extractor._table in file.columns: - output = extractor(file[extractor._table][k]) - extracted_event[extractor._extractor_name] = output + for extractor in extractors: + output = extractor(file[extractor._table][k]) + extracted_event[extractor._extractor_name] = output outputs.append(extracted_event) return outputs diff --git a/tests/data/test_prometheus_reader.py b/tests/data/test_prometheus_reader.py new file mode 100644 index 000000000..f4a9863a5 --- /dev/null +++ b/tests/data/test_prometheus_reader.py @@ -0,0 +1,40 @@ +"""Tests for the PrometheusReader.""" + +import os + +import pytest + +from graphnet.constants import TEST_DATA_DIR +from graphnet.data.extractors.prometheus import ( + PrometheusFeatureExtractor, + PrometheusTruthExtractor, +) +from graphnet.data.readers import PrometheusReader + +FILE_PATH = os.path.join( + TEST_DATA_DIR, "prometheus", "22980001_photons.parquet" +) + + +def test_prometheus_reader_extracts_configured_tables() -> None: + """Reader extracts every configured table when all are present.""" + reader = PrometheusReader() + reader.set_extractors( + [PrometheusTruthExtractor(), PrometheusFeatureExtractor()] + ) + events = reader(FILE_PATH) + assert len(events) > 0 + assert set(events[0].keys()) == {"mc_truth", "photons"} + + +def test_prometheus_reader_raises_on_missing_table() -> None: + """Reader raises if an extractor's table is not in the file.""" + reader = PrometheusReader() + reader.set_extractors( + [ + PrometheusTruthExtractor(), + PrometheusFeatureExtractor(table_name="not_a_table"), + ] + ) + with pytest.raises(ValueError, match="not_a_table"): + reader(FILE_PATH) From 376c54777ab2cde3e2ce59ea7f8f7e263b597d56 Mon Sep 17 00:00:00 2001 From: Severin Magel Date: Tue, 14 Jul 2026 10:22:08 -0400 Subject: [PATCH 2/4] Sync TRUTH.PROMETHEUS with current Prometheus output TRUTH.PROMETHEUS still described a truth schema (injection_*, primary_lepton_1_*, ...) that Prometheus stopped writing long ago; using it against a freshly converted file silently strips every truth column. Rename that list to TRUTH.PROMETHEUS_LEGACY (it matches older files such as the bundled example database, which the 04_training examples run on) and set TRUTH.PROMETHEUS to the schema Prometheus writes today. PrometheusTruthExtractor now takes its default columns from TRUTH.PROMETHEUS so the constant and the extractor cannot drift apart again. Part of #909. Co-Authored-By: Claude Fable 5 --- examples/04_training/01_train_dynedge.py | 2 +- examples/04_training/02_train_tito_model.py | 2 +- examples/04_training/05_train_RNN_TITO.py | 2 +- examples/04_training/06_train_icemix_model.py | 2 +- examples/04_training/07_train_normalizing_flow.py | 2 +- examples/04_training/08_train_grit_model.py | 2 +- src/graphnet/data/constants.py | 14 ++++++++++++++ .../extractors/prometheus/prometheus_extractor.py | 15 ++++----------- tests/data/test_prometheus_reader.py | 2 ++ 9 files changed, 26 insertions(+), 17 deletions(-) diff --git a/examples/04_training/01_train_dynedge.py b/examples/04_training/01_train_dynedge.py index cf45f11a3..a2cc886e4 100644 --- a/examples/04_training/01_train_dynedge.py +++ b/examples/04_training/01_train_dynedge.py @@ -24,7 +24,7 @@ # Constants features = FEATURES.PROMETHEUS -truth = TRUTH.PROMETHEUS +truth = TRUTH.PROMETHEUS_LEGACY def main( diff --git a/examples/04_training/02_train_tito_model.py b/examples/04_training/02_train_tito_model.py index 4dfac2f49..8a96605b3 100644 --- a/examples/04_training/02_train_tito_model.py +++ b/examples/04_training/02_train_tito_model.py @@ -26,7 +26,7 @@ # Constants features = FEATURES.PROMETHEUS -truth = TRUTH.PROMETHEUS +truth = TRUTH.PROMETHEUS_LEGACY def main( diff --git a/examples/04_training/05_train_RNN_TITO.py b/examples/04_training/05_train_RNN_TITO.py index 948d4a3c8..0eee3063b 100644 --- a/examples/04_training/05_train_RNN_TITO.py +++ b/examples/04_training/05_train_RNN_TITO.py @@ -30,7 +30,7 @@ # Constants features = FEATURES.PROMETHEUS -truth = TRUTH.PROMETHEUS +truth = TRUTH.PROMETHEUS_LEGACY def main( diff --git a/examples/04_training/06_train_icemix_model.py b/examples/04_training/06_train_icemix_model.py index b76c91aa4..7395f7c8f 100644 --- a/examples/04_training/06_train_icemix_model.py +++ b/examples/04_training/06_train_icemix_model.py @@ -32,7 +32,7 @@ # Constants features = FEATURES.PROMETHEUS -truth = TRUTH.PROMETHEUS +truth = TRUTH.PROMETHEUS_LEGACY def main( diff --git a/examples/04_training/07_train_normalizing_flow.py b/examples/04_training/07_train_normalizing_flow.py index b78b34607..56bf0ff35 100644 --- a/examples/04_training/07_train_normalizing_flow.py +++ b/examples/04_training/07_train_normalizing_flow.py @@ -30,7 +30,7 @@ # Constants features = FEATURES.PROMETHEUS -truth = TRUTH.PROMETHEUS +truth = TRUTH.PROMETHEUS_LEGACY def main( diff --git a/examples/04_training/08_train_grit_model.py b/examples/04_training/08_train_grit_model.py index 24a5394a2..da18d99d8 100644 --- a/examples/04_training/08_train_grit_model.py +++ b/examples/04_training/08_train_grit_model.py @@ -24,7 +24,7 @@ # Constants features = FEATURES.PROMETHEUS -truth = TRUTH.PROMETHEUS +truth = TRUTH.PROMETHEUS_LEGACY def main( diff --git a/src/graphnet/data/constants.py b/src/graphnet/data/constants.py index 702974da9..3b8736e0b 100644 --- a/src/graphnet/data/constants.py +++ b/src/graphnet/data/constants.py @@ -87,7 +87,21 @@ class TRUTH: ] DEEPCORE = ICECUBE86 UPGRADE = DEEPCORE + # Truth schema written by current Prometheus versions; the default + # column list of `PrometheusTruthExtractor`. PROMETHEUS = [ + "interaction", + "initial_state_energy", + "initial_state_type", + "initial_state_zenith", + "initial_state_azimuth", + "initial_state_x", + "initial_state_y", + "initial_state_z", + ] + # Truth schema written by older Prometheus versions, e.g. the bundled + # example database `prometheus-events.db`. + PROMETHEUS_LEGACY = [ "injection_energy", "injection_type", "injection_interaction_type", diff --git a/src/graphnet/data/extractors/prometheus/prometheus_extractor.py b/src/graphnet/data/extractors/prometheus/prometheus_extractor.py index b9481d483..e7d9f36bc 100644 --- a/src/graphnet/data/extractors/prometheus/prometheus_extractor.py +++ b/src/graphnet/data/extractors/prometheus/prometheus_extractor.py @@ -4,6 +4,7 @@ import pandas as pd import numpy as np +from graphnet.data.constants import TRUTH from graphnet.data.extractors import Extractor @@ -59,17 +60,9 @@ def __init__(self, table_name: str = "mc_truth") -> None: table_name: Name of the table in the parquet files that contain event-level truth. Defaults to "mc_truth". """ - columns = [ - "interaction", - "initial_state_energy", - "initial_state_type", - "initial_state_zenith", - "initial_state_azimuth", - "initial_state_x", - "initial_state_y", - "initial_state_z", - ] - super().__init__(extractor_name=table_name, columns=columns) + super().__init__( + extractor_name=table_name, columns=list(TRUTH.PROMETHEUS) + ) class PrometheusFeatureExtractor(PrometheusExtractor): diff --git a/tests/data/test_prometheus_reader.py b/tests/data/test_prometheus_reader.py index f4a9863a5..35e187295 100644 --- a/tests/data/test_prometheus_reader.py +++ b/tests/data/test_prometheus_reader.py @@ -5,6 +5,7 @@ import pytest from graphnet.constants import TEST_DATA_DIR +from graphnet.data.constants import TRUTH from graphnet.data.extractors.prometheus import ( PrometheusFeatureExtractor, PrometheusTruthExtractor, @@ -25,6 +26,7 @@ def test_prometheus_reader_extracts_configured_tables() -> None: events = reader(FILE_PATH) assert len(events) > 0 assert set(events[0].keys()) == {"mc_truth", "photons"} + assert set(TRUTH.PROMETHEUS) <= set(events[0]["mc_truth"].keys()) def test_prometheus_reader_raises_on_missing_table() -> None: From e9b905593d46099b06a9a595e2810015599ba78e Mon Sep 17 00:00:00 2001 From: Severin Magel Date: Tue, 14 Jul 2026 10:30:22 -0400 Subject: [PATCH 3/4] PrometheusTruthExtractor: allow overriding the default truth columns The default (TRUTH.PROMETHEUS) covers the event-level scalars every Prometheus injection type writes; a columns override lets users add injection-specific fields such as LeptonInjector's bjorken_x/bjorken_y/ column_depth without subclassing. Part of #909. Co-Authored-By: Claude Fable 5 --- .../prometheus/prometheus_extractor.py | 19 ++++++++++++++----- tests/data/test_prometheus_reader.py | 13 +++++++++++++ 2 files changed, 27 insertions(+), 5 deletions(-) diff --git a/src/graphnet/data/extractors/prometheus/prometheus_extractor.py b/src/graphnet/data/extractors/prometheus/prometheus_extractor.py index e7d9f36bc..f3427a92d 100644 --- a/src/graphnet/data/extractors/prometheus/prometheus_extractor.py +++ b/src/graphnet/data/extractors/prometheus/prometheus_extractor.py @@ -1,6 +1,6 @@ """Parquet Extractor for conversion of simulation files from PROMETHEUS.""" -from typing import List +from typing import List, Optional import pandas as pd import numpy as np @@ -53,16 +53,25 @@ class PrometheusTruthExtractor(PrometheusExtractor): This Extractor will "initial_state" i.e. neutrino truth. """ - def __init__(self, table_name: str = "mc_truth") -> None: + def __init__( + self, + table_name: str = "mc_truth", + columns: Optional[List[str]] = None, + ) -> None: """Construct PrometheusTruthExtractor. Args: table_name: Name of the table in the parquet files that contain event-level truth. Defaults to "mc_truth". + columns: Columns to extract from the table. Defaults to + `TRUTH.PROMETHEUS`, the truth fields written by every + Prometheus injection type. Override to include truth-level + data not present per default, e.g. LeptonInjector's + "bjorken_x", "bjorken_y" and "column_depth". """ - super().__init__( - extractor_name=table_name, columns=list(TRUTH.PROMETHEUS) - ) + if columns is None: + columns = list(TRUTH.PROMETHEUS) + super().__init__(extractor_name=table_name, columns=columns) class PrometheusFeatureExtractor(PrometheusExtractor): diff --git a/tests/data/test_prometheus_reader.py b/tests/data/test_prometheus_reader.py index 35e187295..12d2f96da 100644 --- a/tests/data/test_prometheus_reader.py +++ b/tests/data/test_prometheus_reader.py @@ -1,12 +1,14 @@ """Tests for the PrometheusReader.""" import os +from typing import List import pytest from graphnet.constants import TEST_DATA_DIR from graphnet.data.constants import TRUTH from graphnet.data.extractors.prometheus import ( + PrometheusExtractor, PrometheusFeatureExtractor, PrometheusTruthExtractor, ) @@ -29,6 +31,17 @@ def test_prometheus_reader_extracts_configured_tables() -> None: assert set(TRUTH.PROMETHEUS) <= set(events[0]["mc_truth"].keys()) +def test_prometheus_truth_extractor_columns_override() -> None: + """Truth extractor extracts only the overridden columns.""" + reader = PrometheusReader() + extractors: List[PrometheusExtractor] = [ + PrometheusTruthExtractor(columns=["initial_state_energy"]) + ] + reader.set_extractors(extractors) + events = reader(FILE_PATH) + assert list(events[0]["mc_truth"].keys()) == ["initial_state_energy"] + + def test_prometheus_reader_raises_on_missing_table() -> None: """Reader raises if an extractor's table is not in the file.""" reader = PrometheusReader() From 068603fe39f3073d1c967aadfc087e4750ca33ee Mon Sep 17 00:00:00 2001 From: Severin Magel Date: Tue, 14 Jul 2026 10:47:08 -0400 Subject: [PATCH 4/4] Point remaining TRUTH.PROMETHEUS consumers at the right schema The datamodule tests run on the bundled legacy example database, so they need TRUTH.PROMETHEUS_LEGACY; PublicPrometheusDataset duplicated the current schema as a literal list and now derives it from TRUTH.PROMETHEUS instead. Part of #909. Co-Authored-By: Claude Fable 5 --- src/graphnet/datasets/prometheus_datasets.py | 13 ++----------- tests/data/test_datamodule.py | 2 +- 2 files changed, 3 insertions(+), 12 deletions(-) diff --git a/src/graphnet/datasets/prometheus_datasets.py b/src/graphnet/datasets/prometheus_datasets.py index 8dc028278..3dcc555b1 100644 --- a/src/graphnet/datasets/prometheus_datasets.py +++ b/src/graphnet/datasets/prometheus_datasets.py @@ -8,7 +8,7 @@ from graphnet.training.labels import Direction, Track from graphnet.data import ERDAHostedDataset -from graphnet.data.constants import FEATURES +from graphnet.data.constants import FEATURES, TRUTH from graphnet.data.utilities import query_database @@ -18,16 +18,7 @@ class PublicPrometheusDataset(ERDAHostedDataset): # Static Member Variables: _pulsemaps = ["photons"] _truth_table = "mc_truth" - _event_truth = [ - "interaction", - "initial_state_energy", - "initial_state_type", - "initial_state_zenith", - "initial_state_azimuth", - "initial_state_x", - "initial_state_y", - "initial_state_z", - ] + _event_truth = TRUTH.PROMETHEUS _pulse_truth = None _features = FEATURES.PROMETHEUS diff --git a/tests/data/test_datamodule.py b/tests/data/test_datamodule.py index f35784644..1d159af97 100644 --- a/tests/data/test_datamodule.py +++ b/tests/data/test_datamodule.py @@ -85,7 +85,7 @@ def dataset_setup(dataset_ref: pytest.FixtureRequest) -> tuple: dataset_kwargs = { "truth_table": "mc_truth", "pulsemaps": "total", - "truth": TRUTH.PROMETHEUS, + "truth": TRUTH.PROMETHEUS_LEGACY, "features": FEATURES.PROMETHEUS, "path": data_path, "graph_definition": graph_definition,