lars/giant
1
0
Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Rewrite analysis as streaming rollout-vs-reference plotting pipeline
CI / Lint (ruff check) (push) Failing after 4s
Details
CI / Format (ruff format) (push) Failing after 4s
Details
CI / Type check (ty) (push) Failing after 3s
Details
CI / Tests (push) Failing after 4s
Details
CI / Bump version, build & publish wheel (push) Has been skipped
Details

Browse Source 
Replace the monolithic giant/analysis.py (predict-local + RolloutVsTruth
diagnostics) with a lean giant/analysis/ package that compares one
autoregressive `giant rollout` for a checkpoint against a held-out
miniCaloSim reference file, and generates publication-styled plots in
parallel on HTCondor.

Rollout output and a raw reference file share a world-frame physical
column subset under identical names, so the old ALR/local-frame decode
machinery is gone — everything is world-frame mm/MeV.

- sources.py: canonical LazyFrames, synthetic-termination-row filtering,
  the secondary view (rollout generation>0 tracks vs reference sec_*_list).
- reduce.py: streaming primitives — a single hist1d group_by pass, per-event
  scalars, edep-weighted depth/transverse profiles, species share, leakage.
- context.py/grouping.py: prep resolves fixed bin edges + energy/pdg/material
  group sets once into shared.json, so each compute job is one pass, no range
  scan (histogram efficiency).
- catalog.py: declarative PlotSpec registry — marginals x {overall,energy,pdg,
  material}, per-event totals, shower profiles, species/leakage, secondaries.
- render.py: the only plotstyle/LaTeX importer; PDFs + gallery metadata.
- condor.py + `giant analyze` CLI (prep/compute-one/list/render/submit):
  one job per plot, compute/render split (workers polars-only, no LaTeX).

