lars
a14a4f973a
Whitespace-only reflow (line wrapping, blank lines between defs); no logic changes. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
49 lines
1.6 KiB
Python
49 lines
1.6 KiB
Python
"""Export shower-level plots from a `giant rollout` steps parquet.
|
|
|
|
Not part of the package; run manually. Optionally overlays the real showers
|
|
seeded from the same events (a `giant predict --coord local` file) by passing a
|
|
reference path. Usage::
|
|
|
|
python analysis/export_rollout_observables.py ROLLOUT.parquet [REFERENCE_local.parquet]
|
|
"""
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import giant.analysis as a
|
|
|
|
rollout_file = sys.argv[1] if len(sys.argv) > 1 else "rollout.parquet"
|
|
reference_file = (
|
|
sys.argv[2] if len(sys.argv) > 2 and sys.argv[2] not in ("", "-") else None
|
|
)
|
|
OUT = Path(sys.argv[3]) if len(sys.argv) > 3 else Path(".")
|
|
|
|
print(f"=== computing rollout observables: {rollout_file} ===")
|
|
obs = a.compute_rollout_observables(rollout_file)
|
|
tbl = obs.event_table
|
|
print(f"n events: {len(tbl)}")
|
|
print(
|
|
f"total_edep/event: mean={tbl['total_edep'].mean():.4g} MeV "
|
|
f"leaked_E/event: mean={tbl['leaked_E'].mean():.4g} MeV "
|
|
f"n_tracks/event: mean={tbl['n_tracks'].mean():.1f} "
|
|
f"n_steps/event: mean={tbl['n_steps'].mean():.1f}"
|
|
)
|
|
|
|
reference = None
|
|
if reference_file is not None:
|
|
print(f"=== computing real reference: {reference_file} ===")
|
|
reference = a.compute_event_observables_pl(reference_file)
|
|
|
|
print("=== plots ===")
|
|
for name, fn in [
|
|
("longitudinal", a.plot_rollout_longitudinal),
|
|
("transverse", a.plot_rollout_transverse),
|
|
("total-energy", a.plot_rollout_total_energy),
|
|
]:
|
|
fig = fn(obs, reference=reference)
|
|
path = OUT / f"rollout-{name}.png"
|
|
fig.savefig(path, dpi=150, bbox_inches="tight")
|
|
print(f"wrote {path}")
|
|
|
|
print("DONE")
|