From aebcb1de991baba41efec848ef630c6c04fd0bb1 Mon Sep 17 00:00:00 2001 From: Severin Magel Date: Sun, 2 Aug 2026 00:14:52 -0400 Subject: [PATCH 1/5] Add a runnable graphnet demo on the bundled Prometheus data examples/graphnet_demo.py runs the full loop on the 50-event example file shipped in a graphnet checkout: a SQLiteDataset is adapted to the RawEvent contract via an identity detector, the geometry asset and the guaranteed-splittable event selection are built from the file itself, a tiny DeepIce (92K params) is pretrained on CPU in about a minute, and the exported checkpoint is loaded back into a stock graphnet DeepIce. fit() now selects the accelerator automatically so the demo also runs without a GPU. Co-Authored-By: Claude Fable 5 --- README.md | 11 +- examples/graphnet_demo.py | 330 ++++++++++++++++++++++++++++++++++++++ src/spine/train.py | 4 +- 3 files changed, 342 insertions(+), 3 deletions(-) create mode 100644 examples/graphnet_demo.py diff --git a/README.md b/README.md index 8037e57..4032cdd 100644 --- a/README.md +++ b/README.md @@ -40,7 +40,7 @@ src/spine/ train.py reader-agnostic fit() assembly configs/ Hydra groups: backbone/ task/ optimizer/ scheduler/ callbacks/ trainer/ data/ integrations/spine_graphnet/ graphnet integration: DeepIce backbone + reader adapter -examples/ Hydra launcher (train_curtain.py) +examples/ Hydra launcher (train_curtain.py) + runnable graphnet demo (graphnet_demo.py) tests/ core-independence gate (spine imports no graphnet) ``` @@ -65,6 +65,15 @@ python examples/train_curtain.py \ data.train_selection=train.parquet data.val_selection=val.parquet \ callbacks=curtain_auc task/objectives=v2 trainer.devices=4 ``` + +An end-to-end demo of the graphnet frame needs no data of your own: it runs on +the Prometheus example file bundled with a graphnet checkout, pretrains a tiny +DeepIce on CPU in about a minute, and loads the exported encoder back into a +stock graphnet DeepIce: +``` +python examples/graphnet_demo.py --out curtain_demo_out +``` + ## 📥 Reading data SPINE mandates **no reader**. Provide any PyTorch `Dataset` satisfying the contract stated canonically in `spine/data/datamodule.py`: diff --git a/examples/graphnet_demo.py b/examples/graphnet_demo.py new file mode 100644 index 0000000..ff10108 --- /dev/null +++ b/examples/graphnet_demo.py @@ -0,0 +1,330 @@ +"""End-to-end CURTAIN pretraining on graphnet's bundled Prometheus mock data. + +Demonstrates the graphnet frame around the spine core using the 50-event demo +file shipped inside a graphnet checkout: a graphnet SQLiteDataset is adapted +to the RawEvent contract, the geometry asset and the event selection are built +from the file itself, a tiny DeepIce is pretrained (about a minute on CPU), +and the exported checkpoint is loaded back into a stock graphnet DeepIce ready +for supervised fine-tuning. + +Usage (spine installed per the README, graphnet checkout on the import path): + python examples/graphnet_demo.py --out curtain_demo_out +""" + +from __future__ import annotations + +import argparse +import os +import sqlite3 +from collections.abc import Callable +from functools import partial + +import numpy as np +import torch +from graphnet.constants import EXAMPLE_DATA_DIR +from graphnet.data.dataset import SQLiteDataset +from graphnet.models.data_representation import EdgelessGraph, NodesAsPulses +from graphnet.models.detector import Detector +from graphnet.models.gnn import DeepIce +from spine_graphnet.deepice_backbone import DeepIceBackbone +from spine_graphnet.readers import GraphNetRawDataset +from torch.utils.data import Dataset + +from spine.data.geometry import load_geometry +from spine.data.scaling import FeatureLayout, FeatureScaler +from spine.pretext.curtain.callbacks import CurtainValAUC +from spine.pretext.curtain.objectives import OccupancyObjective +from spine.pretext.curtain.sampler import can_always_split +from spine.pretext.curtain.task import CurtainTask +from spine.train import fit + +# raw columns read from the demo file; sensor_id maps 1:1 to a position there, +# so it doubles as the data-carried sensor key +FEATURES = ["sensor_pos_x", "sensor_pos_y", "sensor_pos_z", "t", "sensor_id"] +LAYOUT = FeatureLayout() + +# the selection filter guarantees make_sample cannot raise only under the very +# same knob values the task later samples with -- single-source them +SPLIT_KNOBS = { + "holdout_mode": "temporal", + "min_visible": 8, + "min_future": 4, + "random_vis_frac": 0.5, +} + +TINY_BACKBONE = { + "d_model": 32, + "depth": 2, + "head_size": 8, + "depth_rel": 1, + "n_rel": 1, + "seq_length": 32, +} + + +def _identity(x: torch.Tensor) -> torch.Tensor: + """Return the input unchanged. + + Args: + x: Feature column values. + + Returns: + The same values. + """ + return x + + +class RawPrometheus(Detector): + """Identity detector: features stay raw; spine scales after sampling.""" + + xyz = ["sensor_pos_x", "sensor_pos_y", "sensor_pos_z"] + string_id_column = "sensor_string_id" + sensor_id_column = "sensor_id" + + def feature_map(self) -> dict[str, Callable]: + """Map every read column to the identity. + + Returns: + Identity standardization for each column in FEATURES. + """ + return {name: _identity for name in FEATURES} + + +class UnitCharge(Dataset): + """Append charge = 1 to each pulse row. + + Prometheus demo rows are single photons with no charge column; unit charge + per row completes the (x, y, z, t, charge) layout and turns the sampler's + charge-weighted mean time into the plain mean of the visible photon times. + """ + + def __init__(self, raw: Dataset): + """Wrap a RawEvent dataset whose pulses lack the charge column. + + Args: + raw: Dataset yielding RawEvents with (x, y, z, t) pulses. + """ + self.raw = raw + + def __len__(self) -> int: + return len(self.raw) + + def __getitem__(self, idx: int) -> dict: + ev = dict(self.raw[idx]) + p = ev["pulses"] + ev["pulses"] = np.concatenate([p, np.ones((len(p), 1), np.float32)], axis=1) + return ev + + +class PrometheusDemoScaler(FeatureScaler): + """Demo-detector scaling: positions span about 100 m, times about 1 us.""" + + def scale_pulses(self, x: torch.Tensor) -> torch.Tensor: + """Scale xyz and t to O(1); the unit charge passes through. + + Args: + x: [..., F] raw pulse features, columns per `self.layout`. + + Returns: + Standardized features, same shape and column order. + """ + lay = self.layout + out = x.clone() + out[..., list(lay.pos)] = x[..., list(lay.pos)] / 100.0 + out[..., lay.t] = x[..., lay.t] / 1000.0 + return out + + def scale_positions(self, p: torch.Tensor) -> torch.Tensor: + """Scale raw positions with the same factor as the pulse xyz. + + Args: + p: [..., 3] raw positions in metres. + + Returns: + Standardized coordinates on the same scale as scaled pulse xyz. + """ + return p / 100.0 + + +def build_geometry_asset(db: str, out_path: str) -> None: + """Write the per-sensor geometry asset derived from the demo file. + + The demo detector is taken to be every sensor appearing in the file. + Neighbour lists cover all other sensors sorted by distance, so the + sampler's nearest-dark lookup always finds a dark sensor. + + Args: + db: Path to the demo SQLite file. + out_path: Destination .npz path. + """ + con = sqlite3.connect(db) + rows = con.execute( + "SELECT DISTINCT sensor_id, sensor_pos_x, sensor_pos_y, sensor_pos_z " + "FROM total ORDER BY sensor_id" + ).fetchall() + con.close() + arr = np.asarray(rows, dtype=np.float64) + xyz = arr[:, 1:].astype(np.float32) + dist = np.linalg.norm(xyz[:, None] - xyz[None], axis=-1) + np.savez( + out_path, + xyz=xyz, + knn_idx=np.argsort(dist, axis=1)[:, 1:], + sensor_id=arr[:, 0].astype(np.int64), + ) + + +def make_raw_dataset(db: str, selection: list[int] | None = None) -> Dataset: + """Adapt a graphnet SQLiteDataset to the RawEvent contract. + + The identity detector plus NodesAsPulses keeps node features raw and in + FEATURES order; truth stays empty because pretraining needs no labels. + + Args: + db: Path to the demo SQLite file. + selection: Event numbers to read; None reads all. + + Returns: + RawEvent dataset with (x, y, z, t, charge) pulses and sensor keys. + """ + gn = SQLiteDataset( + path=db, + pulsemaps=["total"], + features=FEATURES, + truth=[], + truth_table="mc_truth", + data_representation=EdgelessGraph( + detector=RawPrometheus(), + node_definition=NodesAsPulses(), + input_feature_names=FEATURES, + ), + selection=selection, + ) + return UnitCharge( + GraphNetRawDataset(gn, sensor_key_index=FEATURES.index("sensor_id")) + ) + + +def build_selection( + raw: Dataset, val_frac: float, seed: int +) -> tuple[list[int], list[int]]: + """Filter to guaranteed-splittable events and split into train/val. + + Args: + raw: RawEvent dataset over the full file. + val_frac: Fraction of the kept events used for validation. + seed: Shuffle seed. + + Returns: + Train and validation event-number lists. + """ + keep = [] + for i in range(len(raw)): + ev = raw[i] + pt = ev["pulses"][:, LAYOUT.t] + if can_always_split(pt, ev["sensor_key"], **SPLIT_KNOBS): + keep.append(ev["event_no"]) + order = np.random.default_rng(seed).permutation(len(keep)) + n_val = max(1, round(val_frac * len(keep))) + val = [keep[i] for i in order[:n_val]] + train = [keep[i] for i in order[n_val:]] + return train, val + + +def parse_args() -> argparse.Namespace: + """Parse the demo's command line. + + Returns: + Parsed arguments. + """ + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument( + "--db", + default=None, + help="demo SQLite file; default: the graphnet checkout's bundled copy", + ) + ap.add_argument("--out", default="curtain_demo_out", help="output directory") + ap.add_argument("--epochs", type=int, default=15) + ap.add_argument("--seed", type=int, default=7) + return ap.parse_args() + + +def main() -> None: + """Run the demo end to end. + + Raises: + SystemExit: If the demo data is absent and no --db was given. + """ + args = parse_args() + db = args.db or os.path.join( + EXAMPLE_DATA_DIR, "sqlite", "prometheus", "prometheus-events.db" + ) + if not os.path.exists(db): + raise SystemExit( + f"demo data not found at {db} -- clone graphnet (the file ships in " + "its repo under data/examples) or pass --db" + ) + os.makedirs(args.out, exist_ok=True) + + geo_path = os.path.join(args.out, "prometheus_demo_geometry.npz") + build_geometry_asset(db, geo_path) + geo = load_geometry(geo_path, sensor_key="sensor_id") + + train_ev, val_ev = build_selection( + make_raw_dataset(db), val_frac=0.2, seed=args.seed + ) + print( + f"selection: {len(train_ev)} train / {len(val_ev)} val " + "guaranteed-splittable events", + flush=True, + ) + + task = CurtainTask( + geo=geo, + # v2 is one line more: append DtObjective(weight=1.0) from + # spine.pretext.curtain.objectives + objectives=[OccupancyObjective()], + scaler=PrometheusDemoScaler(), + dt_scale=100.0, + **SPLIT_KNOBS, + ) + backbone = DeepIceBackbone(**TINY_BACKBONE) + ckpt_path = os.path.join(args.out, "curtain_prometheus_demo.pth") + fit( + make_raw_dataset(db, selection=train_ev), + make_raw_dataset(db, selection=val_ev), + task, + backbone, + ckpt_path, + optimizer=partial(torch.optim.AdamW, lr=1e-3, weight_decay=1e-4), + batch=8, + num_workers=0, + devices=1, + max_epochs=args.epochs, + patience=10, + callbacks=[CurtainValAUC()], + ) + + # the payoff of the graphnet connection: the exported backbone weights + # drop straight into a stock graphnet DeepIce for supervised fine-tuning + ckpt = torch.load(ckpt_path, map_location="cpu") + stock = DeepIce( + hidden_dim=TINY_BACKBONE["d_model"], + depth=TINY_BACKBONE["depth"], + head_size=TINY_BACKBONE["head_size"], + depth_rel=TINY_BACKBONE["depth_rel"], + n_rel=TINY_BACKBONE["n_rel"], + seq_length=TINY_BACKBONE["seq_length"], + include_dynedge=False, + n_features=5, + ) + stock.load_state_dict(ckpt["backbone"]) + print( + f"best val loss {ckpt['val_loss']:.4f} (step {ckpt['step']}); " + f"backbone from {ckpt_path} loaded into a stock graphnet DeepIce", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/src/spine/train.py b/src/spine/train.py index 5b6d509..ef3e4e1 100644 --- a/src/spine/train.py +++ b/src/spine/train.py @@ -60,7 +60,7 @@ def fit( batch: Events per batch. num_workers: Training-loader worker processes. val_num_workers: Validation-loader workers; None uses num_workers. - devices: GPUs; more than one trains with DDP. + devices: Accelerator devices; more than one trains with DDP. precision: Lightning precision string. max_epochs: Epoch ceiling (early stopping usually ends the run). patience: EarlyStopping patience in epochs on the val loss. @@ -127,7 +127,7 @@ def fit( else "auto" ) trainer = pl.Trainer( - accelerator="gpu", + accelerator="auto", devices=devices, strategy=strategy, precision=precision, From 4e5df8c130c3f6f1898f53b84b4fb584ae49c3c4 Mon Sep 17 00:00:00 2001 From: Severin Magel Date: Sun, 2 Aug 2026 12:06:25 -0400 Subject: [PATCH 2/5] Remove the random holdout mode from the CURTAIN sampler The temporal cutoff is the pretext: every benchmarked run trained with it, and the never-used random sensor split cost a second branch in sample_event and can_always_split, two pass-through CurtainTask arguments, and a dead random_vis_frac knob in every caller's selection filter. Co-Authored-By: Claude Fable 5 --- configs/task/curtain.yaml | 2 +- examples/graphnet_demo.py | 2 - src/spine/pretext/curtain/sampler.py | 68 ++++++++-------------------- src/spine/pretext/curtain/task.py | 18 ++------ 4 files changed, 25 insertions(+), 65 deletions(-) diff --git a/configs/task/curtain.yaml b/configs/task/curtain.yaml index bf7470f..a15ed7b 100644 --- a/configs/task/curtain.yaml +++ b/configs/task/curtain.yaml @@ -1,6 +1,6 @@ # The CURTAIN pretext task -- a self-contained plugin config that pulls in its # own objectives. The launcher injects the runtime geo + scaler and -# recursive-instantiates the rest. Sampler knobs (holdout_mode, q_lo/q_hi, +# recursive-instantiates the rest. Sampler knobs (q_lo/q_hi, # pos_k, neg_anchor, rand_neg_frac, min_visible, min_future, resample_tries) # are task arguments with their defaults in CurtainTask -- override like # task.q_lo=0.5. diff --git a/examples/graphnet_demo.py b/examples/graphnet_demo.py index ff10108..2008549 100644 --- a/examples/graphnet_demo.py +++ b/examples/graphnet_demo.py @@ -46,10 +46,8 @@ # the selection filter guarantees make_sample cannot raise only under the very # same knob values the task later samples with -- single-source them SPLIT_KNOBS = { - "holdout_mode": "temporal", "min_visible": 8, "min_future": 4, - "random_vis_frac": 0.5, } TINY_BACKBONE = { diff --git a/src/spine/pretext/curtain/sampler.py b/src/spine/pretext/curtain/sampler.py index b60c4fd..0382955 100644 --- a/src/spine/pretext/curtain/sampler.py +++ b/src/spine/pretext/curtain/sampler.py @@ -86,10 +86,8 @@ def can_always_split( pt: np.ndarray, sensor_key: np.ndarray, *, - holdout_mode: str, min_visible: int, min_future: int, - random_vis_frac: float, ) -> bool: """Report whether sample_event is guaranteed to split this event. @@ -101,30 +99,20 @@ def can_always_split( Args: pt: Pulse times of the event. sensor_key: [P] integer sensor identity per pulse. - holdout_mode: "temporal" or "random". min_visible: Minimum visible sensors for a valid split. min_future: Minimum future-new sensors for a valid split. - random_vis_frac: Visible fraction of hit sensors ("random" mode only). Returns: True iff the event always yields a valid split. - - Raises: - ValueError: On an unknown `holdout_mode`. """ _, inverse = np.unique(sensor_key, return_inverse=True) n = int(inverse.max()) + 1 if len(sensor_key) else 0 if n < min_visible + min_future: return False - if holdout_mode == "temporal": - first_t = np.full(n, np.inf) - np.minimum.at(first_t, inverse, pt) - s = np.sort(first_t) - return bool(s[n - min_future] > s[min_visible - 1]) - if holdout_mode == "random": - n_vis = int(round(random_vis_frac * n)) - return n_vis >= min_visible and (n - n_vis) >= min_future - raise ValueError(f"holdout_mode {holdout_mode!r} not implemented") + first_t = np.full(n, np.inf) + np.minimum.at(first_t, inverse, pt) + s = np.sort(first_t) + return bool(s[n - min_future] > s[min_visible - 1]) def sample_event( @@ -134,8 +122,6 @@ def sample_event( rng: np.random.Generator, sensor: np.ndarray, *, - holdout_mode: str, - random_vis_frac: float, q_lo: float, q_hi: float, pos_k: int, @@ -153,8 +139,6 @@ def sample_event( geo: Geometry asset (xyz, knn_idx). rng: Generator; the cutoff and negatives re-randomize per call. sensor: [P] geometry-row index per pulse (from the data's keys). - holdout_mode: "temporal" (time cutoff) or "random" (sensor split). - random_vis_frac: Visible fraction of hit sensors ("random" mode only). q_lo: Lower bound of the cutoff-quantile window. q_hi: Upper bound of the cutoff-quantile window. pos_k: Maximum positives per event (capped by supply). @@ -169,9 +153,6 @@ def sample_event( `query_label [Q]`, `query_hard [Q]` (False = random negative), `query_dt [Q]` and `t_cwm` (charge-weighted mean visible time -- a deterministic dt reference with no future leakage). - - Raises: - ValueError: On an unknown `holdout_mode`. """ n_sensors = geo["xyz"].shape[0] knn_idx = geo["knn_idx"] @@ -182,32 +163,21 @@ def sample_event( is_dark[hit_sensors] = False dark_pool = np.flatnonzero(is_dark) - if holdout_mode == "temporal": - split = _temporal_split( - pt, - hit_sensors, - first_t, - rng, - q_lo=q_lo, - q_hi=q_hi, - min_visible=min_visible, - min_future=min_future, - resample_tries=resample_tries, - ) - if split is None: - return None - T, visible, future = split - vis_pulse_mask = pt < T - elif holdout_mode == "random": - T = float("nan") - perm = rng.permutation(hit_sensors) - n_vis = int(round(random_vis_frac * len(hit_sensors))) - if n_vis < min_visible or len(hit_sensors) - n_vis < min_future: - return None - visible, future = perm[:n_vis], perm[n_vis:] - vis_pulse_mask = np.isin(sensor, visible) - else: - raise ValueError(f"holdout_mode {holdout_mode!r} not implemented") + split = _temporal_split( + pt, + hit_sensors, + first_t, + rng, + q_lo=q_lo, + q_hi=q_hi, + min_visible=min_visible, + min_future=min_future, + resample_tries=resample_tries, + ) + if split is None: + return None + T, visible, future = split + vis_pulse_mask = pt < T if len(future) > pos_k: future = rng.choice(future, pos_k, replace=False) diff --git a/src/spine/pretext/curtain/task.py b/src/spine/pretext/curtain/task.py index d5e9c78..bdf0938 100644 --- a/src/spine/pretext/curtain/task.py +++ b/src/spine/pretext/curtain/task.py @@ -45,8 +45,6 @@ def __init__( max_pulses: int = 768, center_time: bool = True, dt_scale: float = 500.0, - holdout_mode: str = "temporal", - random_vis_frac: float = 0.5, q_lo: float = 0.3, q_hi: float = 0.7, pos_k: int = 32, @@ -65,8 +63,6 @@ def __init__( max_pulses: Cap on visible pulses fed to the encoder per event. center_time: Reference times to the charge-weighted mean. dt_scale: Divisor bringing the dt target to O(1). - holdout_mode: "temporal" (time cutoff) or "random" (sensor split). - random_vis_frac: Visible fraction ("random" mode only). q_lo: Lower bound of the cutoff-quantile window. q_hi: Upper bound of the cutoff-quantile window. pos_k: Maximum positive queries per event (capped by supply). @@ -91,8 +87,6 @@ def __init__( self.max_pulses = max_pulses self.center_time = center_time self.dt_scale = dt_scale - self.holdout_mode = holdout_mode - self.random_vis_frac = random_vis_frac self.q_lo = q_lo self.q_hi = q_hi self.pos_k = pos_k @@ -149,8 +143,6 @@ def make_sample( self.geo, rng, sensor, - holdout_mode=self.holdout_mode, - random_vis_frac=self.random_vis_frac, q_lo=self.q_lo, q_hi=self.q_hi, pos_k=self.pos_k, @@ -164,11 +156,11 @@ def make_sample( # the sampler's deterministic fallback guarantees a split for any # event with enough hit sensors, so None means under-filtered input raise ValueError( - f"event {event['event_no']} is not splittable in " - f"{self.holdout_mode!r} mode: fewer than min_visible + " - f"min_future = {self.min_visible + self.min_future} hit " - "sensors, or a degenerate first-hit-time spread; pre-filter " - "the selection with sampler.can_always_split" + f"event {event['event_no']} is not splittable: fewer than " + f"min_visible + min_future = " + f"{self.min_visible + self.min_future} hit sensors, or a " + "degenerate first-hit-time spread; pre-filter the selection " + "with sampler.can_always_split" ) vis = p[res["vis_pulse_mask"]] if self.center_time: From 7ec875d7a159dac2b064fabfbeaa31bf56f336c5 Mon Sep 17 00:00:00 2001 From: Severin Magel Date: Sun, 2 Aug 2026 12:13:44 -0400 Subject: [PATCH 3/5] Make the Prometheus demo runnable through the Hydra launcher The demo's graphnet chain moves to an importable home (spine_graphnet.prometheus: identity detector, unit-charge wrapper, and a demo_reader factory matching the data group's reader contract), and its scaler joins spine.data.scaling next to HexagonScaler. graphnet_demo.py now always stages a self-contained demo dir (db copy, geometry asset, selection parquets) and gains --prepare-only, so the same run also launches from config via the new configs/experiment pattern: python examples/graphnet_demo.py --prepare-only python examples/train_curtain.py +experiment=prometheus_demo Co-Authored-By: Claude Fable 5 --- README.md | 6 + configs/data/prometheus_demo.yaml | 11 ++ configs/experiment/prometheus_demo.yaml | 32 ++++ examples/graphnet_demo.py | 202 +++++++--------------- integrations/spine_graphnet/prometheus.py | 108 ++++++++++++ src/spine/data/scaling.py | 34 ++++ 6 files changed, 250 insertions(+), 143 deletions(-) create mode 100644 configs/data/prometheus_demo.yaml create mode 100644 configs/experiment/prometheus_demo.yaml create mode 100644 integrations/spine_graphnet/prometheus.py diff --git a/README.md b/README.md index 4032cdd..debbb4b 100644 --- a/README.md +++ b/README.md @@ -73,6 +73,12 @@ stock graphnet DeepIce: ``` python examples/graphnet_demo.py --out curtain_demo_out ``` +The same run is also available through the Hydra path: stage the demo inputs +once, then launch from the experiment config: +``` +python examples/graphnet_demo.py --prepare-only +python examples/train_curtain.py +experiment=prometheus_demo +``` ## 📥 Reading data SPINE mandates **no reader**. Provide any PyTorch `Dataset` satisfying the diff --git a/configs/data/prometheus_demo.yaml b/configs/data/prometheus_demo.yaml new file mode 100644 index 0000000..bd17bfe --- /dev/null +++ b/configs/data/prometheus_demo.yaml @@ -0,0 +1,11 @@ +# graphnet's bundled Prometheus demo file, staged locally by +# `python examples/graphnet_demo.py --prepare-only` (copies the db and writes +# geometry + selections into curtain_demo_out/). Paths are relative to the +# invocation cwd (the launcher keeps cwd via hydra.job.chdir=false). +db: curtain_demo_out/prometheus-events.db +train_selection: curtain_demo_out/train_selection.parquet +val_selection: curtain_demo_out/val_selection.parquet +n_train: 0 +reader: + _target_: spine_graphnet.prometheus.demo_reader + db: ${data.db} diff --git a/configs/experiment/prometheus_demo.yaml b/configs/experiment/prometheus_demo.yaml new file mode 100644 index 0000000..d51aed0 --- /dev/null +++ b/configs/experiment/prometheus_demo.yaml @@ -0,0 +1,32 @@ +# @package _global_ +# The bundled-Prometheus demo through the Hydra path -- same tiny run as +# examples/graphnet_demo.py. Stage its inputs once, then launch: +# python examples/graphnet_demo.py --prepare-only +# python examples/train_curtain.py +experiment=prometheus_demo +defaults: + - override /data: prometheus_demo + - override /callbacks: curtain_auc + +geo: curtain_demo_out/prometheus_demo_geometry.npz +geo_sensor_key: sensor_id +out: curtain_demo_out/curtain_prometheus_demo.pth + +scaler: + _target_: spine.data.scaling.PrometheusDemoScaler + +task: + dt_scale: 100.0 + +backbone: + d_model: 32 + depth: 2 + head_size: 8 + depth_rel: 1 + n_rel: 1 + seq_length: 32 + +trainer: + batch: 8 + num_workers: 0 + max_epochs: 15 + patience: 10 diff --git a/examples/graphnet_demo.py b/examples/graphnet_demo.py index 2008549..d985473 100644 --- a/examples/graphnet_demo.py +++ b/examples/graphnet_demo.py @@ -7,40 +7,39 @@ and the exported checkpoint is loaded back into a stock graphnet DeepIce ready for supervised fine-tuning. -Usage (spine installed per the README, graphnet checkout on the import path): - python examples/graphnet_demo.py --out curtain_demo_out +The script always stages a self-contained demo directory (db copy, geometry +asset, selection parquets), so the same run is also available through the +Hydra launcher: + python examples/graphnet_demo.py --prepare-only + python examples/train_curtain.py +experiment=prometheus_demo """ from __future__ import annotations import argparse import os +import shutil import sqlite3 -from collections.abc import Callable from functools import partial import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq import torch from graphnet.constants import EXAMPLE_DATA_DIR -from graphnet.data.dataset import SQLiteDataset -from graphnet.models.data_representation import EdgelessGraph, NodesAsPulses -from graphnet.models.detector import Detector from graphnet.models.gnn import DeepIce from spine_graphnet.deepice_backbone import DeepIceBackbone -from spine_graphnet.readers import GraphNetRawDataset +from spine_graphnet.prometheus import demo_reader from torch.utils.data import Dataset from spine.data.geometry import load_geometry -from spine.data.scaling import FeatureLayout, FeatureScaler +from spine.data.scaling import FeatureLayout, PrometheusDemoScaler from spine.pretext.curtain.callbacks import CurtainValAUC from spine.pretext.curtain.objectives import OccupancyObjective from spine.pretext.curtain.sampler import can_always_split from spine.pretext.curtain.task import CurtainTask from spine.train import fit -# raw columns read from the demo file; sensor_id maps 1:1 to a position there, -# so it doubles as the data-carried sensor key -FEATURES = ["sensor_pos_x", "sensor_pos_y", "sensor_pos_z", "t", "sensor_id"] LAYOUT = FeatureLayout() # the selection filter guarantees make_sample cannot raise only under the very @@ -60,90 +59,6 @@ } -def _identity(x: torch.Tensor) -> torch.Tensor: - """Return the input unchanged. - - Args: - x: Feature column values. - - Returns: - The same values. - """ - return x - - -class RawPrometheus(Detector): - """Identity detector: features stay raw; spine scales after sampling.""" - - xyz = ["sensor_pos_x", "sensor_pos_y", "sensor_pos_z"] - string_id_column = "sensor_string_id" - sensor_id_column = "sensor_id" - - def feature_map(self) -> dict[str, Callable]: - """Map every read column to the identity. - - Returns: - Identity standardization for each column in FEATURES. - """ - return {name: _identity for name in FEATURES} - - -class UnitCharge(Dataset): - """Append charge = 1 to each pulse row. - - Prometheus demo rows are single photons with no charge column; unit charge - per row completes the (x, y, z, t, charge) layout and turns the sampler's - charge-weighted mean time into the plain mean of the visible photon times. - """ - - def __init__(self, raw: Dataset): - """Wrap a RawEvent dataset whose pulses lack the charge column. - - Args: - raw: Dataset yielding RawEvents with (x, y, z, t) pulses. - """ - self.raw = raw - - def __len__(self) -> int: - return len(self.raw) - - def __getitem__(self, idx: int) -> dict: - ev = dict(self.raw[idx]) - p = ev["pulses"] - ev["pulses"] = np.concatenate([p, np.ones((len(p), 1), np.float32)], axis=1) - return ev - - -class PrometheusDemoScaler(FeatureScaler): - """Demo-detector scaling: positions span about 100 m, times about 1 us.""" - - def scale_pulses(self, x: torch.Tensor) -> torch.Tensor: - """Scale xyz and t to O(1); the unit charge passes through. - - Args: - x: [..., F] raw pulse features, columns per `self.layout`. - - Returns: - Standardized features, same shape and column order. - """ - lay = self.layout - out = x.clone() - out[..., list(lay.pos)] = x[..., list(lay.pos)] / 100.0 - out[..., lay.t] = x[..., lay.t] / 1000.0 - return out - - def scale_positions(self, p: torch.Tensor) -> torch.Tensor: - """Scale raw positions with the same factor as the pulse xyz. - - Args: - p: [..., 3] raw positions in metres. - - Returns: - Standardized coordinates on the same scale as scaled pulse xyz. - """ - return p / 100.0 - - def build_geometry_asset(db: str, out_path: str) -> None: """Write the per-sensor geometry asset derived from the demo file. @@ -172,37 +87,6 @@ def build_geometry_asset(db: str, out_path: str) -> None: ) -def make_raw_dataset(db: str, selection: list[int] | None = None) -> Dataset: - """Adapt a graphnet SQLiteDataset to the RawEvent contract. - - The identity detector plus NodesAsPulses keeps node features raw and in - FEATURES order; truth stays empty because pretraining needs no labels. - - Args: - db: Path to the demo SQLite file. - selection: Event numbers to read; None reads all. - - Returns: - RawEvent dataset with (x, y, z, t, charge) pulses and sensor keys. - """ - gn = SQLiteDataset( - path=db, - pulsemaps=["total"], - features=FEATURES, - truth=[], - truth_table="mc_truth", - data_representation=EdgelessGraph( - detector=RawPrometheus(), - node_definition=NodesAsPulses(), - input_feature_names=FEATURES, - ), - selection=selection, - ) - return UnitCharge( - GraphNetRawDataset(gn, sensor_key_index=FEATURES.index("sensor_id")) - ) - - def build_selection( raw: Dataset, val_frac: float, seed: int ) -> tuple[list[int], list[int]]: @@ -229,6 +113,43 @@ def build_selection( return train, val +def prepare(db: str, out: str, seed: int) -> tuple[dict, list[int], list[int]]: + """Stage the self-contained demo directory. + + Copies the demo db and writes the geometry asset plus the train/val + selection parquets, so both this script and the Hydra experiment + (+experiment=prometheus_demo) run entirely from `out`. + + Args: + db: Source demo SQLite file. + out: Demo directory to stage into. + seed: Selection shuffle seed. + + Returns: + Staged input paths and the train/val event-number lists. + """ + os.makedirs(out, exist_ok=True) + paths = { + "db": os.path.join(out, "prometheus-events.db"), + "geo": os.path.join(out, "prometheus_demo_geometry.npz"), + "train": os.path.join(out, "train_selection.parquet"), + "val": os.path.join(out, "val_selection.parquet"), + } + shutil.copyfile(db, paths["db"]) + build_geometry_asset(paths["db"], paths["geo"]) + train_ev, val_ev = build_selection( + demo_reader(paths["db"]), val_frac=0.2, seed=seed + ) + pq.write_table(pa.table({"event_no": train_ev}), paths["train"]) + pq.write_table(pa.table({"event_no": val_ev}), paths["val"]) + print( + f"selection: {len(train_ev)} train / {len(val_ev)} val " + f"guaranteed-splittable events; demo dir staged at {out}", + flush=True, + ) + return paths, train_ev, val_ev + + def parse_args() -> argparse.Namespace: """Parse the demo's command line. @@ -241,9 +162,14 @@ def parse_args() -> argparse.Namespace: default=None, help="demo SQLite file; default: the graphnet checkout's bundled copy", ) - ap.add_argument("--out", default="curtain_demo_out", help="output directory") + ap.add_argument("--out", default="curtain_demo_out", help="demo directory") ap.add_argument("--epochs", type=int, default=15) ap.add_argument("--seed", type=int, default=7) + ap.add_argument( + "--prepare-only", + action="store_true", + help="stage the demo directory and exit (for the Hydra launcher)", + ) return ap.parse_args() @@ -262,20 +188,10 @@ def main() -> None: f"demo data not found at {db} -- clone graphnet (the file ships in " "its repo under data/examples) or pass --db" ) - os.makedirs(args.out, exist_ok=True) - - geo_path = os.path.join(args.out, "prometheus_demo_geometry.npz") - build_geometry_asset(db, geo_path) - geo = load_geometry(geo_path, sensor_key="sensor_id") - - train_ev, val_ev = build_selection( - make_raw_dataset(db), val_frac=0.2, seed=args.seed - ) - print( - f"selection: {len(train_ev)} train / {len(val_ev)} val " - "guaranteed-splittable events", - flush=True, - ) + paths, train_ev, val_ev = prepare(db, args.out, args.seed) + if args.prepare_only: + return + geo = load_geometry(paths["geo"], sensor_key="sensor_id") task = CurtainTask( geo=geo, @@ -289,8 +205,8 @@ def main() -> None: backbone = DeepIceBackbone(**TINY_BACKBONE) ckpt_path = os.path.join(args.out, "curtain_prometheus_demo.pth") fit( - make_raw_dataset(db, selection=train_ev), - make_raw_dataset(db, selection=val_ev), + demo_reader(paths["db"], train_ev), + demo_reader(paths["db"], val_ev), task, backbone, ckpt_path, diff --git a/integrations/spine_graphnet/prometheus.py b/integrations/spine_graphnet/prometheus.py new file mode 100644 index 0000000..d51923e --- /dev/null +++ b/integrations/spine_graphnet/prometheus.py @@ -0,0 +1,108 @@ +"""RawEvent adapters for graphnet's bundled Prometheus demo file. + +The 50-event example file shipped in a graphnet checkout has per-photon rows +with no charge column and a globally unique `sensor_id` that serves as the +data-carried sensor key. `demo_reader` builds the full reader chain over it; +it backs both the plain demo script (examples/graphnet_demo.py) and the Hydra +data group (configs/data/prometheus_demo.yaml). +""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence + +import numpy as np +import torch +from graphnet.data.dataset import SQLiteDataset +from graphnet.models.data_representation import EdgelessGraph, NodesAsPulses +from graphnet.models.detector import Detector +from torch.utils.data import Dataset + +from spine_graphnet.readers import GraphNetRawDataset + +FEATURES = ["sensor_pos_x", "sensor_pos_y", "sensor_pos_z", "t", "sensor_id"] + + +def _identity(x: torch.Tensor) -> torch.Tensor: + """Return the input unchanged. + + Args: + x: Feature column values. + + Returns: + The same values. + """ + return x + + +class RawPrometheus(Detector): + """Identity detector: features stay raw; spine scales after sampling.""" + + xyz = ["sensor_pos_x", "sensor_pos_y", "sensor_pos_z"] + string_id_column = "sensor_string_id" + sensor_id_column = "sensor_id" + + def feature_map(self) -> dict[str, Callable]: + """Map every read column to the identity. + + Returns: + Identity standardization for each column in FEATURES. + """ + return {name: _identity for name in FEATURES} + + +class UnitCharge(Dataset): + """Append charge = 1 to each pulse row. + + Prometheus demo rows are single photons with no charge column; unit charge + per row completes the (x, y, z, t, charge) layout and turns the sampler's + charge-weighted mean time into the plain mean of the visible photon times. + """ + + def __init__(self, raw: Dataset): + """Wrap a RawEvent dataset whose pulses lack the charge column. + + Args: + raw: Dataset yielding RawEvents with (x, y, z, t) pulses. + """ + self.raw = raw + + def __len__(self) -> int: + return len(self.raw) + + def __getitem__(self, idx: int) -> dict: + ev = dict(self.raw[idx]) + p = ev["pulses"] + ev["pulses"] = np.concatenate([p, np.ones((len(p), 1), np.float32)], axis=1) + return ev + + +def demo_reader(db: str, event_nos: Sequence[int] | None = None) -> Dataset: + """Build the RawEvent reader chain over the Prometheus demo file. + + The identity detector plus NodesAsPulses keeps node features raw and in + FEATURES order; truth stays empty because pretraining needs no labels. + + Args: + db: Path to the demo SQLite file. + event_nos: Event numbers to read; None reads all. + + Returns: + RawEvent dataset with (x, y, z, t, charge) pulses and sensor keys. + """ + gn = SQLiteDataset( + path=db, + pulsemaps=["total"], + features=FEATURES, + truth=[], + truth_table="mc_truth", + data_representation=EdgelessGraph( + detector=RawPrometheus(), + node_definition=NodesAsPulses(), + input_feature_names=FEATURES, + ), + selection=None if event_nos is None else [int(e) for e in event_nos], + ) + return UnitCharge( + GraphNetRawDataset(gn, sensor_key_index=FEATURES.index("sensor_id")) + ) diff --git a/src/spine/data/scaling.py b/src/spine/data/scaling.py index e2dbad2..14df854 100644 --- a/src/spine/data/scaling.py +++ b/src/spine/data/scaling.py @@ -111,3 +111,37 @@ def scale_positions(self, p: Tensor) -> Tensor: Standardized coordinates on the same scale as scaled pulse xyz. """ return p / self._POS.to(p.device) + + +class PrometheusDemoScaler(FeatureScaler): + """Scaling for graphnet's bundled Prometheus demo file. + + That detector spans about 100 m and its photon times about 1 us; the + per-photon unit charge passes through unscaled. + """ + + def scale_pulses(self, x: Tensor) -> Tensor: + """Scale xyz and t to O(1); the charge column passes through. + + Args: + x: [..., F] raw pulse features, columns per `self.layout`. + + Returns: + Standardized features, same shape and column order. + """ + lay = self.layout + out = x.clone() + out[..., list(lay.pos)] = x[..., list(lay.pos)] / 100.0 + out[..., lay.t] = x[..., lay.t] / 1000.0 + return out + + def scale_positions(self, p: Tensor) -> Tensor: + """Scale raw positions with the same factor as the pulse xyz. + + Args: + p: [..., 3] raw positions in metres. + + Returns: + Standardized coordinates on the same scale as scaled pulse xyz. + """ + return p / 100.0 From 4de892e9c09c11b643eb74b0bf467e112f34f1c7 Mon Sep 17 00:00:00 2001 From: Severin Magel Date: Sun, 2 Aug 2026 12:25:12 -0400 Subject: [PATCH 4/5] Reuse graphnet's Prometheus detector as the demo scaler DetectorScaler (spine_graphnet.scaling) turns any graphnet Detector's feature_map into a spine FeatureScaler: named pulse columns are standardized per the detector, unmapped columns (the demo's unit charge) pass through, and query positions go through the detector's xyz entries. Pretraining in the detector's own standardization keeps the encoder's feature space identical to what a downstream graphnet fine-tune applies; the demo-specific PrometheusDemoScaler leaves the core. Co-Authored-By: Claude Fable 5 --- configs/experiment/prometheus_demo.yaml | 7 ++- examples/graphnet_demo.py | 12 ++++- integrations/spine_graphnet/scaling.py | 67 +++++++++++++++++++++++++ src/spine/data/scaling.py | 34 ------------- 4 files changed, 83 insertions(+), 37 deletions(-) create mode 100644 integrations/spine_graphnet/scaling.py diff --git a/configs/experiment/prometheus_demo.yaml b/configs/experiment/prometheus_demo.yaml index d51aed0..3a21b80 100644 --- a/configs/experiment/prometheus_demo.yaml +++ b/configs/experiment/prometheus_demo.yaml @@ -11,8 +11,13 @@ geo: curtain_demo_out/prometheus_demo_geometry.npz geo_sensor_key: sensor_id out: curtain_demo_out/curtain_prometheus_demo.pth +# graphnet's own detector for this file supplies the standardization, so the +# encoder pretrained here matches a downstream graphnet fine-tune's inputs scaler: - _target_: spine.data.scaling.PrometheusDemoScaler + _target_: spine_graphnet.scaling.DetectorScaler + detector: + _target_: graphnet.models.detector.prometheus.Prometheus + feature_names: [sensor_pos_x, sensor_pos_y, sensor_pos_z, t, charge] task: dt_scale: 100.0 diff --git a/examples/graphnet_demo.py b/examples/graphnet_demo.py index d985473..a09422e 100644 --- a/examples/graphnet_demo.py +++ b/examples/graphnet_demo.py @@ -27,13 +27,15 @@ import pyarrow.parquet as pq import torch from graphnet.constants import EXAMPLE_DATA_DIR +from graphnet.models.detector.prometheus import Prometheus from graphnet.models.gnn import DeepIce from spine_graphnet.deepice_backbone import DeepIceBackbone from spine_graphnet.prometheus import demo_reader +from spine_graphnet.scaling import DetectorScaler from torch.utils.data import Dataset from spine.data.geometry import load_geometry -from spine.data.scaling import FeatureLayout, PrometheusDemoScaler +from spine.data.scaling import FeatureLayout from spine.pretext.curtain.callbacks import CurtainValAUC from spine.pretext.curtain.objectives import OccupancyObjective from spine.pretext.curtain.sampler import can_always_split @@ -42,6 +44,12 @@ LAYOUT = FeatureLayout() +# pulse columns after the reader (sensor_id swapped for unit charge); scaling +# reuses graphnet's own Prometheus detector -- the one graphnet's examples +# pair with this file -- so pretraining runs in the same feature space a +# downstream graphnet fine-tune applies +PULSE_FEATURES = ["sensor_pos_x", "sensor_pos_y", "sensor_pos_z", "t", "charge"] + # the selection filter guarantees make_sample cannot raise only under the very # same knob values the task later samples with -- single-source them SPLIT_KNOBS = { @@ -198,7 +206,7 @@ def main() -> None: # v2 is one line more: append DtObjective(weight=1.0) from # spine.pretext.curtain.objectives objectives=[OccupancyObjective()], - scaler=PrometheusDemoScaler(), + scaler=DetectorScaler(Prometheus(), PULSE_FEATURES), dt_scale=100.0, **SPLIT_KNOBS, ) diff --git a/integrations/spine_graphnet/scaling.py b/integrations/spine_graphnet/scaling.py new file mode 100644 index 0000000..ead67c5 --- /dev/null +++ b/integrations/spine_graphnet/scaling.py @@ -0,0 +1,67 @@ +"""graphnet Detector -> spine FeatureScaler bridge.""" + +from __future__ import annotations + +from graphnet.models.detector import Detector +from torch import Tensor + +from spine.data.scaling import FeatureLayout, FeatureScaler + + +class DetectorScaler(FeatureScaler): + """Scale features with a graphnet Detector's feature_map. + + Pretraining in the detector's own standardization keeps the encoder's + feature space identical to what a downstream graphnet StandardModel + applies at fine-tune time. Columns without a feature_map entry (e.g. a + synthetic charge) pass through unscaled; query positions go through the + detector's xyz entries so they stay on the pulse coordinate scale. + """ + + def __init__( + self, + detector: Detector, + feature_names: list[str], + layout: FeatureLayout | None = None, + ): + """Bind the detector's standardization to the pulse columns. + + Args: + detector: graphnet Detector supplying feature_map() and xyz. + feature_names: Name per raw pulse column, in column order. + layout: Column layout; None uses the default order. + """ + super().__init__(layout) + self.detector = detector + self.feature_names = feature_names + + def scale_pulses(self, x: Tensor) -> Tensor: + """Apply the detector's per-column standardization. + + Args: + x: [..., F] raw pulse features, columns per `self.layout`. + + Returns: + Standardized features, same shape and column order. + """ + fm = self.detector.feature_map() + out = x.clone() + for j, name in enumerate(self.feature_names): + if name in fm: + out[..., j] = fm[name](x[..., j]) + return out + + def scale_positions(self, p: Tensor) -> Tensor: + """Apply the detector's xyz standardization to raw positions. + + Args: + p: [..., 3] raw positions, same units as the pulse xyz columns. + + Returns: + Standardized coordinates on the same scale as scaled pulse xyz. + """ + fm = self.detector.feature_map() + out = p.clone() + for k, name in enumerate(self.detector.xyz): + out[..., k] = fm[name](p[..., k]) + return out diff --git a/src/spine/data/scaling.py b/src/spine/data/scaling.py index 14df854..e2dbad2 100644 --- a/src/spine/data/scaling.py +++ b/src/spine/data/scaling.py @@ -111,37 +111,3 @@ def scale_positions(self, p: Tensor) -> Tensor: Standardized coordinates on the same scale as scaled pulse xyz. """ return p / self._POS.to(p.device) - - -class PrometheusDemoScaler(FeatureScaler): - """Scaling for graphnet's bundled Prometheus demo file. - - That detector spans about 100 m and its photon times about 1 us; the - per-photon unit charge passes through unscaled. - """ - - def scale_pulses(self, x: Tensor) -> Tensor: - """Scale xyz and t to O(1); the charge column passes through. - - Args: - x: [..., F] raw pulse features, columns per `self.layout`. - - Returns: - Standardized features, same shape and column order. - """ - lay = self.layout - out = x.clone() - out[..., list(lay.pos)] = x[..., list(lay.pos)] / 100.0 - out[..., lay.t] = x[..., lay.t] / 1000.0 - return out - - def scale_positions(self, p: Tensor) -> Tensor: - """Scale raw positions with the same factor as the pulse xyz. - - Args: - p: [..., 3] raw positions in metres. - - Returns: - Standardized coordinates on the same scale as scaled pulse xyz. - """ - return p / 100.0 From da24ed68e742c5946b3ec87502730f09b619299e Mon Sep 17 00:00:00 2001 From: Severin Magel Date: Sun, 2 Aug 2026 12:36:15 -0400 Subject: [PATCH 5/5] Keep all demo code in the example file spine_graphnet.prometheus is gone: the stock Prometheus detector already reads raw through replace_with_identity (it adds identity entries even for columns its feature_map lacks, e.g. sensor_id), and the Hydra data group can target graphnet_demo.demo_reader directly because running examples/train_curtain.py puts examples/ on sys.path. The integration package keeps only general surface: backbone, reader adapter, DetectorScaler. Co-Authored-By: Claude Fable 5 --- configs/data/prometheus_demo.yaml | 5 +- examples/graphnet_demo.py | 69 +++++++++++++- integrations/spine_graphnet/prometheus.py | 108 ---------------------- 3 files changed, 72 insertions(+), 110 deletions(-) delete mode 100644 integrations/spine_graphnet/prometheus.py diff --git a/configs/data/prometheus_demo.yaml b/configs/data/prometheus_demo.yaml index bd17bfe..d8ede06 100644 --- a/configs/data/prometheus_demo.yaml +++ b/configs/data/prometheus_demo.yaml @@ -7,5 +7,8 @@ train_selection: curtain_demo_out/train_selection.parquet val_selection: curtain_demo_out/val_selection.parquet n_train: 0 reader: - _target_: spine_graphnet.prometheus.demo_reader + # the demo module next to the launcher: `python examples/train_curtain.py` + # puts examples/ on sys.path, so the demo's reader resolves without living + # in an installed package + _target_: graphnet_demo.demo_reader db: ${data.db} diff --git a/examples/graphnet_demo.py b/examples/graphnet_demo.py index a09422e..3ebac10 100644 --- a/examples/graphnet_demo.py +++ b/examples/graphnet_demo.py @@ -20,6 +20,7 @@ import os import shutil import sqlite3 +from collections.abc import Sequence from functools import partial import numpy as np @@ -27,10 +28,12 @@ import pyarrow.parquet as pq import torch from graphnet.constants import EXAMPLE_DATA_DIR +from graphnet.data.dataset import SQLiteDataset +from graphnet.models.data_representation import EdgelessGraph, NodesAsPulses from graphnet.models.detector.prometheus import Prometheus from graphnet.models.gnn import DeepIce from spine_graphnet.deepice_backbone import DeepIceBackbone -from spine_graphnet.prometheus import demo_reader +from spine_graphnet.readers import GraphNetRawDataset from spine_graphnet.scaling import DetectorScaler from torch.utils.data import Dataset @@ -44,6 +47,10 @@ LAYOUT = FeatureLayout() +# raw columns read from the demo file; sensor_id maps 1:1 to a position there, +# so it doubles as the data-carried sensor key +FEATURES = ["sensor_pos_x", "sensor_pos_y", "sensor_pos_z", "t", "sensor_id"] + # pulse columns after the reader (sensor_id swapped for unit charge); scaling # reuses graphnet's own Prometheus detector -- the one graphnet's examples # pair with this file -- so pretraining runs in the same feature space a @@ -67,6 +74,66 @@ } +class UnitCharge(Dataset): + """Append charge = 1 to each pulse row. + + Prometheus demo rows are single photons with no charge column; unit charge + per row completes the (x, y, z, t, charge) layout and turns the sampler's + charge-weighted mean time into the plain mean of the visible photon times. + """ + + def __init__(self, raw: Dataset): + """Wrap a RawEvent dataset whose pulses lack the charge column. + + Args: + raw: Dataset yielding RawEvents with (x, y, z, t) pulses. + """ + self.raw = raw + + def __len__(self) -> int: + return len(self.raw) + + def __getitem__(self, idx: int) -> dict: + ev = dict(self.raw[idx]) + p = ev["pulses"] + ev["pulses"] = np.concatenate([p, np.ones((len(p), 1), np.float32)], axis=1) + return ev + + +def demo_reader(db: str, event_nos: Sequence[int] | None = None) -> Dataset: + """Build the RawEvent reader chain over the Prometheus demo file. + + The stock Prometheus detector reads every column through the identity + (replace_with_identity covers sensor_id too), keeping node features raw + and in FEATURES order for the sampler; truth stays empty because + pretraining needs no labels. The Hydra data group targets this function + (configs/data/prometheus_demo.yaml). + + Args: + db: Path to the demo SQLite file. + event_nos: Event numbers to read; None reads all. + + Returns: + RawEvent dataset with (x, y, z, t, charge) pulses and sensor keys. + """ + gn = SQLiteDataset( + path=db, + pulsemaps=["total"], + features=FEATURES, + truth=[], + truth_table="mc_truth", + data_representation=EdgelessGraph( + detector=Prometheus(replace_with_identity=list(FEATURES)), + node_definition=NodesAsPulses(), + input_feature_names=FEATURES, + ), + selection=None if event_nos is None else [int(e) for e in event_nos], + ) + return UnitCharge( + GraphNetRawDataset(gn, sensor_key_index=FEATURES.index("sensor_id")) + ) + + def build_geometry_asset(db: str, out_path: str) -> None: """Write the per-sensor geometry asset derived from the demo file. diff --git a/integrations/spine_graphnet/prometheus.py b/integrations/spine_graphnet/prometheus.py deleted file mode 100644 index d51923e..0000000 --- a/integrations/spine_graphnet/prometheus.py +++ /dev/null @@ -1,108 +0,0 @@ -"""RawEvent adapters for graphnet's bundled Prometheus demo file. - -The 50-event example file shipped in a graphnet checkout has per-photon rows -with no charge column and a globally unique `sensor_id` that serves as the -data-carried sensor key. `demo_reader` builds the full reader chain over it; -it backs both the plain demo script (examples/graphnet_demo.py) and the Hydra -data group (configs/data/prometheus_demo.yaml). -""" - -from __future__ import annotations - -from collections.abc import Callable, Sequence - -import numpy as np -import torch -from graphnet.data.dataset import SQLiteDataset -from graphnet.models.data_representation import EdgelessGraph, NodesAsPulses -from graphnet.models.detector import Detector -from torch.utils.data import Dataset - -from spine_graphnet.readers import GraphNetRawDataset - -FEATURES = ["sensor_pos_x", "sensor_pos_y", "sensor_pos_z", "t", "sensor_id"] - - -def _identity(x: torch.Tensor) -> torch.Tensor: - """Return the input unchanged. - - Args: - x: Feature column values. - - Returns: - The same values. - """ - return x - - -class RawPrometheus(Detector): - """Identity detector: features stay raw; spine scales after sampling.""" - - xyz = ["sensor_pos_x", "sensor_pos_y", "sensor_pos_z"] - string_id_column = "sensor_string_id" - sensor_id_column = "sensor_id" - - def feature_map(self) -> dict[str, Callable]: - """Map every read column to the identity. - - Returns: - Identity standardization for each column in FEATURES. - """ - return {name: _identity for name in FEATURES} - - -class UnitCharge(Dataset): - """Append charge = 1 to each pulse row. - - Prometheus demo rows are single photons with no charge column; unit charge - per row completes the (x, y, z, t, charge) layout and turns the sampler's - charge-weighted mean time into the plain mean of the visible photon times. - """ - - def __init__(self, raw: Dataset): - """Wrap a RawEvent dataset whose pulses lack the charge column. - - Args: - raw: Dataset yielding RawEvents with (x, y, z, t) pulses. - """ - self.raw = raw - - def __len__(self) -> int: - return len(self.raw) - - def __getitem__(self, idx: int) -> dict: - ev = dict(self.raw[idx]) - p = ev["pulses"] - ev["pulses"] = np.concatenate([p, np.ones((len(p), 1), np.float32)], axis=1) - return ev - - -def demo_reader(db: str, event_nos: Sequence[int] | None = None) -> Dataset: - """Build the RawEvent reader chain over the Prometheus demo file. - - The identity detector plus NodesAsPulses keeps node features raw and in - FEATURES order; truth stays empty because pretraining needs no labels. - - Args: - db: Path to the demo SQLite file. - event_nos: Event numbers to read; None reads all. - - Returns: - RawEvent dataset with (x, y, z, t, charge) pulses and sensor keys. - """ - gn = SQLiteDataset( - path=db, - pulsemaps=["total"], - features=FEATURES, - truth=[], - truth_table="mc_truth", - data_representation=EdgelessGraph( - detector=RawPrometheus(), - node_definition=NodesAsPulses(), - input_feature_names=FEATURES, - ), - selection=None if event_nos is None else [int(e) for e in event_nos], - ) - return UnitCharge( - GraphNetRawDataset(gn, sensor_key_index=FEATURES.index("sensor_id")) - )