Styling via ETPlot's plotstyle (added to the analysis extra). New tests cover
the reduce primitives, catalog id uniqueness + compute, condor submit, and a
guarded render smoke test. Delete the two predict-diagnostics notebooks.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Lars Bogner
2026-07-23 17:38:09 +02:00
parent 4f785c43e6
commit f4c2545e8b
22 changed files with 2394 additions and 4680 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CLAUDE.md
+6 -1
View File
@@ -14,6 +14,9 @@ giant train path/to/steps.parquet --mode flow # train (flow matching)
giant train path/to/steps.parquet --mode ddpm # train (DDPM baseline)
giant predict path/to/steps.parquet --checkpoint ckpt/best.pt # per-step predictions
giant rollout path/to/steps.parquet --checkpoint ckpt/best.pt --geometry oracle.pkl # full showers
giant analyze submit --rollout roll.parquet --reference test.parquet --out-dir run/ \
--accounting-group cms # parallel rollout-vs-reference analysis on HTCondor
giant analyze render --reduced-dir run/reduced --out run/plots --gallery # render PDFs + HTML gallery
dwarf --help # dataset/tooling CLI: convert, migrate, bump-gen,
# bump-schema, status, update-manifest, create-manifest,
# make-root, build-geometry-oracle, hparam-scan
@@ -54,7 +57,9 @@ GIANT is a conditional generative surrogate for the Geant4 step function. It rep
**Samplers** (`giant/sample.py`): DDPM, DDIM, and flow matching (ODE integration, ~10 steps). Flow matching is the primary mode.
**Validation** (`giant/validate.py`): step-level marginal comparisons. `giant/analysis.py` is a fully-streaming (lazy polars) diagnostics module, sized for predict/rollout files larger than RAM, with no in-memory `SampleCollection` and no full-array materialization. It covers one-step-ahead `giant predict --coord local` output (`compute_event_observables_pl` + `plot_total_energy`/`plot_longitudinal_profile`/etc. for shower-level observables, plus the marginal/correlation/constraint tiers) and, via the `RolloutVsTruth` source type, a full autoregressive `giant rollout` shower compared against held-out truth data (`compute_rollout_vs_truth_observables_pl` for shower-level observables, reusing the same plot functions) — see the module docstring.
**Validation** (`giant/validate.py`): step-level marginal comparisons.
**Analysis** (`giant/analysis/`, `giant analyze` CLI): a lean, streaming rollout-vs-reference plotting pipeline that compares one autoregressive `giant rollout` (for a given checkpoint) against a held-out miniCaloSim reference steps file, and produces publication-styled PDFs assembled into an HTML gallery. It exploits the fact that rollout output and a raw reference file share a world-frame physical column subset under identical names (`pre_*`/`post_*`/`edep`/`step_length`/`pdg`/`material`/`event_id`), so no ALR/local-frame decode is needed — everything is world-frame mm/MeV. Structure: `sources.py` (canonical LazyFrames + synthetic-termination-row filtering + the secondary view, which is `generation>0 & step_no==0` rollout tracks vs exploded `sec_*_list` reference columns), `reduce.py` (the streaming primitives — a single `hist1d` `group_by([group,bin]).len()` pass, per-event scalars, edep-weighted depth/transverse profiles, species share, leakage), `grouping.py`/`context.py` (fixed bin edges + energy-quantile/pdg/material group sets resolved once by `prep` into `shared.json`, so every compute job is one pass with no range scan), `catalog.py` (the declarative `PlotSpec` registry — marginals × {overall,energy,pdg,material}, per-event totals, shower profiles, species/leakage, secondaries), and `render.py` (the only module importing ETPlot's `plotstyle`/LaTeX; dispatches on `Reduced.kind`, writes PDFs + `metadata.yaml`). **Compute/render split:** `giant analyze submit` runs `prep` then submits one HTCondor job per plot (`compute-one`, polars/numpy only — no LaTeX on workers), each writing a small `reduced/<id>.json`; the local `giant analyze render` turns those into the styled PDF/gallery tree. See `giant/analysis/__init__.py`.
**Shower rollout** (`giant/rollout.py`, `giant rollout` CLI): autoregressively steps the two-stage model into a full shower — each primary post-step becomes the next pre-step, secondaries are pushed as new tracks, and per-step `material`/`layer_id` come from a `GeometryOracle` (`giant/geometry.py`, built via `dwarf build-geometry-oracle`) that learns position → (material, layer_id) from data and flags detector escape by nearest-neighbour distance. Tracks terminate 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 by construction.
analysis/rollout_validation.ipynb
-438
View File
@@ -1,438 +0,0 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"id": "f91460f3",
"metadata": {},
"outputs": [],
"source": [
"# Auto-reload edited modules (e.g. giant.analysis) without restarting the kernel.\n",
"%load_ext autoreload\n",
"%autoreload 2"
]
},
{
"cell_type": "markdown",
"id": "0f10da93",
"metadata": {},
"source": [
"# GIANT rollout-vs-truth validation notebook\n",
"\n",
"Diagnostics for a full autoregressive `giant rollout` shower, compared against a held-out ground-truth steps file (the same schema `giant train` consumes — see `giant.data.loader.load_steps`) rather than one-step-ahead `giant predict` output.\n",
"\n",
"This is the sibling of `validation.ipynb`: that notebook checks whether one-step generation (conditioned on the *real* preceding state, every row) reproduces real marginals/correlations/shower observables. This one checks the thing that actually matters for deployment — whether a shower **rolled out autoregressively from the model's own outputs** still looks physical, which is where covariate shift (small per-step errors compounding across a track) would show up.\n",
"\n",
"Built on `RolloutVsTruth`, which treats the rollout file as \"generated\" and the truth file as \"real\". Unlike the paired predict-parquet `source` (`pred_*`/`true_*` columns of the same row), the two files here are **independent, unpaired datasets** — a rollout doesn't replay real events row-for-row, so real/generated may have different lengths and there's no per-row correspondence. Everything below only ever compares real-vs-generated *distributions*, never individual paired rows, and every check still streams (no `SampleCollection`, no full-file materialization) — see `giant.analysis`'s module docstring for the `RolloutVsTruth` mechanics.\n",
"\n",
"Same four tiers as `validation.ipynb`, all built on the same functions — pass a `RolloutVsTruth` in place of the predict-parquet path/LazyFrame everywhere:\n",
"\n",
"1. **stratified marginals** — per-dimension real-vs-generated, sliced by pdg/material/energy\n",
"2. **joint structure** — correlation matrices, physically-coupled pairwise plots, direction alignment\n",
"3. **physical constraints** — unit-norm directions, non-negative step_length/delta_e/edep (checked on the rollout's own output — with autoregression, a constraint violation early in a track can compound into later steps, unlike one-step-ahead validation)\n",
"4. **event-level (shower) observables** — total/mean/median energy and length per event, longitudinal/transverse profiles, shower-max depth, computed directly from the rollout shower against the truth file's own events (`compute_rollout_vs_truth_observables_pl`, the Tier 4 counterpart to `RolloutVsTruth`)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "251dc1f7",
"metadata": {},
"outputs": [],
"source": [
"from giant.analysis import RolloutVsTruth, plot_kl_bars_pl\n",
"\n",
"# `giant rollout` output for the shower(s) under test.\n",
"ROLLOUT_FILE = (\n",
" \"/ceph/lbogner/geant_steps/predictions/9e76bc2c-f4ef-4488-9f62-b6d14e1f298e.parquet\"\n",
")\n",
"# Any held-out file sharing giant train's input schema (real miniCaloSim\n",
"# steps) — e.g. the val split the rollout's seed events were drawn from.\n",
"TRUTH_FILE = (\n",
" \"/ceph/lbogner/geant_steps/processed/steps/gen3/schema2/pbwo4/shard-009.parquet\"\n",
")\n",
"\n",
"# sample_frac subsamples each side of the Tier 1-3 checks independently\n",
"# (kept memory-bounded for large files); defaults to every row. Tier 4\n",
"# (compute_rollout_vs_truth_observables_pl, below) always streams every row\n",
"# regardless — per-event sums would be silently corrupted by row subsampling.\n",
"SOURCE = RolloutVsTruth(rollout=ROLLOUT_FILE, truth=TRUTH_FILE)"
]
},
{
"cell_type": "markdown",
"id": "df4bf24b",
"metadata": {},
"source": [
"## Tier 1: stratified marginals\n",
"\n",
"KL(real || generated) per target dimension, streamed straight from both files."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "13cd0838",
"metadata": {},
"outputs": [],
"source": [
"for grouping in [None, \"energy\", \"pdg\", \"material\"]:\n",
" fig = plot_kl_bars_pl(SOURCE, group_by=grouping)\n",
" fig.show()"
]
},
{
"cell_type": "markdown",
"id": "aba92bf9",
"metadata": {},
"source": [
"## Detailed marginals (Tier 1, overlaid histograms)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "7bab2e35",
"metadata": {},
"outputs": [],
"source": [
"from giant.analysis import plot_marginals, plot_correlation_matrices, plot_pairwise\n",
"from giant.analysis import plot_direction_alignment, plot_constraint_violations\n",
"\n",
"_ = plot_marginals(SOURCE)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6c2958e7",
"metadata": {},
"outputs": [],
"source": [
"_ = plot_marginals(SOURCE, group_by=\"energy\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0fe70835",
"metadata": {},
"outputs": [],
"source": [
"_ = plot_marginals(SOURCE, group_by=\"pdg\")"
]
},
{
"cell_type": "markdown",
"id": "038dda3b",
"metadata": {},
"source": [
"## Tier 2: joint structure"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3c463d50",
"metadata": {},
"outputs": [],
"source": [
"# Real vs. generated Pearson correlation matrices (+ their difference) over\n",
"# the 9 raw target dims — catches a model that decorrelates targets that are\n",
"# physically coupled even when every individual marginal looks clean.\n",
"_ = plot_correlation_matrices(SOURCE)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "393c6845",
"metadata": {},
"outputs": [],
"source": [
"# Scatter for physically-coupled pairs (step_length/delta_e/edep) — the\n",
"# joint-structure check correlation matrices alone can't fully capture.\n",
"_ = plot_pairwise(SOURCE, n_sample=10000)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1dab5305",
"metadata": {},
"outputs": [],
"source": [
"# cos(angle) between post_dir and travel_dir — coupled through the\n",
"# scattering physics, so this is another joint-structure check.\n",
"_ = plot_direction_alignment(SOURCE)"
]
},
{
"cell_type": "markdown",
"id": "7232d3b4",
"metadata": {},
"source": [
"## Tier 3: physical constraints\n",
"\n",
"Unit-norm direction vectors, non-negative step_length/delta_e/edep. `constraint_report_pl`/`plot_constraint_violations` only ever check the *generated* side (here the rollout output) — under autoregression a violation isn't just a one-off artifact, it can feed the next step's conditioning, so this is worth watching more closely here than in one-step-ahead validation."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "86be25fa",
"metadata": {},
"outputs": [],
"source": [
"_ = plot_constraint_violations(SOURCE)"
]
},
{
"cell_type": "markdown",
"id": "ed4d3037",
"metadata": {},
"source": [
"## Tier 4: event-level (shower) observables\n",
"\n",
"Built on `compute_rollout_vs_truth_observables_pl`, not `compute_event_observables_pl` — the rollout file carries its own `track_id`/`termination_reason` columns the event-level aggregation needs, and the shower here already *is* a full autoregressive rollout rather than one-step generations re-aggregated by event. Entry axis/point and per-event totals are computed separately per side (rollout and truth events are unrelated), but depth/transverse bin edges are shared across both so the profiles below overlay on one binning.\n",
"\n",
"Returns the same `EventObservables` `compute_event_observables_pl` does, so every plot function from `validation.ipynb` works unchanged here too."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9175876a",
"metadata": {},
"outputs": [],
"source": [
"from giant.analysis import compute_rollout_vs_truth_observables_pl\n",
"from giant.analysis import plot_total_energy, plot_total_length\n",
"from giant.analysis import plot_mean_energy_per_step, plot_mean_length_per_step\n",
"from giant.analysis import plot_longitudinal_profile, plot_transverse_profile\n",
"from giant.analysis import plot_shower_max_depth\n",
"\n",
"obs = compute_rollout_vs_truth_observables_pl(ROLLOUT_FILE, TRUTH_FILE)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "83e6dd0e",
"metadata": {},
"outputs": [],
"source": [
"_ = plot_total_energy(obs)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1e35a387",
"metadata": {},
"outputs": [],
"source": [
"_ = plot_total_length(obs)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9221e682",
"metadata": {},
"outputs": [],
"source": [
"_ = plot_mean_energy_per_step(obs)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4536acb0",
"metadata": {},
"outputs": [],
"source": [
"_ = plot_mean_length_per_step(obs)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3a03c04f",
"metadata": {},
"outputs": [],
"source": [
"_ = plot_longitudinal_profile(obs)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "eae0536e",
"metadata": {},
"outputs": [],
"source": [
"_ = plot_transverse_profile(obs)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "98c221e7",
"metadata": {},
"outputs": [],
"source": [
"_ = plot_shower_max_depth(obs)"
]
},
{
"cell_type": "markdown",
"id": "208ca6e4",
"metadata": {},
"source": [
"---\n",
"\n",
"For the dataset-wide breakdown of which particle species contributed how much of the total energy/length (`pdg_contribution_table_pl`), see `validation.ipynb` — it needs the paired predict schema, which this rollout-vs-truth comparison doesn't have."
]
},
{
"cell_type": "markdown",
"id": "7fb27b941602401d91542211134fc71a",
"metadata": {},
"source": [
"## Router gating showcase (MoE)\n",
"\n",
"Every other section above is file-only — it reads `ROLLOUT_FILE` and never touches\n",
"a checkpoint (see `giant.analysis`'s module docstring). This section is the one\n",
"deliberate exception: soft gate weights only exist inside the trained `Router`,\n",
"not in the rollout parquet, so this loads the checkpoint that produced\n",
"`ROLLOUT_FILE` and calls `model.router.gate(...)` directly on that shower's\n",
"pre-step conditioning.\n",
"\n",
"`model.router` is Stage 1's router; Stage 2 (`sec_decoder.router`) is a separate,\n",
"independently trained `Router` instance over the same axis (see\n",
"`giant.model.network.build_models`) and isn't shown here.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "acae54e37e7d407bbb7b55eff062a284",
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"import polars as pl\n",
"import torch\n",
"\n",
"from giant.analysis import plot_router_gating\n",
"from giant.data.transforms import Normalizer, build_cond_features\n",
"from giant.model.network import build_models\n",
"\n",
"# Checkpoint that produced ROLLOUT_FILE (needs `model.router` enabled at\n",
"# train time, i.e. trained with `--router` / `model.router.enabled = true`).\n",
"CHECKPOINT = \"/ceph/lbogner/geant_steps/checkpoints/REPLACE_ME/best.pt\"\n",
"\n",
"ckpt = torch.load(CHECKPOINT, map_location=\"cpu\", weights_only=False)\n",
"model_cfg = ckpt[\"model_config\"]\n",
"conditioning = model_cfg.get(\"conditioning\", \"embedding\")\n",
"pdg_map = {int(k): v for k, v in ckpt[\"pdg_map\"].items()}\n",
"mat_map = {str(k): v for k, v in ckpt[\"mat_map\"].items()}\n",
"cond_norm = Normalizer.from_dict(ckpt[\"normalizer\"][\"cond\"])\n",
"\n",
"model, _sec_decoder = build_models(model_cfg)\n",
"model.load_state_dict(ckpt[\"model\"])\n",
"model.eval()\n",
"\n",
"if not hasattr(model, \"router\"):\n",
" raise RuntimeError(\n",
" f\"{CHECKPOINT} has no router — it was trained with model.router.enabled=False\"\n",
" )"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9a63283cbaf04dbcab1f6479b197f3a8",
"metadata": {},
"outputs": [],
"source": [
"# Pre-step conditioning for every row of the rollout shower, reconstructed\n",
"# the same way `giant predict`/`giant rollout` do (giant.data.transforms).\n",
"cols = [\n",
" \"pdg\",\n",
" \"pre_x\",\n",
" \"pre_y\",\n",
" \"pre_z\",\n",
" \"pre_E\",\n",
" \"pre_dx\",\n",
" \"pre_dy\",\n",
" \"pre_dz\",\n",
" \"material\",\n",
" \"layer_id\",\n",
"]\n",
"df = pl.read_parquet(ROLLOUT_FILE, columns=cols)\n",
"\n",
"# Rows whose pdg/material fell outside the training vocab can't be encoded\n",
"# (mirrors the pdg_mask filtering in `giant predict`'s CLI path).\n",
"known = df[\"pdg\"].map_elements(\n",
" lambda p: int(p) in pdg_map, return_dtype=pl.Boolean\n",
") & df[\"material\"].map_elements(lambda m: str(m) in mat_map, return_dtype=pl.Boolean)\n",
"n_dropped = (~known).sum()\n",
"if n_dropped:\n",
" print(f\"dropping {n_dropped}/{len(df)} rows with unknown pdg/material\")\n",
"df = df.filter(known)\n",
"\n",
"data = {\n",
" \"pre_pos\": df.select(\"pre_x\", \"pre_y\", \"pre_z\").to_numpy().astype(np.float32),\n",
" \"pre_E\": df[\"pre_E\"].to_numpy().astype(np.float32),\n",
" \"pre_dir\": df.select(\"pre_dx\", \"pre_dy\", \"pre_dz\").to_numpy().astype(np.float32),\n",
" \"layer_id\": df[\"layer_id\"].to_numpy(),\n",
" \"pdg\": df[\"pdg\"].to_numpy(),\n",
" \"material\": df[\"material\"].to_numpy(),\n",
"}\n",
"cond_cont, cond_cat = build_cond_features(\n",
" data, pdg_map, mat_map, cond_norm, conditioning=conditioning\n",
")\n",
"cc = torch.from_numpy(cond_cont).float()\n",
"ck = torch.from_numpy(cond_cat).long()\n",
"\n",
"with torch.no_grad():\n",
" gate_weights = model.router.gate(cc, ck).numpy() # (N, n_experts), rows sum to 1"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "8dd0d8092fe74a7c96281538738b07e2",
"metadata": {},
"outputs": [],
"source": [
"# EnergyRouter gates on pre-step energy, so that's the natural x-axis here —\n",
"# swap for a categorical plot if this checkpoint used a different router type.\n",
"_ = plot_router_gating(data[\"pre_E\"], gate_weights, x_label=\"pre_E\")"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "giant (3.12.13.final.0)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.13"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
analysis/validation.ipynb
-706
View File
File diff suppressed because one or more lines are too long
giant/analysis.py
-2359
View File
File diff suppressed because it is too large Load Diff
giant/analysis/__init__.py
+30
View File
@@ -0,0 +1,30 @@
"""Rollout-vs-reference analysis: streaming compute + plotstyle rendering.
Compares one autoregressive ``giant rollout`` against a held-out miniCaloSim
reference file, producing publication-styled comparison plots generated in
parallel on HTCondor (one job per plot, compute/render split).
Only ``render`` (and the ``render`` CLI path) imports plotstyle/LaTeX; everything
re-exported here is plotstyle-free so it runs on a compute worker. Import
``giant.analysis.render`` explicitly for the local render step.
"""
from giant.analysis.catalog import build_catalog, catalog_ids, get_spec
from giant.analysis.condor import SubmitConfig, compute_one, prep, write_submit
from giant.analysis.context import Context, build_context
from giant.analysis.reduced import Reduced
from giant.analysis.sources import Side
__all__ = [
"build_catalog",
"catalog_ids",
"get_spec",
"SubmitConfig",
"compute_one",
"prep",
"write_submit",
"Context",
"build_context",
"Reduced",
"Side",
]
giant/analysis/catalog.py
+574
View File
@@ -0,0 +1,574 @@
"""The declarative plot catalog: one ``PlotSpec`` per figure.
Each spec knows its stable ``id`` (used for the reduced-data filename, the PDF
stem and the condor queue item), its gallery ``family`` (subdirectory), and a
``compute(bundle) -> Reduced`` that runs the streaming reduction. Rendering lives
in ``render.py`` and dispatches on ``Reduced.kind`` — the catalog itself never
imports plotstyle, so ``compute-one`` jobs stay LaTeX-free.
The registry is built by expanding parametric families (marginals over
variable x grouping, secondaries, ...) into concrete specs.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Callable
import numpy as np
import polars as pl
from giant.analysis.context import Context
from giant.analysis.grouping import (
energy_bin_labels,
event_energy_bins,
material_label,
pdg_label,
)
from giant.analysis.reduce import (
attach_entry_axis,
depth_expr,
entry_axis,
event_scalars,
hist1d,
leakage_fraction,
species_share,
transverse_expr,
weighted_profile,
)
from giant.analysis.reduced import Reduced
from giant.analysis.sources import Side, open_side, physical_steps, secondaries
from giant.analysis.variables import RANGED_VARS, cos_scatter_expr
@dataclass
class Bundle:
"""Everything a compute runs against — built once per ``compute-one`` job."""
ctx: Context
r_all: pl.LazyFrame # rollout, all rows (incl. synthetic termination rows)
t_all: pl.LazyFrame # reference, all rows
r_phys: pl.LazyFrame # rollout, physical steps only
t_phys: pl.LazyFrame # reference, physical steps only
@classmethod
def open(cls, rollout, reference, ctx: Context) -> "Bundle":
r_all = open_side(rollout, Side.rollout)
t_all = open_side(reference, Side.reference)
return cls(
ctx=ctx,
r_all=r_all,
t_all=t_all,
r_phys=physical_steps(r_all, Side.rollout),
t_phys=physical_steps(t_all, Side.reference),
)
@dataclass
class PlotSpec:
id: str
family: str
compute: Callable[[Bundle], Reduced]
# ---------------------------------------------------------------------------
# small numpy/hist helpers
# ---------------------------------------------------------------------------
_ROLL = "rollout"
_REF = "reference"
def _counts(h: dict, key, nbins: int) -> list[int]:
return h.get(key, np.zeros(nbins, dtype=np.int64)).astype(np.int64).tolist()
def _np_hist_pair(
r: np.ndarray, t: np.ndarray, nbins: int
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""Shared-edge histogram of two small per-event arrays (robust range)."""
both = np.concatenate([r, t]) if (len(r) or len(t)) else np.array([0.0, 1.0])
lo, hi = float(np.quantile(both, 0.001)), float(np.quantile(both, 0.999))
if not (hi - lo > 1e-6 * max(abs(hi), 1.0)):
lo, hi = lo - 0.5, hi + 0.5
edges = np.linspace(lo, hi, nbins + 1)
return edges, np.histogram(r, edges)[0], np.histogram(t, edges)[0]
# Human-readable figure titles per marginal variable (the axis labels carry units;
# these read cleanly as a title without them).
_TITLE_NAMES = {
"step_length": "Step length",
"edep": "Deposited energy per step",
"delta_e": "Energy loss per step",
"post_E": "Post-step energy",
"cos_scatter": "Scattering cosine",
}
def _var(var: str):
"""(axis label, value expr) for a marginal variable name."""
if var == "cos_scatter":
return ("cos of scattering angle", cos_scatter_expr())
label, expr = RANGED_VARS[var]
return (label, expr)
def _marginal_edges(ctx: Context, var: str) -> np.ndarray:
if var == "cos_scatter":
return np.linspace(-1.0, 1.0, ctx.n_marginal_bins + 1)
return ctx.marginal_edges(var)
# ---------------------------------------------------------------------------
# marginals: variable x {overall, energy, pdg, material}
# ---------------------------------------------------------------------------
def _marginal_overall(b: Bundle, var: str) -> Reduced:
label, expr = _var(var)
edges = _marginal_edges(b.ctx, var)
nb = len(edges) - 1
r = hist1d(b.r_phys, expr, edges)
t = hist1d(b.t_phys, expr, edges)
return Reduced(
id=f"marginal_{var}",
family="marginals",
kind="overlay_hist",
title=_TITLE_NAMES[var],
xlabel=label,
payload={
"edges": edges.tolist(),
_ROLL: _counts(r, 0, nb),
_REF: _counts(t, 0, nb),
"log_y": True,
},
)
def _energy_group_expr(lf: pl.LazyFrame, edges: np.ndarray) -> pl.Expr:
ids, bins = event_energy_bins(lf, edges)
return pl.col("event_id").replace_strict(
ids, bins, default=-1, return_dtype=pl.Int64
)
def _marginal_grouped(b: Bundle, var: str, axis: str) -> Reduced:
label, expr = _var(var)
edges = _marginal_edges(b.ctx, var)
nb = len(edges) - 1
groups: dict[str, dict] = {}
if axis == "pdg":
r = hist1d(b.r_phys, expr, edges, group=pl.col("pdg"))
t = hist1d(b.t_phys, expr, edges, group=pl.col("pdg"))
for k in b.ctx.top_pdgs:
groups[pdg_label(k)] = {_ROLL: _counts(r, k, nb), _REF: _counts(t, k, nb)}
elif axis == "material":
r = hist1d(b.r_phys, expr, edges, group=pl.col("material"))
t = hist1d(b.t_phys, expr, edges, group=pl.col("material"))
for m in b.ctx.materials:
groups[material_label(m)] = {
_ROLL: _counts(r, m, nb),
_REF: _counts(t, m, nb),
}
else: # energy
e_edges = np.asarray(b.ctx.energy_edges)
r = hist1d(b.r_phys, expr, edges, group=_energy_group_expr(b.r_phys, e_edges))
t = hist1d(b.t_phys, expr, edges, group=_energy_group_expr(b.t_phys, e_edges))
for bi, lbl in enumerate(energy_bin_labels(e_edges)):
groups[lbl] = {_ROLL: _counts(r, bi, nb), _REF: _counts(t, bi, nb)}
return Reduced(
id=f"marginal_{var}_by_{axis}",
family="marginals",
kind="grouped_hist",
title=f"{_TITLE_NAMES[var]} by {axis}",
xlabel=label,
payload={"edges": edges.tolist(), "groups": groups, "log_y": True},
)
# ---------------------------------------------------------------------------
# per-event scalar observables
# ---------------------------------------------------------------------------
def _event_scalar(
b: Bundle, spec_id: str, title: str, xlabel: str, col: str, use_all: bool
) -> Reduced:
r_lf, t_lf = (b.r_all, b.t_all) if use_all else (b.r_phys, b.t_phys)
r = event_scalars(r_lf)[col].to_numpy()
t = event_scalars(t_lf)[col].to_numpy()
edges, rc, tc = _np_hist_pair(r, t, b.ctx.n_marginal_bins)
return Reduced(
id=spec_id,
family="event",
kind="overlay_hist",
title=title,
xlabel=xlabel,
payload={
"edges": edges.tolist(),
_ROLL: rc.astype(np.int64).tolist(),
_REF: tc.astype(np.int64).tolist(),
"log_y": False,
},
)
def _event_total_edep_by_energy(b: Bundle) -> Reduced:
e_edges = np.asarray(b.ctx.energy_edges)
r = event_scalars(b.r_all)
t = event_scalars(b.t_all)
r_bin = np.clip(
np.digitize(r["incident_E"].to_numpy(), e_edges[1:-1]), 0, len(e_edges) - 2
)
t_bin = np.clip(
np.digitize(t["incident_E"].to_numpy(), e_edges[1:-1]), 0, len(e_edges) - 2
)
r_val, t_val = r["total_edep"].to_numpy(), t["total_edep"].to_numpy()
edges, _, _ = _np_hist_pair(r_val, t_val, b.ctx.n_marginal_bins)
groups: dict[str, dict] = {}
for bi, lbl in enumerate(energy_bin_labels(e_edges)):
rc = np.histogram(r_val[r_bin == bi], edges)[0]
tc = np.histogram(t_val[t_bin == bi], edges)[0]
groups[lbl] = {
_ROLL: rc.astype(np.int64).tolist(),
_REF: tc.astype(np.int64).tolist(),
}
return Reduced(
id="event_total_edep_by_energy",
family="event",
kind="grouped_hist",
title="Total deposited energy per event by incident energy",
xlabel="total deposited energy [MeV]",
payload={"edges": edges.tolist(), "groups": groups, "log_y": False},
)
# ---------------------------------------------------------------------------
# shower shape profiles
# ---------------------------------------------------------------------------
def _profile(
b: Bundle, spec_id: str, title: str, xlabel: str, coord_fn, edges_key: str
) -> Reduced:
edges = np.asarray(getattr(b.ctx, edges_key))
r_ea, t_ea = entry_axis(b.r_all), entry_axis(b.t_all)
r_lf = attach_entry_axis(b.r_all, r_ea)
t_lf = attach_entry_axis(b.t_all, t_ea)
r_mean, r_std = weighted_profile(r_lf, coord_fn(), edges, pl.col("edep"))
t_mean, t_std = weighted_profile(t_lf, coord_fn(), edges, pl.col("edep"))
return Reduced(
id=spec_id,
family="shower",
kind="profile",
title=title,
xlabel=xlabel,
payload={
"edges": edges.tolist(),
"rollout_mean": r_mean.tolist(),
"rollout_std": r_std.tolist(),
"reference_mean": t_mean.tolist(),
"reference_std": t_std.tolist(),
"ylabel": "mean deposited energy per event [MeV]",
},
)
# ---------------------------------------------------------------------------
# species share + leakage
# ---------------------------------------------------------------------------
def _species_share(b: Bundle) -> Reduced:
r = species_share(b.r_all)
t = species_share(b.t_all)
r_map = dict(zip(r["pdg"].to_list(), r["total_edep"].to_list()))
t_map = dict(zip(t["pdg"].to_list(), t["total_edep"].to_list()))
r_tot = sum(r_map.values()) or 1.0
t_tot = sum(t_map.values()) or 1.0
labels = [pdg_label(k) for k in b.ctx.top_pdgs]
return Reduced(
id="species_edep_share",
family="species",
kind="bar",
title="Deposited-energy share by species",
xlabel="species",
payload={
"labels": labels,
_ROLL: [r_map.get(k, 0.0) / r_tot for k in b.ctx.top_pdgs],
_REF: [t_map.get(k, 0.0) / t_tot for k in b.ctx.top_pdgs],
"ylabel": "fraction of total deposited energy",
},
)
def _leakage(b: Bundle) -> Reduced:
frac = leakage_fraction(b.r_all)
edges = np.linspace(
0.0,
max(float(frac.max()) if len(frac) else 1.0, 1e-3),
b.ctx.n_marginal_bins + 1,
)
counts = np.histogram(frac, edges)[0]
return Reduced(
id="leakage_fraction",
family="species",
kind="single_hist",
title="Escaped (leakage) energy fraction per shower",
xlabel="escaped energy fraction",
payload={
"edges": edges.tolist(),
_ROLL: counts.astype(np.int64).tolist(),
"log_y": True,
"note": "rollout only; the reference has no detector-escape concept",
},
)
# ---------------------------------------------------------------------------
# secondaries
# ---------------------------------------------------------------------------
def _sec_frames(b: Bundle):
return (
secondaries(b.r_phys, Side.rollout),
secondaries(b.t_all, Side.reference),
)
def _sec_count_per_event(b: Bundle) -> Reduced:
r_sec, t_sec = _sec_frames(b)
r = (
r_sec.group_by("event_id")
.agg(pl.len().alias("n"))
.collect(engine="streaming")["n"]
.to_numpy()
)
t = (
t_sec.group_by("event_id")
.agg(pl.len().alias("n"))
.collect(engine="streaming")["n"]
.to_numpy()
)
edges, rc, tc = _np_hist_pair(
r.astype(float), t.astype(float), min(b.ctx.n_marginal_bins, 40)
)
return Reduced(
id="sec_count_per_event",
family="secondaries",
kind="overlay_hist",
title="Number of secondaries per event",
xlabel="secondaries per event",
payload={
"edges": edges.tolist(),
_ROLL: rc.astype(np.int64).tolist(),
_REF: tc.astype(np.int64).tolist(),
"log_y": False,
},
)
def _sec_count_per_species(b: Bundle) -> Reduced:
r_sec, t_sec = _sec_frames(b)
r = dict(
zip(
*[
r_sec.group_by("pdg")
.agg(pl.len().alias("n"))
.collect(engine="streaming")[c]
.to_list()
for c in ("pdg", "n")
]
)
)
t = dict(
zip(
*[
t_sec.group_by("pdg")
.agg(pl.len().alias("n"))
.collect(engine="streaming")[c]
.to_list()
for c in ("pdg", "n")
]
)
)
keys = sorted(set(r) | set(t), key=lambda k: -(r.get(k, 0) + t.get(k, 0)))[
: len(b.ctx.top_pdgs)
]
return Reduced(
id="sec_count_per_species",
family="secondaries",
kind="bar",
title="Secondary count by species",
xlabel="species",
payload={
"labels": [pdg_label(k) for k in keys],
_ROLL: [float(r.get(k, 0)) for k in keys],
_REF: [float(t.get(k, 0)) for k in keys],
"ylabel": "secondary count",
},
)
def _sec_energy(b: Bundle) -> Reduced:
r_sec, t_sec = _sec_frames(b)
edges = np.linspace(*b.ctx.sec_energy_range, b.ctx.n_sec_bins + 1)
r = hist1d(r_sec, pl.col("energy"), edges)
t = hist1d(t_sec, pl.col("energy"), edges)
nb = len(edges) - 1
return Reduced(
id="sec_energy",
family="secondaries",
kind="overlay_hist",
title="Secondary birth energy",
xlabel="secondary energy [MeV]",
payload={
"edges": edges.tolist(),
_ROLL: _counts(r, 0, nb),
_REF: _counts(t, 0, nb),
"log_y": True,
},
)
def _sec_cos_angle(b: Bundle) -> Reduced:
edges = np.linspace(-1.0, 1.0, b.ctx.n_sec_bins + 1)
nb = len(edges) - 1
cos = (
pl.col("sdx") * pl.col("axis_x")
+ pl.col("sdy") * pl.col("axis_y")
+ pl.col("sdz") * pl.col("axis_z")
).clip(-1.0, 1.0)
def _side(sec_lf: pl.LazyFrame, steps_lf: pl.LazyFrame) -> list[int]:
ea = entry_axis(steps_lf)
return _counts(hist1d(attach_entry_axis(sec_lf, ea), cos, edges), 0, nb)
r_sec, t_sec = _sec_frames(b)
return Reduced(
id="sec_cos_angle",
family="secondaries",
kind="overlay_hist",
title="Secondary emission angle relative to the shower axis",
xlabel="cos of emission angle",
payload={
"edges": edges.tolist(),
_ROLL: _side(r_sec, b.r_phys),
_REF: _side(t_sec, b.t_all),
"log_y": False,
},
)
# ---------------------------------------------------------------------------
# registry assembly
# ---------------------------------------------------------------------------
MARGINAL_VARS = ["step_length", "edep", "delta_e", "post_E", "cos_scatter"]
GROUPING_AXES = ["energy", "pdg", "material"]
def build_catalog() -> list[PlotSpec]:
"""All concrete plot specs, each with a unique id."""
specs: list[PlotSpec] = []
for var in MARGINAL_VARS:
specs.append(
PlotSpec(
f"marginal_{var}", "marginals", lambda b, v=var: _marginal_overall(b, v)
)
)
for axis in GROUPING_AXES:
specs.append(
PlotSpec(
f"marginal_{var}_by_{axis}",
"marginals",
lambda b, v=var, a=axis: _marginal_grouped(b, v, a),
)
)
specs += [
PlotSpec(
"event_total_edep",
"event",
lambda b: _event_scalar(
b,
"event_total_edep",
"Total deposited energy per event",
"total deposited energy [MeV]",
"total_edep",
use_all=True,
),
),
PlotSpec("event_total_edep_by_energy", "event", _event_total_edep_by_energy),
PlotSpec(
"event_mean_length",
"event",
lambda b: _event_scalar(
b,
"event_mean_length",
"Mean step length per event",
"mean step length [mm]",
"mean_length",
use_all=False,
),
),
PlotSpec(
"event_n_steps",
"event",
lambda b: _event_scalar(
b,
"event_n_steps",
"Number of steps per event",
"steps per event",
"n_steps",
use_all=False,
),
),
PlotSpec(
"shower_longitudinal",
"shower",
lambda b: _profile(
b,
"shower_longitudinal",
"Longitudinal shower profile",
"depth along shower axis [mm]",
depth_expr,
"depth_edges",
),
),
PlotSpec(
"shower_transverse",
"shower",
lambda b: _profile(
b,
"shower_transverse",
"Transverse shower profile",
"radius from shower axis [mm]",
transverse_expr,
"transverse_edges",
),
),
PlotSpec("species_edep_share", "species", _species_share),
PlotSpec("leakage_fraction", "species", _leakage),
PlotSpec("sec_count_per_event", "secondaries", _sec_count_per_event),
PlotSpec("sec_count_per_species", "secondaries", _sec_count_per_species),
PlotSpec("sec_energy", "secondaries", _sec_energy),
PlotSpec("sec_cos_angle", "secondaries", _sec_cos_angle),
]
return specs
def catalog_ids() -> list[str]:
return [s.id for s in build_catalog()]
def get_spec(spec_id: str) -> PlotSpec:
for s in build_catalog():
if s.id == spec_id:
return s
raise KeyError(f"unknown plot id: {spec_id!r}")
giant/analysis/condor.py
+148
View File
@@ -0,0 +1,148 @@
"""HTCondor orchestration: prep, per-plot compute, and the submit description.
Job model (one condor job per plot, compute/render split):
1. ``prep`` runs once on the submit node — resolves the shared context (fixed bin
edges, energy quantiles, top species/materials) from a subsample and writes
``shared.json``. Cheap; no LaTeX.
2. one job per catalog id runs ``giant analyze compute-one`` on a worker — a
single streaming pass producing ``reduced/<id>.json``. polars/numpy only, no
LaTeX, so it needs no plotstyle in the container.
3. a final *local* ``giant analyze render`` turns every reduced artifact into a
styled PDF + gallery metadata (that step imports plotstyle/LaTeX).
Files on ``/ceph`` or ``/work`` are reached via ``ProvidesETPResources``; no
HTCondor file transfer of the multi-GB inputs.
"""
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
from giant.analysis.catalog import Bundle, catalog_ids, get_spec
from giant.analysis.context import Context, build_context
# ---------------------------------------------------------------------------
# per-plot compute (what each condor job runs)
# ---------------------------------------------------------------------------
def compute_one(
spec_id: str,
rollout: str | Path,
reference: str | Path,
shared: str | Path,
out: str | Path,
) -> Path:
"""Run one plot's streaming reduction and write its ``Reduced`` JSON."""
ctx = Context.load(shared)
bundle = Bundle.open(rollout, reference, ctx)
reduced = get_spec(spec_id).compute(bundle)
out = Path(out)
reduced.save(out)
return out
# ---------------------------------------------------------------------------
# submit description
# ---------------------------------------------------------------------------
@dataclass
class SubmitConfig:
rollout: Path
reference: Path
out_dir: Path
accounting_group: str
repo_dir: Path
docker_image: str = "mschnepf/slc7-condocker"
request_memory_mb: int = 4096
request_cpus: int = 1
request_walltime_s: int = 3600
remote: bool = False # +RemoteJob (grid I/O) vs ProvidesETPResources (local files)
_WRAPPER = """#!/bin/bash
set -euo pipefail
cd {repo_dir}
exec uv run giant analyze compute-one \\
--id "$1" \\
--rollout {rollout} \\
--reference {reference} \\
--shared {shared} \\
--out {reduced_dir}/"$1".json
"""
def _submit_description(cfg: SubmitConfig, wrapper: Path, ids_file: Path) -> str:
reqs_attrs = (
"+RemoteJob = True\nrequest_walltime = {wt}\n".format(wt=cfg.request_walltime_s)
if cfg.remote
else "requirements = TARGET.ProvidesETPResources\n"
)
return (
"universe = docker\n"
f"docker_image = {cfg.docker_image}\n"
f"executable = {wrapper}\n"
"arguments = $(plotid)\n"
"should_transfer_files = YES\n"
"when_to_transfer_output = ON_EXIT\n"
f"request_memory = {cfg.request_memory_mb}\n"
f"request_cpus = {cfg.request_cpus}\n"
f"+RequestWalltime = {cfg.request_walltime_s}\n"
f"accounting_group = {cfg.accounting_group}\n"
f"{reqs_attrs}"
"output = logs/$(plotid).out\n"
"error = logs/$(plotid).err\n"
"log = logs/condor.log\n"
f"queue plotid from {ids_file}\n"
)
def write_submit(cfg: SubmitConfig, ids: list[str] | None = None) -> Path:
"""Write the wrapper script, plot-id list, and HTCondor submit description.
Returns the path to the submit description (``<out_dir>/analyze.sub``). Does
not submit — call ``condor_submit`` on the returned file, or use ``submit``.
"""
ids = ids or catalog_ids()
out_dir = cfg.out_dir
reduced_dir = out_dir / "reduced"
(out_dir / "logs").mkdir(parents=True, exist_ok=True)
reduced_dir.mkdir(parents=True, exist_ok=True)
wrapper = out_dir / "run_compute.sh"
wrapper.write_text(
_WRAPPER.format(
repo_dir=cfg.repo_dir,
rollout=cfg.rollout,
reference=cfg.reference,
shared=out_dir / "shared.json",
reduced_dir=reduced_dir,
)
)
wrapper.chmod(0o755)
ids_file = out_dir / "plotids.txt"
ids_file.write_text("\n".join(ids) + "\n")
sub = out_dir / "analyze.sub"
sub.write_text(_submit_description(cfg, wrapper, ids_file))
return sub
def prep(
rollout: str | Path,
reference: str | Path,
out_dir: str | Path,
**kwargs,
) -> Path:
"""Build and save the shared context (``<out_dir>/shared.json``)."""
out_dir = Path(out_dir)
out_dir.mkdir(parents=True, exist_ok=True)
ctx = build_context(rollout, reference, **kwargs)
shared = out_dir / "shared.json"
ctx.save(shared)
return shared
giant/analysis/context.py
+199
View File
@@ -0,0 +1,199 @@
"""Shared analysis context (the ``prep`` step): fixed bin edges + group sets.
Every histogram in the catalog bins against **fixed** edges so each compute job
is a single streaming pass with no min/max range scan. Those edges — plus the
energy-bin quantiles, the top PDG species and the material list to stratify by,
and the shower depth/transverse ranges — are resolved *once* here, on the submit
node, from a hash-subsample plus a few cheap exact ``group_by`` passes, and shipped
in ``shared.json``. Tiny and self-describing; no per-event arrays.
plotstyle-free (runs on the submit node, but also importable by workers).
"""
from __future__ import annotations
import json
from dataclasses import asdict, dataclass, field
from pathlib import Path
import numpy as np
import polars as pl
from giant.analysis.grouping import energy_bin_edges
from giant.analysis.reduce import (
attach_entry_axis,
depth_expr,
entry_axis,
transverse_expr,
)
from giant.analysis.sources import Side, open_side, physical_steps, secondaries
from giant.analysis.variables import RANGED_VARS
@dataclass
class Context:
"""Resolved bin edges and grouping sets shared by every compute job."""
n_marginal_bins: int
var_ranges: dict[str, tuple[float, float]] # ranged var -> (lo, hi)
energy_edges: list[float]
top_pdgs: list[int]
materials: list[str]
depth_edges: list[float]
transverse_edges: list[float]
sec_energy_range: tuple[float, float]
n_sec_bins: int
n_events: dict[str, int] = field(default_factory=dict)
# -- (de)serialization -------------------------------------------------
def save(self, path: str | Path) -> None:
Path(path).write_text(json.dumps(asdict(self), indent=2))
@classmethod
def load(cls, path: str | Path) -> "Context":
d = json.loads(Path(path).read_text())
d["var_ranges"] = {k: tuple(v) for k, v in d["var_ranges"].items()}
d["sec_energy_range"] = tuple(d["sec_energy_range"])
return cls(**d)
# -- convenience -------------------------------------------------------
def marginal_edges(self, var: str) -> np.ndarray:
lo, hi = self.var_ranges[var]
return np.linspace(lo, hi, self.n_marginal_bins + 1)
_LO_Q, _HI_Q = 0.001, 0.999
def _row_subsample(lf: pl.LazyFrame, sample_rows: int, seed: int) -> pl.LazyFrame:
"""Hash-subsample ~``sample_rows`` rows (for range estimation only)."""
n_total = lf.select(pl.len()).collect(engine="streaming").item()
if n_total <= sample_rows:
return lf
threshold = int(sample_rows / n_total * 2**32)
return lf.filter((pl.col("pre_E").hash(seed=seed) % 2**32) < threshold)
def _combined_quantiles(
r_vals: np.ndarray, t_vals: np.ndarray, lo_q: float, hi_q: float
) -> tuple[float, float]:
"""Robust (lo_q, hi_q) range over the union of two value samples."""
both = np.concatenate([r_vals, t_vals])
lo, hi = float(np.quantile(both, lo_q)), float(np.quantile(both, hi_q))
if not (hi - lo > 1e-6 * max(abs(hi), 1.0)):
lo, hi = lo - 0.5, hi + 0.5
return lo, hi
def build_context(
rollout: str | Path | pl.LazyFrame,
reference: str | Path | pl.LazyFrame,
*,
n_energy_bins: int = 4,
n_marginal_bins: int = 50,
n_sec_bins: int = 40,
top_k_pdg: int = 6,
sample_rows: int = 1_000_000,
seed: int = 0,
) -> Context:
"""Resolve the shared context from the two files (the ``prep`` step)."""
r_all = open_side(rollout, Side.rollout)
t_all = open_side(reference, Side.reference)
r_lf = physical_steps(r_all, Side.rollout)
t_lf = physical_steps(t_all, Side.reference)
# Ranged marginal variables: robust ranges over a shared row subsample.
exprs = [e.alias(n) for n, (_, e) in RANGED_VARS.items()]
r_s = (
_row_subsample(r_lf, sample_rows, seed)
.select(exprs)
.collect(engine="streaming")
)
t_s = (
_row_subsample(t_lf, sample_rows, seed)
.select(exprs)
.collect(engine="streaming")
)
var_ranges = {
name: _combined_quantiles(
r_s[name].to_numpy(), t_s[name].to_numpy(), _LO_Q, _HI_Q
)
for name in RANGED_VARS
}
# Energy-bin edges from exact per-event incident energies (cheap group_by).
def _incident(lf: pl.LazyFrame) -> np.ndarray:
return (
lf.group_by("event_id")
.agg(pl.col("pre_E").max())
.collect(engine="streaming")["pre_E"]
.to_numpy()
)
r_inc, t_inc = _incident(r_lf), _incident(t_lf)
energy_edges = energy_bin_edges(np.concatenate([r_inc, t_inc]), n_energy_bins)
# Top PDG species and material list (cheap single-column group_bys).
def _counts(lf: pl.LazyFrame, col: str) -> pl.DataFrame:
return lf.group_by(col).agg(pl.len().alias("n")).collect(engine="streaming")
pdg_counts = (
pl.concat([_counts(r_lf, "pdg"), _counts(t_lf, "pdg")])
.group_by("pdg")
.agg(pl.col("n").sum())
.sort("n", descending=True)
)
top_pdgs = [int(x) for x in pdg_counts["pdg"].to_list()[:top_k_pdg]]
materials = sorted(
set(_counts(r_lf, "material")["material"].to_list())
| set(_counts(t_lf, "material")["material"].to_list())
)
# Shower depth / transverse ranges from a subsampled proxy.
def _proxy(lf: pl.LazyFrame) -> tuple[np.ndarray, np.ndarray]:
ea = entry_axis(lf)
sub = (
attach_entry_axis(_row_subsample(lf, sample_rows, seed), ea)
.select(depth_expr().alias("d"), transverse_expr().alias("t"))
.collect(engine="streaming")
)
return sub["d"].to_numpy(), sub["t"].to_numpy()
r_d, r_t = _proxy(r_lf)
t_d, t_t = _proxy(t_lf)
d_lo, d_hi = _combined_quantiles(r_d, t_d, _LO_Q, _HI_Q)
depth_edges = np.linspace(d_lo, d_hi, n_marginal_bins + 1)
t_hi = max(float(np.quantile(np.concatenate([r_t, t_t]), _HI_Q)), 1e-6)
transverse_edges = np.linspace(0.0, t_hi, n_marginal_bins + 1)
# Secondary energy range.
r_se = secondaries(r_lf, Side.rollout).select("energy")
t_se = secondaries(t_all, Side.reference).select("energy")
r_se = _row_sample_col(r_se, sample_rows, seed)
t_se = _row_sample_col(t_se, sample_rows, seed)
sec_energy_range = _combined_quantiles(r_se, t_se, _LO_Q, _HI_Q)
return Context(
n_marginal_bins=n_marginal_bins,
var_ranges=var_ranges,
energy_edges=[float(x) for x in energy_edges],
top_pdgs=top_pdgs,
materials=materials,
depth_edges=[float(x) for x in depth_edges],
transverse_edges=[float(x) for x in transverse_edges],
sec_energy_range=sec_energy_range,
n_sec_bins=n_sec_bins,
n_events={
"rollout": len(r_inc),
"reference": len(t_inc),
},
)
def _row_sample_col(lf: pl.LazyFrame, sample_rows: int, seed: int) -> np.ndarray:
"""Collect a subsample of a single-column ``energy`` LazyFrame to numpy."""
vals = lf.collect(engine="streaming")["energy"].to_numpy()
if len(vals) > sample_rows:
rng = np.random.default_rng(seed)
vals = vals[rng.choice(len(vals), size=sample_rows, replace=False)]
return vals
giant/analysis/grouping.py
+108
View File
@@ -0,0 +1,108 @@
"""Grouping axes (overall / energy / pdg / material) and their labels.
Energy grouping is by the event's **incident (primary) energy** — the largest
``pre_E`` in the event — so every step of a shower lands in one bin, the physically
meaningful stratification for a calorimeter surrogate. The quantile bin *edges* are
sized once in ``prep`` (from a subsample) and shipped in ``shared.json``; a compute
job that needs them re-derives the small per-event ``event_id -> bin`` map itself
(one bounded streaming ``group_by`` over ``pre_E``), so ``shared.json`` stays tiny.
Pure/plotstyle-free so it can run on the compute workers.
"""
from __future__ import annotations
import numpy as np
import polars as pl
# Common electromagnetic/hadronic species; anything else falls back to its code.
PDG_NAMES: dict[int, str] = {
11: "e-",
-11: "e+",
22: "gamma",
2112: "n",
2212: "p",
-2212: "pbar",
111: "pi0",
211: "pi+",
-211: "pi-",
13: "mu-",
-13: "mu+",
321: "K+",
-321: "K-",
130: "K0L",
}
def pdg_label(code: int) -> str:
"""Human-readable species label for a PDG code (falls back to the code)."""
code = int(code)
if code in PDG_NAMES:
return PDG_NAMES[code]
if abs(code) > 1_000_000_000:
return f"ion {code}"
return str(code)
def material_label(name: str) -> str:
"""Display label for a Geant4 material, dropping the ``G4_`` prefix."""
return name[3:] if name.startswith("G4_") else name
def energy_bin_edges(incident_E: np.ndarray, n_bins: int = 4) -> np.ndarray:
"""Equal-population (quantile) bin edges over per-event incident energies.
Returns ``n_bins + 1`` monotonically non-decreasing edges. The top edge is
nudged up so the largest value falls inside the last bin under a
right-open convention. Degenerate (single-value) input widens by +/-0.5.
"""
incident_E = np.asarray(incident_E, dtype=np.float64)
edges = np.quantile(incident_E, np.linspace(0.0, 1.0, n_bins + 1))
edges = np.unique(edges)
if edges.size < 2:
v = edges[0] if edges.size else 0.0
edges = np.array([v - 0.5, v + 0.5])
edges[-1] = np.nextafter(edges[-1], np.inf)
return edges
def energy_bin_labels(edges: np.ndarray) -> list[str]:
"""``E in [lo, hi)`` labels for each bin defined by ``edges`` (MeV)."""
return [
f"E in [{edges[i]:.3g}, {edges[i + 1]:.3g}) MeV" for i in range(len(edges) - 1)
]
def digitize_expr(value: pl.Expr, edges: np.ndarray) -> pl.Expr:
"""Bin index of ``value`` under arbitrary (possibly non-uniform) ``edges``.
``bin = (#interior edges <= value)``, clipped to ``[0, n_bins-1]`` — matches
``np.digitize(value, edges[1:-1])`` and works for the quantile energy edges.
Vectorized as a sum of boolean comparisons; no per-row Python.
"""
interior = [float(e) for e in edges[1:-1]]
n_bins = len(edges) - 1
idx = pl.lit(0, dtype=pl.Int32)
for e in interior:
idx = idx + (value >= e).cast(pl.Int32)
return idx.clip(0, n_bins - 1)
def event_energy_bins(
lf: pl.LazyFrame, edges: np.ndarray
) -> tuple[np.ndarray, np.ndarray]:
"""Per-event incident-energy bin: ``(event_ids, bin_idx)`` numpy arrays.
Incident energy is ``max(pre_E)`` per event (the primary). One bounded
streaming ``group_by``; the tiny per-event result is digitized in numpy.
"""
per_event = (
lf.group_by("event_id")
.agg(pl.col("pre_E").max().alias("incident_E"))
.collect(engine="streaming")
.sort("event_id")
)
event_ids = per_event["event_id"].to_numpy()
incident = per_event["incident_E"].to_numpy()
bin_idx = np.clip(np.digitize(incident, edges[1:-1]), 0, len(edges) - 2)
return event_ids, bin_idx.astype(np.int64)
giant/analysis/reduce.py
+218
View File
@@ -0,0 +1,218 @@
"""Streaming compute primitives — the reduce half of the analysis.
Everything here turns a (possibly larger-than-RAM) LazyFrame into a *compact*
numpy/DataFrame artifact in bounded memory, and never imports plotstyle so it can
run on an HTCondor worker. Efficiency rules (see the plan's "Histogram
efficiency" section):
* ``hist1d`` is a single streaming ``group_by([group, bin]).len()`` pass against
**fixed** edges (no min/max range pass) with a strict column projection — only
the columns the value/group expressions reference are read from the parquet.
* the multi-quantity reductions (``event_scalars``, profiles, ``species_share``,
``leakage_fraction``) each emit *all* their outputs from one ``group_by``.
* per-event -> per-row lookups (shower entry/axis) use ``replace_strict`` (a hash
map applied as an expression, bounded memory), never a streaming join.
"""
from __future__ import annotations
import numpy as np
import polars as pl
from giant.constants import TERM_ESCAPED
# ---------------------------------------------------------------------------
# 1-D histogram primitive
# ---------------------------------------------------------------------------
def _bin_expr(value: pl.Expr, lo: float, hi: float, nbins: int) -> pl.Expr:
"""Uniform bin index of ``value`` over ``[lo, hi]`` into ``nbins`` bins."""
return ((value - lo) / (hi - lo) * nbins).floor().cast(pl.Int32).clip(0, nbins - 1)
def hist1d(
lf: pl.LazyFrame,
value: pl.Expr,
edges: np.ndarray,
group: pl.Expr | None = None,
) -> dict[object, np.ndarray]:
"""Streaming histogram of ``value`` over fixed uniform ``edges``, by ``group``.
Returns ``{group_key: counts}`` (counts is an ``int64`` array of length
``len(edges)-1``). One hash pass; runtime is independent of group cardinality,
so every pdg/material/energy stratum falls out together. Only the tiny
``(n_groups x nbins)`` result is materialized.
"""
lo, hi, nbins = float(edges[0]), float(edges[-1]), len(edges) - 1
group = pl.lit(0, dtype=pl.Int64) if group is None else group
res = (
lf.select(group.alias("_g"), _bin_expr(value, lo, hi, nbins).alias("_b"))
.group_by("_g", "_b")
.agg(pl.len().alias("_n"))
.collect(engine="streaming")
)
out: dict[object, np.ndarray] = {}
for g, b, n in res.iter_rows():
out.setdefault(g, np.zeros(nbins, dtype=np.int64))[b] = n
return out
# ---------------------------------------------------------------------------
# Per-event scalar observables (one bounded group_by pass)
# ---------------------------------------------------------------------------
def event_scalars(lf: pl.LazyFrame) -> pl.DataFrame:
"""One row per event: total/mean deposited energy, path length, step count.
Columns: ``event_id, total_edep, total_length, n_steps, mean_length,
incident_E`` (incident = ``max(pre_E)``, the primary). The caller chooses
whether ``lf`` includes the rollout's synthetic termination rows — pass the
full scan for energy totals (they carry the deposited remainder), physical
steps only for step-count / mean-length.
"""
return (
lf.group_by("event_id")
.agg(
pl.col("edep").sum().alias("total_edep"),
pl.col("step_length").sum().alias("total_length"),
pl.len().alias("n_steps"),
pl.col("pre_E").max().alias("incident_E"),
)
.with_columns((pl.col("total_length") / pl.col("n_steps")).alias("mean_length"))
.collect(engine="streaming")
)
# ---------------------------------------------------------------------------
# Shower shape: entry/axis + edep-weighted longitudinal / transverse profiles
# ---------------------------------------------------------------------------
def entry_axis(lf: pl.LazyFrame) -> pl.DataFrame:
"""Per-event shower entry point and axis (from the highest-``pre_E`` step).
One bounded ``group_by``: the primary's ``pre_pos`` becomes the entry point
and its ``pre_dir`` the shower axis.
"""
return (
lf.group_by("event_id")
.agg(
pl.col("pre_x").get(pl.col("pre_E").arg_max()).alias("entry_x"),
pl.col("pre_y").get(pl.col("pre_E").arg_max()).alias("entry_y"),
pl.col("pre_z").get(pl.col("pre_E").arg_max()).alias("entry_z"),
pl.col("pre_dx").get(pl.col("pre_E").arg_max()).alias("axis_x"),
pl.col("pre_dy").get(pl.col("pre_E").arg_max()).alias("axis_y"),
pl.col("pre_dz").get(pl.col("pre_E").arg_max()).alias("axis_z"),
)
.collect(engine="streaming")
.sort("event_id")
)
_ENTRY_AXIS_COLS = ("entry_x", "entry_y", "entry_z", "axis_x", "axis_y", "axis_z")
def attach_entry_axis(lf: pl.LazyFrame, entry: pl.DataFrame) -> pl.LazyFrame:
"""Broadcast each event's entry/axis onto its rows via ``replace_strict``.
A hash map applied as an expression — streams in bounded memory, unlike a
join which would buffer the whole file-sized left side.
"""
ids = entry["event_id"].to_numpy()
return lf.with_columns(
pl.col("event_id")
.replace_strict(ids, entry[col].to_numpy(), return_dtype=pl.Float64)
.alias(col)
for col in _ENTRY_AXIS_COLS
)
def depth_expr() -> pl.Expr:
"""Signed distance of ``post_pos`` from the entry point along the shower axis."""
dx = pl.col("post_x") - pl.col("entry_x")
dy = pl.col("post_y") - pl.col("entry_y")
dz = pl.col("post_z") - pl.col("entry_z")
return dx * pl.col("axis_x") + dy * pl.col("axis_y") + dz * pl.col("axis_z")
def transverse_expr() -> pl.Expr:
"""Perpendicular distance of ``post_pos`` from the shower axis."""
dx = pl.col("post_x") - pl.col("entry_x")
dy = pl.col("post_y") - pl.col("entry_y")
dz = pl.col("post_z") - pl.col("entry_z")
depth = dx * pl.col("axis_x") + dy * pl.col("axis_y") + dz * pl.col("axis_z")
tx = dx - depth * pl.col("axis_x")
ty = dy - depth * pl.col("axis_y")
tz = dz - depth * pl.col("axis_z")
return (tx**2 + ty**2 + tz**2).sqrt()
def weighted_profile(
lf: pl.LazyFrame,
coord: pl.Expr,
edges: np.ndarray,
weight: pl.Expr,
) -> tuple[np.ndarray, np.ndarray]:
"""Event-averaged, ``weight``-summed profile of ``coord``, with an event-RMS band.
One streaming ``group_by(event_id, bin)`` sums ``weight`` per (event, bin);
collapsed in numpy to the per-bin mean over events and its event-to-event std
(the band). ``coord``/``weight`` require the entry/axis columns attached.
Returns ``(mean, std)``, each length ``len(edges)-1``.
"""
lo, hi, nbins = float(edges[0]), float(edges[-1]), len(edges) - 1
grid = (
lf.select(
"event_id",
_bin_expr(coord, lo, hi, nbins).alias("_b"),
weight.alias("_w"),
)
.group_by("event_id", "_b")
.agg(pl.col("_w").sum().alias("_ws"))
.collect(engine="streaming")
)
ev = grid["event_id"].to_numpy()
uniq, inv = np.unique(ev, return_inverse=True)
mat = np.zeros((len(uniq), nbins), dtype=np.float64)
np.add.at(mat, (inv, grid["_b"].to_numpy()), grid["_ws"].to_numpy())
return mat.mean(axis=0), mat.std(axis=0)
# ---------------------------------------------------------------------------
# Species contribution and leakage
# ---------------------------------------------------------------------------
def species_share(lf: pl.LazyFrame) -> pl.DataFrame:
"""Total deposited energy per PDG species (``pdg, total_edep``), one pass."""
return (
lf.group_by("pdg")
.agg(pl.col("edep").sum().alias("total_edep"))
.collect(engine="streaming")
.sort("total_edep", descending=True)
)
def leakage_fraction(lf: pl.LazyFrame) -> np.ndarray:
"""Per-event escaped-energy fraction (rollout only), one bounded pass.
Escaped rows carry the leaked energy in ``pre_E`` (``edep`` is 0 there); the
fraction is ``escaped / (deposited + escaped)`` per event.
"""
per_event = (
lf.group_by("event_id")
.agg(
pl.col("edep").sum().alias("deposited"),
pl.col("pre_E")
.filter(pl.col("termination_reason") == TERM_ESCAPED)
.sum()
.alias("escaped"),
)
.collect(engine="streaming")
)
deposited = per_event["deposited"].to_numpy()
escaped = per_event["escaped"].fill_null(0.0).to_numpy()
total = deposited + escaped
return np.where(total > 0, escaped / total, 0.0)
giant/analysis/reduced.py
+38
View File
@@ -0,0 +1,38 @@
"""The compact, self-describing artifact a compute job produces per plot.
Serialized as small JSON (no pickle, no per-event arrays) so it is trivially
transferable off the batch worker and human-inspectable. ``render.py`` dispatches
on ``kind`` and needs nothing but this file.
"""
from __future__ import annotations
import json
from dataclasses import asdict, dataclass, field
from pathlib import Path
# Reduced.kind values:
# "overlay_hist" rollout vs reference density histogram over shared edges
# "grouped_hist" one panel per group (energy/pdg/material), each an overlay
# "profile" edep-weighted mean +/- event-RMS vs depth/radius, two series
# "bar" per-category rollout vs reference bars (share / counts)
# "single_hist" one series only (e.g. rollout leakage; reference has none)
@dataclass
class Reduced:
id: str
family: str
kind: str
title: str
xlabel: str
payload: dict
meta: dict = field(default_factory=dict)
def save(self, path: str | Path) -> None:
Path(path).parent.mkdir(parents=True, exist_ok=True)
Path(path).write_text(json.dumps(asdict(self)))
@classmethod
def load(cls, path: str | Path) -> "Reduced":
return cls(**json.loads(Path(path).read_text()))
giant/analysis/render.py
+206
View File
@@ -0,0 +1,206 @@
"""Render reduced artifacts to styled PDFs + gallery metadata (the local step).
This is the *only* module that imports ``plotstyle`` (ETPlot's KIT matplotlib
theme), which renders through a real LaTeX toolchain — so it runs on the
submit/login node, never on a compute worker. It reads nothing but the small
``Reduced`` JSON files a run produced, so it is fully decoupled from the heavy
streaming compute.
For each reduced artifact it writes ``<out>/<family>/<id>.pdf`` plus a sibling
``<id>.yaml`` (per-plot gallery metadata) and a per-family ``metadata.yaml``.
Optionally runs ``gallery generate`` to build the static HTML site.
"""
from __future__ import annotations
import subprocess
from pathlib import Path
import numpy as np
import plotstyle as ps # ty: ignore[unresolved-import]
import yaml
from giant.analysis.reduced import Reduced
_SERIES_LABELS = {"rollout": "rollout", "reference": "reference (Geant4)"}
def _density(counts: list[int] | np.ndarray, edges: np.ndarray) -> np.ndarray:
counts = np.asarray(counts, dtype=np.float64)
total = counts.sum()
if total == 0:
return counts
return counts / (total * (edges[1] - edges[0]))
def _overlay(ax, edges: np.ndarray, series: dict[str, list], log_y: bool) -> None:
for key in ("reference", "rollout"):
if key in series:
ax.stairs(_density(series[key], edges), edges, label=_SERIES_LABELS[key])
if log_y:
ax.set_yscale("log")
def _render_overlay(r: Reduced):
edges = np.asarray(r.payload["edges"])
fig, ax = ps.new_figure("thesis-single", title=r.title)
_overlay(ax, edges, r.payload, r.payload.get("log_y", False))
ax.set_xlabel(r.xlabel)
ax.set_ylabel("density")
ps.style_legend(ax, title="source")
return fig
def _render_single(r: Reduced):
edges = np.asarray(r.payload["edges"])
fig, ax = ps.new_figure("thesis-single", title=r.title)
ax.stairs(
_density(r.payload["rollout"], edges), edges, label=_SERIES_LABELS["rollout"]
)
if r.payload.get("log_y"):
ax.set_yscale("log")
ax.set_xlabel(r.xlabel)
ax.set_ylabel("density")
ps.style_legend(ax, title="source")
return fig
def _render_grouped(r: Reduced):
edges = np.asarray(r.payload["edges"])
groups = r.payload["groups"]
labels = list(groups)
n = len(labels)
ncols = min(3, n) or 1
nrows = (n + ncols - 1) // ncols
fig, axes = ps.new_figure(
"slide-16x9", title=r.title, nrows=nrows, ncols=ncols, squeeze=False
)
flat = axes.ravel()
for i, lbl in enumerate(labels):
ax = flat[i]
_overlay(ax, edges, groups[lbl], r.payload.get("log_y", False))
ax.set_title(lbl, fontsize=8)
ax.set_xlabel(r.xlabel)
for j in range(n, len(flat)):
flat[j].set_visible(False)
ps.style_legend(flat[0], title="source")
return fig
def _render_profile(r: Reduced):
edges = np.asarray(r.payload["edges"])
centers = 0.5 * (edges[:-1] + edges[1:])
fig, ax = ps.new_figure("thesis-single", title=r.title)
for key in ("reference", "rollout"):
mean = np.asarray(r.payload[f"{key}_mean"])
std = np.asarray(r.payload[f"{key}_std"])
(line,) = ax.plot(centers, mean, label=_SERIES_LABELS[key])
ax.fill_between(
centers, mean - std, mean + std, alpha=0.2, color=line.get_color()
)
ax.set_xlabel(r.xlabel)
ax.set_ylabel(r.payload.get("ylabel", "mean deposited energy [MeV]"))
ps.style_legend(ax, title="source")
return fig
def _render_bar(r: Reduced):
labels = r.payload["labels"]
x = np.arange(len(labels))
width = 0.4
fig, ax = ps.new_figure("thesis-single", title=r.title)
ax.bar(
x - width / 2, r.payload["reference"], width, label=_SERIES_LABELS["reference"]
)
ax.bar(x + width / 2, r.payload["rollout"], width, label=_SERIES_LABELS["rollout"])
ax.set_xticks(x)
ax.set_xticklabels(labels, rotation=45, ha="right")
ax.set_ylabel(r.payload.get("ylabel", "value"))
ps.style_legend(ax, title="source")
return fig
_RENDERERS = {
"overlay_hist": _render_overlay,
"single_hist": _render_single,
"grouped_hist": _render_grouped,
"profile": _render_profile,
"bar": _render_bar,
}
def render(r: Reduced):
"""Build the matplotlib figure for one reduced artifact (dispatch on kind)."""
return _RENDERERS[r.kind](r)
def _plot_metadata(r: Reduced, run_meta: dict) -> dict:
meta = {
"title": r.title,
"description": f"Rollout vs reference: {r.title}.",
"plot_type": r.kind,
"family": r.family,
}
meta.update(r.meta)
if "note" in r.payload:
meta["note"] = r.payload["note"]
return meta
def render_all(
reduced_dir: str | Path,
out_dir: str | Path,
run_meta: dict | None = None,
*,
run_gallery: bool = False,
) -> list[Path]:
"""Render every reduced artifact under ``reduced_dir`` to a PDF tree.
Writes ``<out>/<family>/<id>.pdf`` + ``<id>.yaml`` and a per-family
``metadata.yaml`` (carrying the run's checkpoint/paths as gallery params).
Returns the list of PDF paths written.
"""
ps.use()
run_meta = run_meta or {}
reduced_dir, out_dir = Path(reduced_dir), Path(out_dir)
pdfs: list[Path] = []
families: set[str] = set()
for jf in sorted(reduced_dir.glob("*.json")):
r = Reduced.load(jf)
family_dir = out_dir / r.family
family_dir.mkdir(parents=True, exist_ok=True)
families.add(r.family)
fig = render(r)
ps.savefig(fig, str(family_dir / r.id), formats=("pdf",))
(family_dir / f"{r.id}.yaml").write_text(
yaml.safe_dump(_plot_metadata(r, run_meta), sort_keys=False)
)
pdfs.append(family_dir / f"{r.id}.pdf")
import matplotlib.pyplot as plt
plt.close(fig)
# Root + per-family gallery metadata.
out_dir.mkdir(parents=True, exist_ok=True)
(out_dir / "metadata.yaml").write_text(
yaml.safe_dump(
{
"title": run_meta.get("title", "GIANT rollout analysis"),
"description": "Autoregressive rollout compared against held-out Geant4 reference steps.",
"experiment": "GIANT",
"parameters": run_meta,
},
sort_keys=False,
)
)
for fam in families:
(out_dir / fam / "metadata.yaml").write_text(
yaml.safe_dump(
{"title": fam, "description": f"{fam} plots."}, sort_keys=False
)
)
if run_gallery:
subprocess.run(["gallery", "generate", "--source", str(out_dir)], check=True)
return pdfs
giant/analysis/sources.py
+165
View File
@@ -0,0 +1,165 @@
"""Canonical world-frame LazyFrame builders for the two sides of a comparison.
The analysis compares one autoregressive ``giant rollout`` (the *generated* side)
against a raw miniCaloSim steps file (the *reference* / real side). Both carry a
**shared world-frame physical column subset** under identical names, so no
renaming or coordinate decode is needed — everything is already in world-frame
mm / MeV:
event_id, track_id, step_no, pdg,
pre_x, pre_y, pre_z, pre_E, pre_dx, pre_dy, pre_dz,
post_x, post_y, post_z, post_E, post_dx, post_dy, post_dz,
edep, step_length, material, layer_id
(rollout: ``rollout.py:_RECORD_KEYS``; reference: minicalosim ``RunAction.cc``
Steps ntuple passed through by ``dwarf convert``.)
The two files differ in their *extra* columns — the rollout adds ``parent_id``,
``generation``, ``n_sec_pred``, ``termination_reason``; the reference adds
``process``, field columns, ``child_track_ids`` and the ``sec_*_list`` secondary
birth-state lists. Those are only touched by the side-specific helpers here
(synthetic-row filtering, the secondary view).
Nothing in this module (or ``reduce.py``) imports plotstyle — compute runs on
HTCondor workers that have no LaTeX toolchain.
"""
from __future__ import annotations
from enum import Enum
from pathlib import Path
import polars as pl
import pyarrow.parquet as pq
from giant.constants import (
PREDICT_COORD_METADATA_KEY,
ROLLOUT_COORD_VALUE,
TERM_ENERGY_CUTOFF,
TERM_ESCAPED,
TERM_MAX_STEPS,
TERM_UNKNOWN_PDG,
)
# The world-frame physical columns both sides share under identical names.
PHYS_COLS: tuple[str, ...] = (
"event_id",
"pdg",
"pre_x",
"pre_y",
"pre_z",
"pre_E",
"pre_dx",
"pre_dy",
"pre_dz",
"post_x",
"post_y",
"post_z",
"post_E",
"post_dx",
"post_dy",
"post_dz",
"edep",
"step_length",
"material",
"layer_id",
)
# Rollout rows written purely for bookkeeping (a track's forced stop): they carry
# step_length=0, post_pos=pre_pos, and — for every reason but escaped — the
# track's whole remaining pre_E dumped into edep so the shower still conserves
# energy. They are not physical steps (the reference has no equivalent), so a
# per-step marginal comparison must drop them; a per-event energy total must keep
# them. See rollout.py's terminal-row handling.
SYNTHETIC_TERMINATION_REASONS: frozenset[str] = frozenset(
{TERM_ESCAPED, TERM_UNKNOWN_PDG, TERM_ENERGY_CUTOFF, TERM_MAX_STEPS}
)
class Side(str, Enum):
"""Which of the two comparison inputs a file is."""
rollout = "rollout"
reference = "reference"
def _check_rollout_metadata(path: Path) -> None:
"""Raise if ``path`` carries coord metadata that isn't the rollout tag.
A missing tag (older rollout output, predating tagging) is allowed through,
matching ``giant rollout``'s own leniency; a tag that is present but wrong is
a real mismatch and worth failing on before the column layout is trusted.
"""
metadata = pq.read_schema(path).metadata or {}
coord = metadata.get(PREDICT_COORD_METADATA_KEY.encode())
if coord is not None and coord.decode() != ROLLOUT_COORD_VALUE:
raise ValueError(
f"{path} is not a rollout file (coord={coord.decode()!r}); "
"expected `giant rollout` output"
)
def open_side(source: str | Path | pl.LazyFrame, side: Side) -> pl.LazyFrame:
"""Lazily scan one side's file, verifying the rollout tag when applicable.
Returns the *full* lazy scan (no column projection) so downstream reductions
can push their own narrow projection into the parquet read — the single
biggest lever on a larger-than-RAM file. ``pl.LazyFrame`` inputs pass straight
through (used by tests).
"""
if isinstance(source, pl.LazyFrame):
return source
path = Path(source)
if side is Side.rollout:
_check_rollout_metadata(path)
return pl.scan_parquet(path)
def physical_steps(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame:
"""Real, physical steps only — drops the rollout's synthetic termination rows.
The predicate is pushed down so the dropped rows are never decoded. The
reference has no such rows, so it is returned unchanged.
"""
if side is Side.reference:
return lf
return lf.filter(
~pl.col("termination_reason").is_in(list(SYNTHETIC_TERMINATION_REASONS))
)
def secondaries(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame:
"""Per-secondary birth state, one row per produced secondary.
Canonical columns: ``event_id, energy, pdg, sdx, sdy, sdz`` (birth energy in
MeV, PDG code, birth unit direction in the world frame). The two sides encode
secondaries differently:
- rollout: each secondary is its own track, so its birth state is the row with
``generation > 0`` and ``step_no == 0`` (``pre_E`` / ``pre_dir`` there).
- reference: secondaries live in per-parent-step ``sec_*_list`` columns; the
lists are exploded together and empty (no-secondary) steps drop out.
"""
if side is Side.rollout:
return lf.filter((pl.col("generation") > 0) & (pl.col("step_no") == 0)).select(
"event_id",
pl.col("pre_E").alias("energy"),
"pdg",
pl.col("pre_dx").alias("sdx"),
pl.col("pre_dy").alias("sdy"),
pl.col("pre_dz").alias("sdz"),
)
lists = ["sec_E_list", "sec_pdg_list", "sec_dx_list", "sec_dy_list", "sec_dz_list"]
return (
lf.select("event_id", *lists)
.explode(lists)
.drop_nulls("sec_E_list")
.select(
"event_id",
pl.col("sec_E_list").alias("energy"),
pl.col("sec_pdg_list").alias("pdg"),
pl.col("sec_dx_list").alias("sdx"),
pl.col("sec_dy_list").alias("sdy"),
pl.col("sec_dz_list").alias("sdz"),
)
)
giant/analysis/variables.py
+28
View File
@@ -0,0 +1,28 @@
"""Per-step value expressions shared by ``context`` (range sizing) and ``catalog``.
Kept separate from both so the range-sizing prep and the plot registry agree on
exactly what each variable *is*, with no import cycle. plotstyle-free.
"""
from __future__ import annotations
import polars as pl
# Ranged marginal variables: name -> (axis label, value expression). Their
# histogram ranges are sized from data in `context.build_context`.
RANGED_VARS: dict[str, tuple[str, pl.Expr]] = {
"step_length": ("step length [mm]", pl.col("step_length")),
"edep": ("deposited energy [MeV]", pl.col("edep")),
"delta_e": ("energy loss [MeV]", pl.col("pre_E") - pl.col("post_E")),
"post_E": ("post-step energy [MeV]", pl.col("post_E")),
}
def cos_scatter_expr() -> pl.Expr:
"""cos of the scattering angle: ``pre_dir . post_dir`` (both unit), in [-1, 1]."""
dot = (
pl.col("pre_dx") * pl.col("post_dx")
+ pl.col("pre_dy") * pl.col("post_dy")
+ pl.col("pre_dz") * pl.col("post_dz")
)
return dot.clip(-1.0, 1.0)
giant/cli.py
+128 -1
View File
@@ -891,7 +891,7 @@ def _seed_from_data(files: list[Path], n_events: int | None) -> dict[str, np.nda
Streams conditioning columns and keeps the highest-pre_E step per event_id —
the codebase's convention for the primary (a secondary always carries less
energy than its parent). See giant/analysis.py:_entry_axis_and_bin_edges.
energy than its parent). See giant/analysis/reduce.py:entry_axis.
"""
best_E: dict[int, float] = {}
best: dict[int, tuple] = {}
@@ -1114,5 +1114,132 @@ def rollout(
typer.echo(f"reference: {ref_path}")
analyze_app = typer.Typer(
no_args_is_help=True,
help="Rollout-vs-reference analysis: parallel compute on HTCondor + local render.",
)
app.add_typer(analyze_app, name="analyze")
@analyze_app.command("prep")
def analyze_prep(
rollout: Annotated[Path, typer.Option("--rollout", help="giant rollout parquet")],
reference: Annotated[
Path, typer.Option("--reference", help="Reference miniCaloSim steps parquet")
],
out_dir: Annotated[
Path,
typer.Option(
"--out-dir", "-o", help="Run directory for shared.json / reduced / plots"
),
],
n_energy_bins: Annotated[int, typer.Option("--energy-bins")] = 4,
n_marginal_bins: Annotated[int, typer.Option("--bins")] = 50,
top_k_pdg: Annotated[int, typer.Option("--top-pdg")] = 6,
) -> None:
"""Resolve the shared context (fixed bin edges / group sets) → shared.json."""
from giant.analysis import prep
shared = prep(
rollout,
reference,
out_dir,
n_energy_bins=n_energy_bins,
n_marginal_bins=n_marginal_bins,
top_k_pdg=top_k_pdg,
)
typer.echo(f"wrote {shared}")
@analyze_app.command("compute-one")
def analyze_compute_one(
id: Annotated[
str, typer.Option("--id", help="Catalog plot id (see `analyze list`)")
],
rollout: Annotated[Path, typer.Option("--rollout")],
reference: Annotated[Path, typer.Option("--reference")],
shared: Annotated[
Path, typer.Option("--shared", help="shared.json from `analyze prep`")
],
out: Annotated[Path, typer.Option("--out", help="Output reduced JSON path")],
) -> None:
"""Run one plot's streaming reduction (this is what each condor job runs)."""
from giant.analysis import compute_one
path = compute_one(id, rollout, reference, shared, out)
typer.echo(f"wrote {path}")
@analyze_app.command("list")
def analyze_list() -> None:
"""Print every catalog plot id."""
from giant.analysis import catalog_ids
for pid in catalog_ids():
typer.echo(pid)
@analyze_app.command("render")
def analyze_render(
reduced_dir: Annotated[
Path, typer.Option("--reduced-dir", help="Directory of reduced *.json")
],
out: Annotated[Path, typer.Option("--out", "-o", help="Output PDF/gallery tree")],
gallery: Annotated[
bool,
typer.Option(
"--gallery/--no-gallery", help="Run `gallery generate` after rendering"
),
] = False,
) -> None:
"""Render reduced artifacts to styled PDFs + gallery metadata (local; needs LaTeX)."""
from giant.analysis.render import render_all
pdfs = render_all(reduced_dir, out, run_gallery=gallery)
typer.echo(f"rendered {len(pdfs)} plots → {out}")
@analyze_app.command("submit")
def analyze_submit(
rollout: Annotated[Path, typer.Option("--rollout")],
reference: Annotated[Path, typer.Option("--reference")],
out_dir: Annotated[Path, typer.Option("--out-dir", "-o")],
accounting_group: Annotated[str, typer.Option("--accounting-group")],
docker_image: Annotated[
str, typer.Option("--docker-image")
] = "mschnepf/slc7-condocker",
request_memory: Annotated[int, typer.Option("--request-memory", help="MB")] = 4096,
remote: Annotated[
bool,
typer.Option("--remote/--local", help="+RemoteJob vs ProvidesETPResources"),
] = False,
dry_run: Annotated[
bool, typer.Option("--dry-run", help="Write files but don't condor_submit")
] = False,
) -> None:
"""prep + write the HTCondor submit description (one job per plot), then submit."""
import subprocess
from giant.analysis import SubmitConfig, prep, write_submit
prep(rollout, reference, out_dir)
cfg = SubmitConfig(
rollout=rollout.resolve(),
reference=reference.resolve(),
out_dir=out_dir,
accounting_group=accounting_group,
repo_dir=Path.cwd(),
docker_image=docker_image,
request_memory_mb=request_memory,
remote=remote,
)
sub = write_submit(cfg)
typer.echo(f"wrote submit description: {sub}")
if dry_run:
typer.echo("dry-run: not submitting")
return
subprocess.run(["condor_submit", str(sub)], check=True)
if __name__ == "__main__":
app()
pyproject.toml
+7
View File
@@ -39,6 +39,9 @@ analysis = [
"matplotlib>=3.8,<4",
"polars>=1.0,<2",
"ipykernel>=7.3.0",
# ETPlot's plotstyle (KIT matplotlib theme) + gallery CLI. Only the local
# `giant analyze render` step imports it; compute workers never do.
"gallery[plotting]",
]
[project.scripts]
@@ -65,6 +68,10 @@ torch = [
{ index = "pytorch-cpu", extra = "cpu" },
{ index = "pytorch-cu118", extra = "cuda" },
]
# ETPlot ships the `gallery` distribution (which provides the `plotstyle`
# package under its `plotting` extra). Local checkout next to this repo; swap
# for `{ git = "https://git.larsbogner.de/lars/ETPlot" }` off-machine.
gallery = { path = "../ETPlot", editable = true }
[[tool.uv.index]]
name = "pytorch-cpu"
tests/test_analysis.py
-1175
View File
File diff suppressed because it is too large Load Diff
tests/test_analysis_reduce.py
+171
View File
@@ -0,0 +1,171 @@
"""Tests for the streaming compute primitives (giant.analysis.reduce/sources/grouping)."""
from __future__ import annotations
import numpy as np
import polars as pl
from giant.analysis import grouping as G
from giant.analysis import reduce as R
from giant.analysis.sources import (
SYNTHETIC_TERMINATION_REASONS,
Side,
physical_steps,
secondaries,
)
def _rollout_frame() -> pl.LazyFrame:
# event 1: primary (2 steps) + 1 secondary track + 1 escaped bookkeeping row
# event 2: primary (1 step)
return pl.DataFrame(
{
"event_id": [1, 1, 1, 1, 2],
"track_id": [0, 0, 1, 0, 0],
"parent_id": [-1, -1, 0, -1, -1],
"generation": [0, 0, 1, 0, 0],
"step_no": [0, 1, 0, 99, 0],
"pdg": [11, 11, 22, 11, 11],
"pre_x": [0.0, 0.0, 0.0, 0.0, 0.0],
"pre_y": [0.0, 0.0, 0.0, 0.0, 0.0],
"pre_z": [0.0, 1.0, 1.0, 2.0, 0.0],
"pre_E": [100.0, 60.0, 20.0, 30.0, 50.0],
"pre_dx": [0.0, 0.0, 1.0, 0.0, 0.0],
"pre_dy": [0.0, 0.0, 0.0, 0.0, 0.0],
"pre_dz": [1.0, 1.0, 0.0, 1.0, 1.0],
"post_x": [0.0, 0.0, 1.0, 0.0, 0.0],
"post_y": [0.0, 0.0, 0.0, 0.0, 0.0],
"post_z": [1.0, 2.0, 1.0, 2.0, 1.0],
"post_E": [60.0, 30.0, 0.0, 0.0, 20.0],
"post_dx": [0.0, 0.0, 1.0, 0.0, 0.0],
"post_dy": [0.0, 0.0, 0.0, 0.0, 0.0],
"post_dz": [1.0, 1.0, 0.0, 1.0, 1.0],
"edep": [40.0, 30.0, 20.0, 0.0, 30.0],
"step_length": [1.0, 1.0, 1.0, 0.0, 1.0],
"material": ["G4_PbWO4"] * 5,
"layer_id": [0, 1, 1, -1, 0],
"n_sec_pred": [1, 0, 0, 0, 0],
"termination_reason": [
"",
"natural_end",
"natural_end",
"escaped",
"natural_end",
],
}
).lazy()
def _reference_frame() -> pl.LazyFrame:
return pl.DataFrame(
{
"event_id": [1, 1, 2],
"track_id": [0, 0, 0],
"step_no": [0, 1, 0],
"pdg": [11, 11, 11],
"pre_x": [0.0, 0.0, 0.0],
"pre_y": [0.0, 0.0, 0.0],
"pre_z": [0.0, 1.0, 0.0],
"pre_E": [100.0, 60.0, 50.0],
"pre_dx": [0.0, 0.0, 0.0],
"pre_dy": [0.0, 0.0, 0.0],
"pre_dz": [1.0, 1.0, 1.0],
"post_x": [0.0, 0.0, 0.0],
"post_y": [0.0, 0.0, 0.0],
"post_z": [1.0, 2.0, 1.0],
"post_E": [60.0, 30.0, 20.0],
"post_dx": [0.0, 0.0, 0.0],
"post_dy": [0.0, 0.0, 0.0],
"post_dz": [1.0, 1.0, 1.0],
"edep": [40.0, 30.0, 30.0],
"step_length": [1.0, 1.0, 1.0],
"material": ["G4_PbWO4", "G4_PbWO4", "G4_Pb"],
"layer_id": [0, 1, 0],
"sec_E_list": [[20.0], [], [10.0]],
"sec_pdg_list": [[22], [], [22]],
"sec_dx_list": [[1.0], [], [0.0]],
"sec_dy_list": [[0.0], [], [0.0]],
"sec_dz_list": [[0.0], [], [1.0]],
}
).lazy()
def test_hist1d_overall_and_grouped():
lf = _rollout_frame()
edges = np.linspace(0.0, 50.0, 6) # width 10
h = R.hist1d(lf, pl.col("edep"), edges)
# edep values: 40,30,20,0,30 -> bins [0),[10),[20),[30),[40)
assert h[0].tolist() == [1, 0, 1, 2, 1]
# grouped by pdg: pdg 22 has a single edep=20
hg = R.hist1d(lf, pl.col("edep"), edges, group=pl.col("pdg"))
assert hg[22].tolist() == [0, 0, 1, 0, 0]
assert hg[11].sum() == 4
def test_physical_steps_drops_synthetic_rollout_rows_only():
lf = _rollout_frame()
phys = physical_steps(lf, Side.rollout).collect()
assert phys.height == 4 # dropped the escaped bookkeeping row
assert "escaped" not in phys["termination_reason"].to_list()
assert SYNTHETIC_TERMINATION_REASONS # non-empty guard
# reference passes through unchanged
ref = _reference_frame()
assert physical_steps(ref, Side.reference).collect().height == ref.collect().height
def test_event_scalars_totals_include_all_rows():
lf = _rollout_frame()
es = R.event_scalars(lf).sort("event_id")
row1 = es.filter(pl.col("event_id") == 1).to_dicts()[0]
assert row1["total_edep"] == 90.0 # 40+30+20+0
assert row1["incident_E"] == 100.0
assert row1["n_steps"] == 4
def test_secondaries_rollout_vs_reference_align():
r = secondaries(_rollout_frame(), Side.rollout).collect().sort("event_id")
assert r["energy"].to_list() == [20.0] # only the generation>0, step_no==0 row
assert r["pdg"].to_list() == [22]
t = secondaries(_reference_frame(), Side.reference).collect().sort("event_id")
# two secondaries (event 1 and event 2); empty list dropped
assert sorted(t["energy"].to_list()) == [10.0, 20.0]
assert t["pdg"].to_list() == [22, 22]
def test_leakage_fraction():
frac = R.leakage_fraction(_rollout_frame())
# event 1: escaped pre_E=30, deposited=90 -> 30/120 = 0.25; event 2: 0
assert sorted(round(f, 6) for f in frac) == [0.0, 0.25]
def test_weighted_profile_matches_manual_bincount():
lf = _rollout_frame()
ea = R.entry_axis(lf)
lf2 = R.attach_entry_axis(lf, ea)
edges = np.linspace(0.0, 3.0, 4) # depth bins along +z
mean, std = R.weighted_profile(lf2, R.depth_expr(), edges, pl.col("edep"))
assert mean.shape == (3,)
# totals conserved: sum over bins == mean total edep per event
assert np.isclose(mean.sum() * 1, (90.0 + 30.0) / 2) # 2 events
def test_energy_bins_edges_and_event_map():
incident = np.array([100.0, 100.0, 1000.0, 1000.0])
edges = G.energy_bin_edges(incident, n_bins=2)
assert len(edges) == 3 and edges[0] <= 100.0 < edges[-1]
ids, bins = G.event_energy_bins(_rollout_frame(), edges)
assert set(bins.tolist()) <= {0, 1}
assert len(ids) == 2
def test_digitize_expr_matches_numpy():
edges = np.array([0.0, 10.0, 100.0, 1000.0])
df = pl.DataFrame({"v": [5.0, 50.0, 500.0, 2000.0]})
got = df.select(G.digitize_expr(pl.col("v"), edges).alias("b"))["b"].to_list()
assert got == np.digitize([5.0, 50.0, 500.0, 2000.0], edges[1:-1]).tolist()
def test_pdg_and_material_labels():
assert G.pdg_label(22) == "gamma"
assert G.pdg_label(999999) == "999999"
assert G.material_label("G4_PbWO4") == "PbWO4"
tests/test_catalog.py
+69
View File
@@ -0,0 +1,69 @@
"""Tests for the plot catalog: id uniqueness + every spec computes a valid Reduced."""
from __future__ import annotations
import pytest
from giant.analysis import build_catalog, catalog_ids, get_spec
from giant.analysis.catalog import Bundle
from giant.analysis.context import build_context
from tests.test_analysis_reduce import _reference_frame, _rollout_frame
@pytest.fixture(scope="module")
def bundle() -> Bundle:
r, t = _rollout_frame(), _reference_frame()
ctx = build_context(
r, t, n_energy_bins=2, n_marginal_bins=10, top_k_pdg=3, sample_rows=1000
)
return Bundle.open(r, t, ctx)
def test_catalog_ids_unique_and_nonempty():
ids = catalog_ids()
assert ids and len(ids) == len(set(ids))
# the required families are all present
fams = {s.family for s in build_catalog()}
assert {"marginals", "event", "shower", "species", "secondaries"} <= fams
def test_get_spec_roundtrip_and_unknown():
spec = get_spec("marginal_edep")
assert spec.id == "marginal_edep" and spec.family == "marginals"
with pytest.raises(KeyError):
get_spec("does_not_exist")
def test_every_spec_computes_valid_reduced(bundle: Bundle):
for spec in build_catalog():
r = spec.compute(bundle)
assert r.id == spec.id
assert r.kind in {
"overlay_hist",
"grouped_hist",
"profile",
"bar",
"single_hist",
}
assert r.title and r.xlabel
_validate_payload(r)
def _validate_payload(r) -> None:
p = r.payload
if r.kind == "overlay_hist":
n = len(p["edges"]) - 1
assert len(p["rollout"]) == n and len(p["reference"]) == n
elif r.kind == "single_hist":
assert len(p["rollout"]) == len(p["edges"]) - 1
elif r.kind == "grouped_hist":
n = len(p["edges"]) - 1
assert p["groups"], "grouped hist must have at least one group"
for g in p["groups"].values():
assert len(g["rollout"]) == n and len(g["reference"]) == n
elif r.kind == "profile":
n = len(p["edges"]) - 1
for k in ("rollout_mean", "rollout_std", "reference_mean", "reference_std"):
assert len(p[k]) == n
elif r.kind == "bar":
assert len(p["labels"]) == len(p["rollout"]) == len(p["reference"])
tests/test_condor.py
+86
View File
@@ -0,0 +1,86 @@
"""Tests for the HTCondor submit description + the compute_one round-trip."""
from __future__ import annotations
from pathlib import Path
from giant.analysis import (
Context,
SubmitConfig,
build_context,
catalog_ids,
compute_one,
prep,
write_submit,
)
from giant.analysis.reduced import Reduced
from tests.test_analysis_reduce import _reference_frame, _rollout_frame
def test_prep_and_compute_one_roundtrip(tmp_path: Path):
r, t = _rollout_frame(), _reference_frame()
ctx = build_context(
r, t, n_energy_bins=2, n_marginal_bins=8, top_k_pdg=3, sample_rows=1000
)
shared = tmp_path / "shared.json"
ctx.save(shared)
assert Context.load(shared).top_pdgs == ctx.top_pdgs
out = compute_one("marginal_edep", r, t, shared, tmp_path / "marginal_edep.json")
reduced = Reduced.load(out)
assert reduced.id == "marginal_edep"
assert len(reduced.payload["rollout"]) == len(reduced.payload["edges"]) - 1
def test_prep_writes_shared_json(tmp_path: Path):
r, t = _rollout_frame(), _reference_frame()
shared = prep(
r,
t,
tmp_path / "run",
n_energy_bins=2,
n_marginal_bins=8,
top_k_pdg=3,
sample_rows=1000,
)
assert shared.exists()
ctx = Context.load(shared)
assert set(ctx.var_ranges) == {"step_length", "edep", "delta_e", "post_E"}
def test_write_submit_description(tmp_path: Path):
cfg = SubmitConfig(
rollout=tmp_path / "r.parquet",
reference=tmp_path / "t.parquet",
out_dir=tmp_path / "run",
accounting_group="cms",
repo_dir=tmp_path,
)
sub = write_submit(cfg)
txt = sub.read_text()
assert "universe = docker" in txt
assert "docker_image = mschnepf/slc7-condocker" in txt
assert "requirements = TARGET.ProvidesETPResources" in txt
assert "accounting_group = cms" in txt
assert "queue plotid from" in txt
# one queue item per catalog id
ids = (cfg.out_dir / "plotids.txt").read_text().split()
assert ids == catalog_ids()
# wrapper is executable and self-contained
wrapper = cfg.out_dir / "run_compute.sh"
assert wrapper.exists() and (wrapper.stat().st_mode & 0o111)
assert "giant analyze compute-one" in wrapper.read_text()
def test_write_submit_remote_flag(tmp_path: Path):
cfg = SubmitConfig(
rollout=tmp_path / "r.parquet",
reference=tmp_path / "t.parquet",
out_dir=tmp_path / "run",
accounting_group="cms",
repo_dir=tmp_path,
remote=True,
)
txt = write_submit(cfg).read_text()
assert "+RemoteJob = True" in txt
assert "ProvidesETPResources" not in txt
tests/test_render.py
+92
View File
@@ -0,0 +1,92 @@
"""Render smoke test — skipped where plotstyle / LaTeX is unavailable."""
from __future__ import annotations
from pathlib import Path
import pytest
pytest.importorskip("plotstyle")
from giant.analysis.reduced import Reduced # noqa: E402
def _try_render(reduced: list[Reduced], out: Path):
from giant.analysis.render import render_all
for r in reduced:
r.save(out / "reduced" / f"{r.id}.json")
return render_all(out / "reduced", out / "plots")
def test_render_one_of_each_kind(tmp_path: Path):
reduced = [
Reduced(
"m",
"marginals",
"overlay_hist",
"Overlay",
"x",
{
"edges": [0, 1, 2, 3],
"rollout": [1, 2, 3],
"reference": [3, 2, 1],
"log_y": False,
},
),
Reduced(
"g",
"marginals",
"grouped_hist",
"Grouped",
"x",
{
"edges": [0, 1, 2],
"groups": {"a": {"rollout": [1, 2], "reference": [2, 1]}},
"log_y": False,
},
),
Reduced(
"p",
"shower",
"profile",
"Profile",
"depth",
{
"edges": [0, 1, 2],
"rollout_mean": [1, 2],
"rollout_std": [0.1, 0.2],
"reference_mean": [1.1, 1.9],
"reference_std": [0.1, 0.1],
"ylabel": "e",
},
),
Reduced(
"b",
"species",
"bar",
"Bar",
"species",
{
"labels": ["e-", "gamma"],
"rollout": [0.6, 0.4],
"reference": [0.5, 0.5],
"ylabel": "frac",
},
),
Reduced(
"s",
"species",
"single_hist",
"Single",
"x",
{"edges": [0, 1, 2], "rollout": [5, 1], "log_y": True},
),
]
try:
pdfs = _try_render(reduced, tmp_path)
except RuntimeError as e: # LaTeX missing at render time
pytest.skip(f"LaTeX rendering unavailable: {e}")
assert len(pdfs) == len(reduced)
assert all(p.exists() for p in pdfs)
assert (tmp_path / "plots" / "metadata.yaml").exists()
uv.lock
Generated
+121
View File
@@ -45,6 +45,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/81/29/5ecc3a15d5a33e31b26c11426c45c501e439cb865d0bff96315d86443b78/appnope-0.1.4-py2.py3-none-any.whl", hash = "sha256:502575ee11cd7a28c0205f379b525beefebab9d161b7c964670864014ed7213c", size = 4321, upload-time = "2024-02-06T09:43:09.663Z" },
]
[[package]]
name = "argcomplete"
version = "3.7.0"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/95/c0/c8e94135e66fabf89a120d9b4b123fe6993506beca6c1938a74c24cfa5fd/argcomplete-3.7.0.tar.gz", hash = "sha256:afde224f753f874807b1dc1414e883ab8fe0cda9c04807b6047dcb8e1ac23913", size = 73284, upload-time = "2026-06-30T22:28:22.249Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/12/f6/5b8ec087cd9cfa9449491ec83f76fb6b7006b4dff57d2ba8aaab330fe8e4/argcomplete-3.7.0-py3-none-any.whl", hash = "sha256:d8f0f22d2a8a7caa383be1e22b6caf1ecaf0ebd10d8f83cc125e36540c95830c", size = 42575, upload-time = "2026-06-30T22:28:20.547Z" },
]
[[package]]
name = "asttokens"
version = "3.0.1"
@@ -441,6 +450,43 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/e5/22/4222d7ddf3da30f363edaa98e329c2bce6c65497c9cb2810931c8b2c0fbc/fsspec-2026.6.0-py3-none-any.whl", hash = "sha256:02e0b71817df9b2169dc30a16832045764def1191b43dcff5bb85bdee212d2a1", size = 203949, upload-time = "2026-06-16T01:57:26.358Z" },
]
[[package]]
name = "gallery"
version = "0.1.3"
source = { editable = "../ETPlot" }
dependencies = [
{ name = "argcomplete" },
{ name = "jinja2" },
{ name = "platformdirs" },
{ name = "pymupdf" },
{ name = "pytest" },
{ name = "pyyaml" },
{ name = "textual" },
]
[package.optional-dependencies]
plotting = [
{ name = "matplotlib" },
]
[package.metadata]
requires-dist = [
{ name = "argcomplete", specifier = ">=3.0" },
{ name = "jinja2", specifier = ">=3.0.0" },
{ name = "matplotlib", marker = "extra == 'plotting'", specifier = ">=3.7" },
{ name = "pip-audit", marker = "extra == 'dev'", specifier = ">=2.7" },
{ name = "platformdirs", specifier = ">=3.0" },
{ name = "pymupdf", specifier = ">=1.23" },
{ name = "pytest" },
{ name = "pytest", marker = "extra == 'dev'", specifier = ">=7.0" },
{ name = "pyyaml", specifier = ">=5.0" },
{ name = "ruff", marker = "extra == 'dev'", specifier = ">=0.6" },
{ name = "setuptools", marker = "extra == 'dev'", specifier = ">=83.0.0" },
{ name = "textual", specifier = ">=0.50" },
{ name = "ty", marker = "extra == 'dev'", specifier = ">=0.0.1" },
]
provides-extras = ["dev", "plotting"]
[[package]]
name = "giant"
version = "0.1.0"
@@ -457,6 +503,7 @@ dependencies = [
[package.optional-dependencies]
analysis = [
{ name = "gallery", extra = ["plotting"] },
{ name = "ipykernel" },
{ name = "matplotlib" },
{ name = "polars" },
@@ -475,6 +522,7 @@ cuda = [
]
dev = [
{ name = "awkward" },
{ name = "gallery", extra = ["plotting"] },
{ name = "ipykernel" },
{ name = "matplotlib" },
{ name = "polars" },
@@ -491,6 +539,7 @@ geometry = [
[package.metadata]
requires-dist = [
{ name = "awkward", marker = "extra == 'convert'", specifier = ">=2.6,<3" },
{ name = "gallery", extras = ["plotting"], marker = "extra == 'analysis'", editable = "../ETPlot" },
{ name = "giant", extras = ["convert", "analysis", "geometry"], marker = "extra == 'dev'" },
{ name = "ipykernel", marker = "extra == 'analysis'", specifier = ">=7.3.0" },
{ name = "matplotlib", marker = "extra == 'analysis'", specifier = ">=3.8,<4" },
@@ -750,6 +799,18 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/b5/91/53255615acd2a1eaca307ede3c90eb550bae9c94581f8c00081b6b1c8f44/kiwisolver-1.5.0-graalpy312-graalpy250_312_native-win_amd64.whl", hash = "sha256:1f1489f769582498610e015a8ef2d36f28f505ab3096d0e16b4858a9ec214f57", size = 75987, upload-time = "2026-03-09T13:15:39.65Z" },
]
[[package]]
name = "linkify-it-py"
version = "2.1.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "uc-micro-py" },
]
sdist = { url = "https://files.pythonhosted.org/packages/2e/c9/06ea13676ef354f0af6169587ae292d3e2406e212876a413bf9eece4eb23/linkify_it_py-2.1.0.tar.gz", hash = "sha256:43360231720999c10e9328dc3691160e27a718e280673d444c38d7d3aaa3b98b", size = 29158, upload-time = "2026-03-01T07:48:47.683Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/b4/de/88b3be5c31b22333b3ca2f6ff1de4e863d8fe45aaea7485f591970ec1d3e/linkify_it_py-2.1.0-py3-none-any.whl", hash = "sha256:0d252c1594ecba2ecedc444053db5d3a9b7ec1b0dd929c8f1d74dce89f86c05e", size = 19878, upload-time = "2026-03-01T07:48:46.098Z" },
]
[[package]]
name = "markdown-it-py"
version = "4.2.0"
@@ -762,6 +823,11 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/b3/81/4da04ced5a082363ecfa159c010d200ecbd959ae410c10c0264a38cac0f5/markdown_it_py-4.2.0-py3-none-any.whl", hash = "sha256:9f7ebbcd14fe59494226453aed97c1070d83f8d24b6fc3a3bcf9a38092641c4a", size = 91687, upload-time = "2026-05-07T12:08:27.182Z" },
]
[package.optional-dependencies]
linkify = [
{ name = "linkify-it-py" },
]
[[package]]
name = "markupsafe"
version = "3.0.3"
@@ -891,6 +957,18 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/41/09/5b161152e2d90f7b87f781c2e1267494aef9c32498df793f73ad0a0a494a/matplotlib_inline-0.2.2-py3-none-any.whl", hash = "sha256:3c821cf1c209f59fb2d2d64abbf5b23b67bcb2210d663f9918dd851c6da1fcf6", size = 9534, upload-time = "2026-05-08T17:33:32.055Z" },
]
[[package]]
name = "mdit-py-plugins"
version = "0.6.1"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "markdown-it-py" },
]
sdist = { url = "https://files.pythonhosted.org/packages/59/fc/f8d0863f8862f25602c0404d75568e89fb6b4109804645e5cdfb1be5cf56/mdit_py_plugins-0.6.1.tar.gz", hash = "sha256:a2bca0f039f39dbd35fb74ae1b5f998608c437463371f0ff7f49a19a17a114d0", size = 56114, upload-time = "2026-05-13T09:03:38.91Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/a5/69/6da5581c6a7fede7dc261bf4e67d6adca4196f176b43288b55b3db395b6e/mdit_py_plugins-0.6.1-py3-none-any.whl", hash = "sha256:214c82fb2ac524472ab6a5bcab1de80f73b50443e187f401bfd77efbc7c6481d", size = 66663, upload-time = "2026-05-13T09:03:37.76Z" },
]
[[package]]
name = "mdurl"
version = "0.1.2"
@@ -1463,6 +1541,23 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/f4/7e/a72dd26f3b0f4f2bf1dd8923c85f7ceb43172af56d63c7383eb62b332364/pygments-2.20.0-py3-none-any.whl", hash = "sha256:81a9e26dd42fd28a23a2d169d86d7ac03b46e2f8b59ed4698fb4785f946d0176", size = 1231151, upload-time = "2026-03-29T13:29:30.038Z" },
]
[[package]]
name = "pymupdf"
version = "1.28.0"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/8e/e9/6d6c5d6c0a3551bffd47681a6240caf941727f195b45593cf20ab36f018f/pymupdf-1.28.0.tar.gz", hash = "sha256:e53f3567403a92da15caa9e7ae0164327fff48817e9f40175367fb9de524258d", size = 87637751, upload-time = "2026-06-29T09:08:47.547Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/c8/b7/88043e38cc7529de070f0c9bd267fa258035cca0b4ad5260536b994594a7/pymupdf-1.28.0-cp310-abi3-macosx_10_15_x86_64.whl", hash = "sha256:892b89ba88e8f98b53133b62877a9dc9b5e7dc6a4aeb837b612db56a8d2e03ac", size = 24597385, upload-time = "2026-06-29T09:03:30.608Z" },
{ url = "https://files.pythonhosted.org/packages/33/f4/23775bbda0781b61fc398cc75079a2b0e64696d8fcf93271748883e9627e/pymupdf-1.28.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:4d692dcf44d3566ae96bc6f6346c6ad432274a29ba617bf7a9fe18009e24adb4", size = 23828292, upload-time = "2026-06-29T09:03:46.129Z" },
{ url = "https://files.pythonhosted.org/packages/1c/f5/bf75fc7a415722f8b33662054f82d88520c0cbfd4c36d0e08aeaec605e49/pymupdf-1.28.0-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:47a5c29ed4eb0744de9c4e37bb49b1259b18d4d75fcc8a7c130f7c9fa15956f6", size = 25045507, upload-time = "2026-06-29T09:04:03.86Z" },
{ url = "https://files.pythonhosted.org/packages/58/69/5d12c9f1f2d76f28383d6110a069c79fbfced5a4f97bb1ee6e8354f52bb7/pymupdf-1.28.0-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:44f0973f5e5edbaec95bc34b64e71d1959d4ee90b1328de1b4f4f5b4fa78673f", size = 25716599, upload-time = "2026-06-29T09:04:19.367Z" },
{ url = "https://files.pythonhosted.org/packages/4d/b4/ec0e017bc42857cc86bd651441dbc41cc18be48d4698ecd27aac491e0c9a/pymupdf-1.28.0-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:4d61ec323a706e153a12e262e51febfb43eeaa20977785ace135d18d48bcdc83", size = 25940489, upload-time = "2026-06-29T09:04:36.624Z" },
{ url = "https://files.pythonhosted.org/packages/06/86/f831fef09013f33b3c9c09fb3923f2ff53e1e437f6ace14b8ae46392f558/pymupdf-1.28.0-cp310-abi3-win32.whl", hash = "sha256:caea2b3b67347fd79e5d15ed7929b0e886aac594ea228073b6d39de0078189da", size = 18489703, upload-time = "2026-06-29T20:50:30.599Z" },
{ url = "https://files.pythonhosted.org/packages/2e/5d/1a03f53eb0449900469335fcfc742ca28e3ba159b7d650e0921d50b8b308/pymupdf-1.28.0-cp310-abi3-win_amd64.whl", hash = "sha256:e01e90fd86abfeb37ceb921eddb951f988a11d45ff6ce6b7664f2039849068ec", size = 19773102, upload-time = "2026-06-29T09:04:49.773Z" },
{ url = "https://files.pythonhosted.org/packages/72/f6/1e52ce243ca792254f6223b4017c5667194c146ce9b88baf37bc5eb3d1c9/pymupdf-1.28.0-cp313-abi3-pyemscripten_2025_0_wasm32.whl", hash = "sha256:74c6d00ba2a9aad3a635db73b07c15db462b480741d831a34a75a56535ebc22b", size = 18357011, upload-time = "2026-06-29T20:50:50.353Z" },
{ url = "https://files.pythonhosted.org/packages/62/b1/46b5b3d8ef3cc71114667cf10c4d8b33f39af97253af32e9a0986775b638/pymupdf-1.28.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:b3e1399c7a64c6914239116a369efcdaac4cfb9e838bde2656d7accc4a85c72d", size = 25753599, upload-time = "2026-06-29T09:05:09.398Z" },
]
[[package]]
name = "pyparsing"
version = "3.3.2"
@@ -1772,6 +1867,23 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/9b/24/84ce997e8ae6296168a74d0d9c4dde572d90fb23fd7c0b219c30ff71e00e/tbb-2021.13.1-py3-none-win_amd64.whl", hash = "sha256:cbf024b2463fdab3ebe3fa6ff453026358e6b903839c80d647e08ad6d0796ee9", size = 286908, upload-time = "2024-08-07T15:09:05.677Z" },
]
[[package]]
name = "textual"
version = "8.2.8"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "markdown-it-py", extra = ["linkify"] },
{ name = "mdit-py-plugins" },
{ name = "platformdirs" },
{ name = "pygments" },
{ name = "rich" },
{ name = "typing-extensions" },
]
sdist = { url = "https://files.pythonhosted.org/packages/00/21/39a76b01bd5eea82a04baaca7580e105d8c59450df03998345bb2cfb307b/textual-8.2.8.tar.gz", hash = "sha256:3f106a9fbc73e39dd266c9712432087de78a6d644084c7c241d6a25c3169115b", size = 1860502, upload-time = "2026-06-30T06:51:24.495Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/fb/be/35261223d9416a0751cdff1c7b4a6f881387218a12d439fe22fefebc8c04/textual-8.2.8-py3-none-any.whl", hash = "sha256:267375fd402dc8d981457212efa71f0e3365fd17bba144ba9bb3ed7563cb374a", size = 731418, upload-time = "2026-06-30T06:51:26.364Z" },
]
[[package]]
name = "threadpoolctl"
version = "3.6.0"
@@ -1960,6 +2072,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/ce/e4/dccd7f47c4b64213ac01ef921a1337ee6e30e8c6466046018326977efd95/tzdata-2026.2-py2.py3-none-any.whl", hash = "sha256:bbe9af844f658da81a5f95019480da3a89415801f6cc966806612cc7169bffe7", size = 349321, upload-time = "2026-04-24T15:22:05.876Z" },
]
[[package]]
name = "uc-micro-py"
version = "2.0.0"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/78/67/9a363818028526e2d4579334460df777115bdec1bb77c08f9db88f6389f2/uc_micro_py-2.0.0.tar.gz", hash = "sha256:c53691e495c8db60e16ffc4861a35469b0ba0821fe409a8a7a0a71864d33a811", size = 6611, upload-time = "2026-03-01T06:31:27.526Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/61/73/d21edf5b204d1467e06500080a50f79d49ef2b997c79123a536d4a17d97c/uc_micro_py-2.0.0-py3-none-any.whl", hash = "sha256:3603a3859af53e5a39bc7677713c78ea6589ff188d70f4fee165db88e22b242c", size = 6383, upload-time = "2026-03-01T06:31:26.257Z" },
]
[[package]]
name = "uproot"
version = "5.7.4"