Add autoregressive shower rollout driver
Closes the loop from single-step prediction into full showers: - giant/geometry.py + `dwarf build-geometry-oracle`: learn position -> (material, layer_id) from data (KNN/SVM) to supply the conditioning the surrogate does not predict; flag detector escape by NN distance. - giant/rollout.py: breadth-first batched frontier that steps all active tracks, spawns secondaries as new tracks, and terminates on energy cutoff, per-track max steps, escape, or natural end. Energy is deposited locally on every stop except escape (leakage), so showers conserve energy exactly. - `giant rollout` CLI: seed from real events (argmax pre_E), load checkpoint, write a world-frame steps parquet + YAML sidecar. - giant/analysis.py: compute_rollout_observables + plot_rollout_* for single-sided longitudinal/transverse/total-energy shower profiles; analysis/export_rollout_observables.py driver. - scikit-learn added as an optional `geometry` extra (lazy-imported). - Tests: tests/test_geometry.py, tests/test_rollout.py. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@@ -20,6 +20,7 @@ from scripts.bump_dataset_version import (
|
||||
run_update_manifest,
|
||||
)
|
||||
from scripts.create_root_files import run_make_root
|
||||
from scripts.geometry_oracle import run_build_geometry_oracle
|
||||
from scripts.hparam_scan import DATA_DEFAULT, SCAN_DIR_DEFAULT, run_hparam_scan
|
||||
from scripts.migrate_geant_steps import run_migration
|
||||
from scripts.steps_to_parquet import convert_steps_to_parquet
|
||||
@@ -385,6 +386,51 @@ def make_root(
|
||||
)
|
||||
|
||||
|
||||
class OracleMethod(str, Enum):
|
||||
knn = "knn"
|
||||
svm = "svm"
|
||||
|
||||
|
||||
@app.command("build-geometry-oracle")
|
||||
def build_geometry_oracle(
|
||||
data: Annotated[
|
||||
Path, typer.Argument(help="Steps parquet file or directory of steps files")
|
||||
],
|
||||
out: Annotated[
|
||||
Path, typer.Option("--out", "-o", help="Output oracle .pkl path")
|
||||
],
|
||||
method: Annotated[
|
||||
OracleMethod,
|
||||
typer.Option("--method", help="Classifier: knn (default) or svm"),
|
||||
] = OracleMethod.knn,
|
||||
k: Annotated[
|
||||
int, typer.Option("--k", help="Neighbours for the knn classifier")
|
||||
] = 1,
|
||||
subsample: Annotated[
|
||||
int,
|
||||
typer.Option("--subsample", help="Max reference points sampled from the data"),
|
||||
] = 500_000,
|
||||
escape_factor: Annotated[
|
||||
float,
|
||||
typer.Option(
|
||||
"--escape-factor",
|
||||
help="escape_threshold = this x median NN spacing of reference points",
|
||||
),
|
||||
] = 5.0,
|
||||
seed: Annotated[int, typer.Option("--seed", help="Sampling seed")] = 0,
|
||||
) -> None:
|
||||
"""Fit a position -> (material, layer_id) oracle for `giant rollout`."""
|
||||
run_build_geometry_oracle(
|
||||
data=data,
|
||||
out=out,
|
||||
method=method.value,
|
||||
k=k,
|
||||
subsample=subsample,
|
||||
escape_factor=escape_factor,
|
||||
seed=seed,
|
||||
)
|
||||
|
||||
|
||||
@app.command("hparam-scan")
|
||||
def hparam_scan(
|
||||
data: Annotated[str, typer.Option("--data")] = DATA_DEFAULT,
|
||||
|
||||
Reference in New Issue
Block a user