2 Commits

Author SHA1 Message Date
lars 7df1945384 perf: replace pandas with polars in the setup-stage scan
CI / Sync project version with tag (pull_request) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 1m6s
CI / Format (ruff format) (pull_request) Successful in 1m11s
CI / Type check (ty) (pull_request) Successful in 1m12s
CI / Tests (pull_request) Successful in 4m16s
CI / Bump version, tag, and update changelog on merge to master (pull_request) Has been skipped
CI / Publish package to Gitea package registry (pull_request) Has been skipped
giant.pipeline.run_setup_stage (used by both giant train and dwarf
warm-cache) previously opened and fully read each parquet file 4-6
separate times via pandas, with per-row Python loops padding the
secondary list columns on every chunk of the normalizer-fitting pass.

- giant/data/loader.py: pandas -> polars throughout; ragged sec_*_list
  padding is now a single vectorized polars expression instead of a
  per-row Python loop (including a .iloc[i] loop for directions).
- giant/data/scan.py (new): a fused metadata scan answering the event
  index, pdg/material vocab, process counts, and pooled-pdg counts in
  one pass per file instead of one pass per section. Frequency-ranking
  ties are now an explicit (-count, first_seen) contract instead of an
  accident of pandas' value_counts iteration order.
- giant/pipeline.py: run_setup_stage restructured to consult the cache
  for every section first, then issue one combined scan request for
  whatever's missing.
- giant/geometry.py: ported the one remaining pandas groupby to polars.
- pyproject.toml: polars promoted to a core dependency, pandas moved
  to dev (only test fixtures still use it).
- giant/tools/profile_setup_scan.py (new): synthetic-data benchmark
  for this scan, mirroring profile_analysis_costs.py's pattern.

Also fixes a real deadlock this surfaced: DataLoader worker
subprocesses fork() on Linux, and polars' native thread pool doesn't
survive a fork — a worker touching polars after the parent already had
hangs instantly. giant/pipeline.py's train/val DataLoaders now use
multiprocessing_context="spawn" whenever num_workers>0.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01DdT32YWNEwnVLZUHsgdeSC
2026-09-02 09:59:37 +02:00
lars deb9e8e7de feat: add WGAN + AR stop-token config variant
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (push) Successful in 1m23s
CI / Format (ruff format) (push) Successful in 1m23s
CI / Type check (ty) (push) Successful in 1m24s
CI / Tests (push) Successful in 3m32s
CI / Publish package to Gitea package registry (push) Has been skipped
CI / Bump version, tag, and update changelog on merge to master (push) Successful in 8s
Branches off configs/baseline.toml with both stages on WGAN-GP and
stage-2 n_sec.mode = stop_token, 30 epochs — combines two unbenchmarked
roadmap axes (post-v0.3.0 WGAN, and the AR stop-token multiplicity mode)
into one variant that stays a single edit away from baseline for
attribution.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01KzPghrmFcAJYrWUvHApY9N
2026-08-31 14:32:26 +02:00
11 changed files with 713 additions and 193 deletions
+122
View File
@@ -0,0 +1,122 @@
# GIANT WGAN-GP + AR stop-token variant of configs/baseline.toml.
#
# Two roadmap axes, combined into one run: WGAN-GP generators for both
# stages (unbenchmarked since the 2026-08-03 pre-v0.3.0 failure, which was
# secondary-species mode collapse — the failure v0.3.0's AR/categorical
# pivot exists to fix) and the AR stop-token multiplicity mode
# (stage2_model.n_sec.mode = "stop_token", never benchmarked at all).
# Everything else is byte-identical to baseline.toml so a rollout compared
# against baseline's analysis_341dfb14 is attributable to these two axes
# alone: conditioning (physical/physical), hidden_dim 512 / n_res_blocks 6 /
# dropout 0.0 per stage, k_max 15, history "markov", teacher_forcing
# "always", particle_type.target "onehot" (n_classes 32, other_policy
# "sample"), lr 3e-4, warmup_epochs 3, weight_decay 0.01, ema_decay 0.9999,
# val_fraction 0.1, num_workers 4, seed 0, validate_steps 10, W&B on.
#
# No [stage1_model.wgan] / [stage2_model.wgan] block: the dataclass defaults
# (noise_dim 64, n_critic 5, gp_weight 10.0, critic_lr 0.0 = inherit
# train.lr, critic_hidden_dim/critic_n_res_blocks 0 = inherit the stage's
# 512/6, stage 2's gumbel_tau_start/_end 1.0/0.1) are what the earlier WGAN
# runs used — writing them out would add keys that don't vary.
#
# particle_type.class_weighting stays "none" (the default): config.py's
# validate_config rejects any other value under stage2_model.generator =
# "wgan", since that path feeds the type slice to the critic via a
# straight-through Gumbel relaxation instead of a weighted cross-entropy.
#
# Prior WGAN writeup (pre-v0.3.0, describes the failure this run re-tests):
# /home/lars/knowledge-base/experiments/giant-wgan-physical-rollout-validation.md
[meta]
# REQUIRED. Without it config.migrate_config reads this file as v0.2 and
# rewrites it from V02_FIXED_FACTS — silently forcing decoder = "one_shot",
# particle_type.target = "physical" and the v0.2 default sizes, while still
# passing validate_config.
config_version = 3
[conditioning]
# Physical-property MLPs rather than learned vocab embeddings: computable for
# any PDG code / material, which is what the held-out-species and
# held-out-material generalization comparisons need.
out_dim = 128
share_stages = false
# n_layers = 2 rather than the v0.3 default of 1: v0.2's conditioning MLP was
# always 2 deep (see _migration.V02_FIXED_FACTS), so this keeps the encoder
# identical to baseline.toml.
[conditioning.particle]
type = "physical"
emb_dim = 16
n_layers = 2
[conditioning.material]
type = "physical"
emb_dim = 16
n_layers = 2
[stage1_model]
generator = "wgan"
hidden_dim = 512
n_res_blocks = 6
dropout = 0.0
[stage2_model]
# Autoregressive in descending-energy order, as baseline.toml — this variant
# only swaps the generator (flow -> wgan) and the multiplicity mode
# (head -> stop_token), not the decoder shape.
decoder = "autoregressive"
generator = "wgan"
hidden_dim = 512
n_res_blocks = 6
dropout = 0.0
k_max = 15
[stage2_model.autoregressive]
history = "markov"
teacher_forcing = "always"
[stage2_model.n_sec]
# EOS-style per-slot stop head on the AR secondary decoder, replacing the
# n_sec classifier entirely (mutually exclusive — see NSecConfig's
# docstring in giant/config.py). Requires decoder = "autoregressive" and
# owner = "stage2" (both already true above/by default); validate_config
# enforces this.
mode = "stop_token"
[stage2_model.particle_type]
target = "onehot"
# Decoupled from conditioning.particle.emb_dim (gitea #29). 32 classes + the
# "other" bucket keeps essentially all real secondary species out of "other"
# without making the head expensive.
n_classes = 32
other_policy = "sample"
[train]
epochs = 30
# Halved from baseline's 36864. That figure came from a measured linear fit
# of the *flow-AR* training step (peak reserved MiB = 0.9736 * batch_size +
# 115); WGAN invalidates it twice over — each stage gains a critic that by
# default inherits the stage's own 512/6 body, and gradient_penalty
# (giant/model/wgan.py, forced fp32 internally) runs a double-backward every
# batch. 18432 is a conservative choice pending a real memory measurement on
# this exact config, not a re-derived fit. Throughput is already flat above
# bs~4096 on the 4070, so this costs occupancy on the L40S, not step
# efficiency.
batch_size = 18432
lr = 3e-4
warmup_epochs = 3
weight_decay = 0.01
ema_decay = 0.9999
val_fraction = 0.1
num_workers = 4
seed = 0
# Tightened from baseline's 10: WGANStageTrainer.supports_val_loss = False,
# and with both stages adversarial there is no per-epoch val loss at all, so
# validate_every's marginal-KL pass (giant/training/trainers.py's
# val_objective) is the only comparable-across-epochs best-checkpoint
# selection signal available. 5 gives 6 evaluations over 30 epochs instead
# of baseline's 3, at ~6x5000s of extra walltime.
validate_every = 5
validate_steps = 10
wandb = true
wandb_project = "giant"
+92 -115
View File
@@ -1,13 +1,16 @@
from dataclasses import dataclass, field
from pathlib import Path
from typing import Iterator
from typing import TYPE_CHECKING, Any, Iterator, Mapping
import numpy as np
import pandas as pd
import polars as pl
import pyarrow.parquet as pq
from giant.constants import K_MAX
if TYPE_CHECKING:
from giant.data.scan import ValueStat
# A manifest is a plain text file listing one parquet path per line, used to
# name a curated subset of files (e.g. a train/holdout pool) without copying
# or symlinking the underlying parquet files. Lines are resolved relative to
@@ -77,88 +80,75 @@ def find_parquet_files(path: str | Path) -> list[Path]:
return [p]
def _pad_list_col(series: pd.Series, K: int, fill: float = 0.0) -> np.ndarray:
"""Pad / truncate a list-valued Series to fixed width K → (N, K) float32."""
out = np.full((len(series), K), fill, dtype=np.float32)
for i, lst in enumerate(series):
if lst is not None and len(lst) > 0:
n = min(len(lst), K)
out[i, :n] = lst[:n]
return out
def _pad_list_column(df: pl.DataFrame, col: str, k: int, fill, dtype: type[pl.DataType] | pl.DataType) -> np.ndarray:
"""Pad / truncate a list-valued column to fixed width `k` → (N, k) numpy array.
def _pad_list_col_int(series: pd.Series, K: int, fill: int = 0) -> np.ndarray:
"""Pad / truncate a list-valued integer Series to fixed width K → (N, K) int64."""
out = np.full((len(series), K), fill, dtype=np.int64)
for i, lst in enumerate(series):
if lst is not None and len(lst) > 0:
n = min(len(lst), K)
out[i, :n] = lst[:n]
return out
def _pad_dir_col(dx: pd.Series, dy: pd.Series, dz: pd.Series, K: int) -> np.ndarray:
"""Pad three list-valued direction columns → (N, K, 3) float32.
Padding direction defaults to (0,0,1) (forward) so it is a valid unit vector.
Concatenating `k` fill values before truncating to `k` guarantees every
row ends up with exactly `k` non-null elements regardless of how short
(including empty) or long the original list was, so `list.to_array(k)`
(a fixed-size-array dtype) converts to a plain 2D numpy array with a
single vectorized expression — no per-row Python loop.
"""
N = len(dx)
out = np.zeros((N, K, 3), dtype=np.float32)
out[:, :, 2] = 1.0
for i in range(N):
lx, ly, lz = dx.iloc[i], dy.iloc[i], dz.iloc[i]
if lx is not None and len(lx) > 0:
n = min(len(lx), K)
out[i, :n, 0] = lx[:n]
out[i, :n, 1] = ly[:n]
out[i, :n, 2] = lz[:n]
return out
fill_tail = pl.lit([fill] * k, dtype=pl.List(dtype))
out = df.select(pl.col(col).cast(pl.List(dtype)).list.concat(fill_tail).list.head(k).list.to_array(k).alias("_p"))
return out["_p"].to_numpy()
def _df_to_dict(df: pd.DataFrame, offset: int = 0, k_max: int = K_MAX) -> dict[str, np.ndarray]:
def _pad_dir_col(df: pl.DataFrame, dx: str, dy: str, dz: str, k: int) -> np.ndarray:
"""Pad three list-valued direction columns → (N, k, 3) float32.
Padding direction defaults to (0, 0, 1) (forward) so it is a valid unit vector.
"""
px = _pad_list_column(df, dx, k, 0.0, pl.Float64)
py = _pad_list_column(df, dy, k, 0.0, pl.Float64)
pz = _pad_list_column(df, dz, k, 1.0, pl.Float64)
return np.stack([px, py, pz], axis=-1).astype(np.float32)
def _df_to_dict(df: pl.DataFrame, offset: int = 0, k_max: int = K_MAX) -> dict[str, np.ndarray]:
has_sec_lists = "sec_E_list" in df.columns
d: dict[str, np.ndarray] = {
"event_id": _offset_event_id(df["event_id"].to_numpy(), offset),
"pdg": df["pdg"].to_numpy(dtype=np.int32),
"pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
"pre_E": df["pre_E"].to_numpy(dtype=np.float32),
"pre_dir": df[["pre_dx", "pre_dy", "pre_dz"]].to_numpy(dtype=np.float32),
"material": df["material"].to_numpy(dtype=object),
"layer_id": df["layer_id"].to_numpy(dtype=np.int32),
"n_sec": df["child_track_ids"].apply(len).to_numpy(dtype=np.int32),
"e_sec": df["e_sec"].to_numpy(dtype=np.float32),
"pdg": df["pdg"].to_numpy().astype(np.int32),
"pre_pos": df.select(["pre_x", "pre_y", "pre_z"]).to_numpy().astype(np.float32),
"pre_E": df["pre_E"].to_numpy().astype(np.float32),
"pre_dir": df.select(["pre_dx", "pre_dy", "pre_dz"]).to_numpy().astype(np.float32),
"material": df["material"].to_numpy().astype(object),
"layer_id": df["layer_id"].to_numpy().astype(np.int32),
"n_sec": df["child_track_ids"].list.len().to_numpy().astype(np.int32),
"e_sec": df["e_sec"].to_numpy().astype(np.float32),
# The physics process that ended the step (e.g. "compt", "phot",
# "eBrem") — a post-step outcome, so it's a router/classifier
# supervision label only, never conditioning (see build_process_map*
# / ProcessRouter). Guarded like has_sec_lists: older parquet
# conversions predating this column still load fine.
"process": (
df["process"].to_numpy(dtype=object) if "process" in df.columns else np.full(len(df), "", dtype=object)
df["process"].to_numpy().astype(object) if "process" in df.columns else np.full(len(df), "", dtype=object)
),
"step_length": df["step_length"].to_numpy(dtype=np.float32),
"post_E": df["post_E"].to_numpy(dtype=np.float32),
"delta_e": (df["pre_E"] - df["post_E"]).to_numpy(dtype=np.float32),
"edep": df["edep"].to_numpy(dtype=np.float32),
"post_dir": df[["post_dx", "post_dy", "post_dz"]].to_numpy(dtype=np.float32),
"post_pos": df[["post_x", "post_y", "post_z"]].to_numpy(dtype=np.float32),
"step_length": df["step_length"].to_numpy().astype(np.float32),
"post_E": df["post_E"].to_numpy().astype(np.float32),
"delta_e": (df["pre_E"] - df["post_E"]).to_numpy().astype(np.float32),
"edep": df["edep"].to_numpy().astype(np.float32),
"post_dir": df.select(["post_dx", "post_dy", "post_dz"]).to_numpy().astype(np.float32),
"post_pos": df.select(["post_x", "post_y", "post_z"]).to_numpy().astype(np.float32),
}
if has_sec_lists:
d["sec_E_list"] = _pad_list_col(df["sec_E_list"], k_max)
d["sec_pdg_list"] = _pad_list_col_int(df["sec_pdg_list"], k_max)
d["sec_dir_list"] = _pad_dir_col(df["sec_dx_list"], df["sec_dy_list"], df["sec_dz_list"], k_max)
d["sec_E_list"] = _pad_list_column(df, "sec_E_list", k_max, 0.0, pl.Float64).astype(np.float32)
d["sec_pdg_list"] = _pad_list_column(df, "sec_pdg_list", k_max, 0, pl.Int64).astype(np.int64)
d["sec_dir_list"] = _pad_dir_col(df, "sec_dx_list", "sec_dy_list", "sec_dz_list", k_max)
return d
def load_steps(path: str | Path, offset: int = 0, k_max: int = K_MAX) -> dict[str, np.ndarray]:
return _df_to_dict(pd.read_parquet(path), offset=offset, k_max=k_max)
return _df_to_dict(pl.read_parquet(path), offset=offset, k_max=k_max)
def load_event_ids(path: str | Path, offset: int = 0) -> np.ndarray:
"""Read only the event_id column — cheap scan for split assignment."""
ids = pd.read_parquet(path, columns=["event_id"])["event_id"].to_numpy()
ids = pl.read_parquet(path, columns=["event_id"])["event_id"].to_numpy()
return _offset_event_id(ids, offset)
@@ -170,7 +160,7 @@ def iter_file_chunks(path: str | Path, offset: int = 0, k_max: int = K_MAX) -> I
module constant for callers that don't care (e.g. Stage-1-only reads)."""
pf = pq.ParquetFile(path)
for i in range(pf.num_row_groups):
yield _df_to_dict(pf.read_row_group(i).to_pandas(), offset=offset, k_max=k_max)
yield _df_to_dict(pl.DataFrame(pf.read_row_group(i)), offset=offset, k_max=k_max)
_COND_COLS = [
@@ -190,17 +180,17 @@ _COND_COLS = [
]
def _cond_df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]:
def _cond_df_to_dict(df: pl.DataFrame, offset: int = 0) -> dict[str, np.ndarray]:
return {
"event_id": _offset_event_id(df["event_id"].to_numpy(), offset),
"pdg": df["pdg"].to_numpy(dtype=np.int32),
"pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
"pre_E": df["pre_E"].to_numpy(dtype=np.float32),
"pre_dir": df[["pre_dx", "pre_dy", "pre_dz"]].to_numpy(dtype=np.float32),
"material": df["material"].to_numpy(dtype=object),
"layer_id": df["layer_id"].to_numpy(dtype=np.int32),
"n_sec": df["child_track_ids"].apply(len).to_numpy(dtype=np.int32),
"e_sec": df["e_sec"].to_numpy(dtype=np.float32),
"pdg": df["pdg"].to_numpy().astype(np.int32),
"pre_pos": df.select(["pre_x", "pre_y", "pre_z"]).to_numpy().astype(np.float32),
"pre_E": df["pre_E"].to_numpy().astype(np.float32),
"pre_dir": df.select(["pre_dx", "pre_dy", "pre_dz"]).to_numpy().astype(np.float32),
"material": df["material"].to_numpy().astype(object),
"layer_id": df["layer_id"].to_numpy().astype(np.int32),
"n_sec": df["child_track_ids"].list.len().to_numpy().astype(np.int32),
"e_sec": df["e_sec"].to_numpy().astype(np.float32),
}
@@ -208,7 +198,7 @@ def iter_cond_chunks(path: str | Path, offset: int = 0) -> Iterator[dict[str, np
"""Yield conditioning-only row-groups (no post-step columns read from disk)."""
pf = pq.ParquetFile(path)
for i in range(pf.num_row_groups):
yield _cond_df_to_dict(pf.read_row_group(i, columns=_COND_COLS).to_pandas(), offset=offset)
yield _cond_df_to_dict(pl.DataFrame(pf.read_row_group(i, columns=_COND_COLS)), offset=offset)
def build_index_maps(
@@ -225,42 +215,28 @@ def build_index_maps(
def build_index_maps_from_files(
files: list[Path],
) -> tuple[dict[int, int], dict[str, int]]:
"""Scan only pdg and material columns across all files (2-column read)."""
pdg_vals: set[int] = set()
mat_vals: set[str] = set()
for path in files:
df = pd.read_parquet(path, columns=["pdg", "material"])
pdg_vals.update(int(v) for v in df["pdg"].unique())
mat_vals.update(str(v) for v in df["material"].unique())
"""Scan only pdg and material columns across all files (fused single-pass scan)."""
from giant.data.scan import ScanRequest, scan_metadata
result = scan_metadata(files, ScanRequest(pdg=True, material=True))
assert result.pdg is not None and result.material is not None
return (
{v: i for i, v in enumerate(sorted(pdg_vals))},
{v: i for i, v in enumerate(sorted(mat_vals))},
{v: i for i, v in enumerate(sorted(result.pdg))},
{v: i for i, v in enumerate(sorted(result.material))},
)
def _accumulate_value_counts(counts: dict, series: pd.Series, cast) -> None:
for name, count in series.value_counts().items():
name = cast(name)
counts[name] = counts.get(name, 0) + int(count)
def _rank_by_frequency_from_files(files: list[Path], column: str, cast) -> dict:
"""Scan `column` across `files` and return `{cast(value): total_count}`,
accumulated in file order (see `fingerprint_files`'s docstring on why
scan order — not a normalized/sorted order — is preserved: it drives
tie-breaking in the frequency ranking below)."""
counts: dict = {}
for path in files:
df = pd.read_parquet(path, columns=[column])
_accumulate_value_counts(counts, df[column], cast)
return counts
def _topn_plus_other_map(counts: dict, n_classes: int) -> tuple[dict, dict, dict]:
def _topn_plus_other_map(counts: "Mapping[Any, ValueStat]", n_classes: int) -> tuple[dict, dict, dict]:
"""Frequency-capped value->index map: the `n_classes - 1` most frequent
keys get their own index; every rarer key is bucketed into a shared
"other" index (`n_classes - 1`).
`counts` maps each key to something with `.count` and `.first_seen`
attributes (`giant.data.scan.ValueStat`) — ties in `.count` are broken by
`.first_seen` (whichever value was scanned first: file order, then row
order within a file — see `giant.data.scan`'s module docstring). This is
an explicit, documented contract, not an accident of iteration order.
Returns `(class_map, other_members, class_counts)` — `other_members` is
`{key: count}` for every key bucketed into "other" (the empirical
within-bucket distribution, for `other_policy = "sample"` at rollout);
@@ -270,15 +246,15 @@ def _topn_plus_other_map(counts: dict, n_classes: int) -> tuple[dict, dict, dict
(gitea #44) needs and that would otherwise be dropped once `counts` is
collapsed into `class_map`.
"""
ranked = sorted(counts, key=lambda k: counts[k], reverse=True)
ranked = sorted(counts, key=lambda k: (-counts[k].count, counts[k].first_seen))
keep = ranked[: max(n_classes - 1, 0)]
class_map = {k: i for i, k in enumerate(keep)}
class_counts = {i: counts[k] for i, k in enumerate(keep)}
class_counts = {i: counts[k].count for i, k in enumerate(keep)}
other_idx = n_classes - 1
other_members: dict = {}
for k in ranked[len(keep) :]:
class_map[k] = other_idx
other_members[k] = counts[k]
other_members[k] = counts[k].count
if other_members:
class_counts[other_idx] = sum(other_members.values())
return class_map, other_members, class_counts
@@ -294,8 +270,11 @@ def build_process_map_from_files(files: list[Path], n_experts: int) -> dict[str,
mirrors how `build_features` clamps the n_sec label to K_MAX for the
fixed-width n_sec_head classifier.
"""
counts = _rank_by_frequency_from_files(files, "process", str)
class_map, _, _ = _topn_plus_other_map(counts, n_experts)
from giant.data.scan import ScanRequest, scan_metadata
result = scan_metadata(files, ScanRequest(process=True))
assert result.process is not None
class_map, _, _ = _topn_plus_other_map(result.process, n_experts)
return class_map
@@ -327,8 +306,13 @@ def build_topn_map_from_files(files: list[Path], column: str, n_classes: int, ca
later for `other_policy = "sample"` at rollout — computed now since it's
free during this same scan.
"""
counts = _rank_by_frequency_from_files(files, column, cast)
class_map, other_members, class_counts = _topn_plus_other_map(counts, n_classes)
from giant.data.scan import ScanRequest, scan_metadata
if column != "material":
raise ValueError(f"build_topn_map_from_files only supports column='material', got {column!r}")
result = scan_metadata(files, ScanRequest(material=True))
assert result.material is not None
class_map, other_members, class_counts = _topn_plus_other_map(result.material, n_classes)
return TopNMap(class_map=class_map, other_members=other_members, class_counts=class_counts)
@@ -349,16 +333,9 @@ def build_pdg_topn_map_from_files(files: list[Path], n_classes: int) -> TopNMap:
join (see `_df_to_dict`'s `has_sec_lists` guard) — silently skipped for
those, same convention as elsewhere in this module.
"""
counts: dict = {}
for path in files:
columns = ["pdg"]
has_sec = "sec_pdg_list" in pq.ParquetFile(path).schema_arrow.names
if has_sec:
columns.append("sec_pdg_list")
df = pd.read_parquet(path, columns=columns)
_accumulate_value_counts(counts, df["pdg"], int)
if has_sec:
exploded = df["sec_pdg_list"].explode().dropna()
_accumulate_value_counts(counts, exploded, int)
class_map, other_members, class_counts = _topn_plus_other_map(counts, n_classes)
from giant.data.scan import ScanRequest, scan_metadata
result = scan_metadata(files, ScanRequest(pooled_pdg=True))
assert result.pooled_pdg is not None
class_map, other_members, class_counts = _topn_plus_other_map(result.pooled_pdg, n_classes)
return TopNMap(class_map=class_map, other_members=other_members, class_counts=class_counts)
+171
View File
@@ -0,0 +1,171 @@
"""Fused metadata scan over one or more parquet files.
`giant.pipeline.run_setup_stage` needs several distinct frequency summaries
before training can start — the event-id → row-count index (for the train/val
split), the pdg/material vocabularies, an optional physics-process count, and
a pooled pdg count (primary + secondary species, for onehot conditioning).
Each of those used to be its own full `pd.read_parquet(path, columns=[...])`
per file (`giant.data.loader`'s old `_rank_by_frequency_from_files` /
`build_index_maps_from_files` / `build_pdg_topn_map_from_files`) — up to five
separate reads of the same file. `scan_metadata` answers all of them in one
`pl.collect_all` per file instead, sharing the file open/decompress cost.
Every requested count comes back keyed by value, as a `ValueStat(count,
first_seen)`. `first_seen` is the value's row ordinal — file order (as given
in `files`), then row order within a file — via `row_index_name` on the
per-file lazy scan plus a running row offset across files. This is what
`giant.data.loader._topn_plus_other_map`'s frequency-ranking tie-break keys
on: among equally-frequent values, whichever was scanned first wins its own
class slot. That is an explicit, documented contract (this module is where
it's implemented), not an accident of iteration order.
"""
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
import numpy as np
import polars as pl
from giant.data.loader import _offset_event_id, event_id_offset
@dataclass(frozen=True)
class ScanRequest:
"""Which aggregations to compute. Every field defaults off so a caller
only pays for what it actually needs."""
event_index: bool = False
pdg: bool = False
material: bool = False
process: bool = False
pooled_pdg: bool = False
"""pdg exploded sec_pdg_list — both roles a PDG code plays (primary
species and secondary species), pooled into one count per code. See
`giant.data.loader.build_pdg_topn_map_from_files`'s docstring for why."""
@dataclass(frozen=True)
class ValueStat:
count: int
first_seen: int
@dataclass
class MetadataScan:
event_index: tuple[np.ndarray, np.ndarray] | None = None
"""(unique_ids, counts), ids ascending — matches
`setup_cache.compute_event_index_from_files`'s return shape."""
pdg: dict[int, ValueStat] | None = None
material: dict[str, ValueStat] | None = None
process: dict[str, ValueStat] | None = None
pooled_pdg: dict[int, ValueStat] | None = None
def _group_lazy(path: Path, column: str) -> pl.LazyFrame:
return (
pl.scan_parquet(path, row_index_name="__row")
.select(column, "__row")
.group_by(column)
.agg(pl.len().alias("__count"), pl.col("__row").min().alias("__first_row"))
)
def _pooled_pdg_lazy(path: Path, has_sec_pdg_list: bool) -> pl.LazyFrame:
lf = pl.scan_parquet(path, row_index_name="__row")
parts = [lf.select(pl.col("pdg").alias("__val"), "__row")]
if has_sec_pdg_list:
parts.append(lf.select(pl.col("sec_pdg_list").alias("__val"), "__row").explode("__val").drop_nulls("__val"))
combined = pl.concat(parts)
return combined.group_by("__val").agg(pl.len().alias("__count"), pl.col("__row").min().alias("__first_row"))
def _merge_counts(acc: dict, df: pl.DataFrame, column: str, row_offset: int, cast) -> None:
for key, count, first_row in zip(
df[column].to_list(), df["__count"].to_list(), df["__first_row"].to_list(), strict=True
):
key = cast(key)
first_seen = row_offset + int(first_row)
if key in acc:
prev_count, prev_first = acc[key]
acc[key] = (prev_count + int(count), min(prev_first, first_seen))
else:
acc[key] = (int(count), first_seen)
def scan_metadata(files: list[Path], request: ScanRequest) -> MetadataScan:
"""Scan `files` once (one `pl.collect_all` per file) and return every
aggregation `request` asks for. Files with zero rows contribute nothing
but still advance nothing (no row_offset change, nothing to merge)."""
event_id_parts: list[tuple[np.ndarray, np.ndarray]] = []
pdg_acc: dict[int, tuple[int, int]] = {}
material_acc: dict[str, tuple[int, int]] = {}
process_acc: dict[str, tuple[int, int]] = {}
pooled_pdg_acc: dict[int, tuple[int, int]] = {}
row_offset = 0
for file_idx, path in enumerate(files):
keys: list[str] = []
lazies: list[pl.LazyFrame] = []
if request.event_index:
keys.append("event_id")
lazies.append(_group_lazy(path, "event_id"))
if request.pdg:
keys.append("pdg")
lazies.append(_group_lazy(path, "pdg"))
if request.material:
keys.append("material")
lazies.append(_group_lazy(path, "material"))
if request.process:
keys.append("process")
lazies.append(_group_lazy(path, "process"))
if request.pooled_pdg:
has_sec = "sec_pdg_list" in pl.scan_parquet(path).collect_schema().names()
keys.append("pooled_pdg")
lazies.append(_pooled_pdg_lazy(path, has_sec))
keys.append("__n")
lazies.append(pl.scan_parquet(path).select(pl.len().alias("__n")))
results = dict(zip(keys, pl.collect_all(lazies, engine="streaming"), strict=True))
n_rows = int(results["__n"].item()) if len(results["__n"]) else 0
if request.event_index:
df = results["event_id"]
ids = _offset_event_id(df["event_id"].to_numpy(), event_id_offset(file_idx))
counts = df["__count"].to_numpy().astype(np.int64)
if ids.size:
event_id_parts.append((ids, counts))
if request.pdg:
_merge_counts(pdg_acc, results["pdg"], "pdg", row_offset, int)
if request.material:
_merge_counts(material_acc, results["material"], "material", row_offset, str)
if request.process:
_merge_counts(process_acc, results["process"], "process", row_offset, str)
if request.pooled_pdg:
_merge_counts(pooled_pdg_acc, results["pooled_pdg"], "__val", row_offset, int)
row_offset += n_rows
event_index = None
if request.event_index:
if event_id_parts:
all_ids = np.concatenate([p[0] for p in event_id_parts])
all_counts = np.concatenate([p[1] for p in event_id_parts])
order = np.argsort(all_ids, kind="stable")
event_index = (all_ids[order], all_counts[order])
else:
event_index = (np.empty(0, dtype=np.int64), np.empty(0, dtype=np.int64))
def _to_stats(acc: dict) -> dict:
return {k: ValueStat(*v) for k, v in acc.items()}
return MetadataScan(
event_index=event_index,
pdg=_to_stats(pdg_acc) if request.pdg else None,
material=_to_stats(material_acc) if request.material else None,
process=_to_stats(process_acc) if request.process else None,
pooled_pdg=_to_stats(pooled_pdg_acc) if request.pooled_pdg else None,
)
+14 -5
View File
@@ -23,7 +23,7 @@ import numpy as np
from giant import config
from giant.constants import COND_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_SLOT_DIM, X_DIM
from giant.data.loader import TopNMap, event_id_offset, load_event_ids
from giant.data.loader import TopNMap
from giant.data.transforms import Normalizer, sorted_membership
# Bump manually on a change to the data-encoding semantics (e.g. a future
@@ -349,12 +349,21 @@ def save(
def compute_event_index_from_files(files: list[Path]) -> tuple[np.ndarray, np.ndarray]:
"""Unique event ids + per-event row (step) counts, across all `files`."""
"""Unique event ids + per-event row (step) counts, across all `files`.
Computed via a streaming per-file `group_by("event_id")` (see
`giant.data.scan.scan_metadata`) rather than concatenating every row's
raw event_id across every file before `np.unique` — the latter's peak
memory is 8 bytes x total row count; this is bounded by the (much
smaller) unique event count instead.
"""
from giant.data.scan import ScanRequest, scan_metadata
if not files:
return np.empty(0, dtype=np.int64), np.empty(0, dtype=np.int64)
all_ids = np.concatenate([load_event_ids(f, offset=event_id_offset(i)) for i, f in enumerate(files)])
unique_ids, counts = np.unique(all_ids, return_counts=True)
return unique_ids, counts
result = scan_metadata(files, ScanRequest(event_index=True))
assert result.event_index is not None
return result.event_index
def n_train_steps_for_split(unique_ids: np.ndarray, counts: np.ndarray, train_events_arr: np.ndarray) -> int:
+12 -8
View File
@@ -28,7 +28,7 @@ from pathlib import Path
from typing import Any, Iterable
import numpy as np
import pandas as pd
import polars as pl
import pyarrow.parquet as pq
_INSTALL_HINT = "the geometry oracle needs scikit-learn — install it with `uv sync --extra cpu --extra geometry`"
@@ -301,14 +301,18 @@ def _fit_slab_lookup(
edges = np.linspace(z_min, z_max, n_bins + 1)
bin_idx = np.clip(np.searchsorted(edges, z, side="right") - 1, 0, n_bins - 1)
# pandas' groupby(...).size() sorts group keys ascending by default, so
# a tie in `n` for the same bin (equal counts split between two
# material/layer_id combos) resolves to the lexicographically-first
# combo — matched here by sorting on the keys first, then a
# maintain_order-stable sort on `n` so ties keep that key order.
counts = (
pd.DataFrame({"bin": bin_idx, "material": mat, "layer_id": lay})
.groupby(["bin", "material", "layer_id"])
.size()
.to_frame("n")
.reset_index()
.sort_values("n", ascending=False)
.drop_duplicates("bin")
pl.DataFrame({"bin": bin_idx, "material": mat, "layer_id": lay})
.group_by(["bin", "material", "layer_id"])
.agg(pl.len().alias("n"))
.sort(["bin", "material", "layer_id"])
.sort("n", descending=True, maintain_order=True)
.unique(subset="bin", keep="first", maintain_order=True)
)
bin_material = np.full(n_bins, "", dtype=object)
+106 -54
View File
@@ -15,14 +15,12 @@ from giant.constants import (
from giant.data import setup_cache
from giant.data.loader import (
TopNMap,
_topn_plus_other_map,
event_id_offset,
find_parquet_files,
iter_file_chunks,
build_index_maps_from_files,
build_pdg_topn_map_from_files,
build_process_map_from_files,
build_topn_map_from_files,
)
from giant.data.scan import MetadataScan, ScanRequest, scan_metadata
from giant.data.transforms import (
Normalizer,
build_features,
@@ -127,12 +125,71 @@ def run_setup_stage(
loaded = setup_cache.load(data, files, echo=echo)
cache = loaded if loaded is not None else setup_cache.SetupCache.empty(files)
if cache is not None and cache.event_index is not None:
# Every section below first asks the cache; whatever's missing is
# collected into one ScanRequest and answered by a single fused scan
# (giant.data.scan.scan_metadata), instead of a separate full pass per
# section (event index, vocab, process counts, pdg/material top-N counts
# used to each re-open and re-read every file on their own).
particle_cfg = cfg["conditioning"]["particle"]
material_cfg = cfg["conditioning"]["material"]
particle_type_cfg = config.ParticleTypeConfig.from_dict(cfg["stage2_model"].get("particle_type"))
particle_type_target = particle_type_cfg.target
# A process map is needed if either stage's router reads the physics
# process label (type="process"). Only one map is built even if both
# stages want one — see the module-level note in giant/cli.py's
# _router_total_experts for why composed-router n_experts isn't a plain
# int; process routers are never composed in practice, so this doesn't
# need that generality.
process_router_cfg = next(
(r for r in (stage1_router, stage2_router) if r.get("enabled") and r.get("type") == "process"),
None,
)
process_n_experts = process_router_cfg["n_experts"] if process_router_cfg is not None else None
need_pdg_onehot = particle_cfg["type"] == "onehot"
need_sec_type_onehot = particle_type_target == "onehot"
sec_type_n_classes = (
resolve_type_n_classes(particle_type_cfg, particle_cfg["emb_dim"]) if need_sec_type_onehot else None
)
need_material_onehot = material_cfg["type"] == "onehot"
material_n_classes = material_cfg["emb_dim"] if need_material_onehot else None
def _topn_cached(axis: str, n_classes: int) -> TopNMap | None:
return cache.topn_maps.get(setup_cache.topn_key(axis, n_classes)) if cache is not None else None
need_event_index = cache is None or cache.event_index is None
need_vocab = cache is None or cache.vocab is None
need_process = process_n_experts is not None and (cache is None or cache.proc_maps.get(process_n_experts) is None)
# The PDG axis is used independently by conditioning.particle.type="onehot"
# (cond_cat's onehot feature) and stage2_model.particle_type.target="onehot"
# (secondary-species decode) — their class counts can now differ (gitea
# #29: stage2_model.particle_type.n_classes, 0 = inherit
# conditioning.particle.emb_dim), but both are built from the same
# pooled pdg-count scan, so a cache miss on either one asks for it.
need_pdg_pooled = (need_pdg_onehot and _topn_cached("pdg", particle_cfg["emb_dim"]) is None) or (
need_sec_type_onehot and sec_type_n_classes is not None and _topn_cached("pdg", sec_type_n_classes) is None
)
need_material_topn = (
need_material_onehot and material_n_classes is not None and _topn_cached("material", material_n_classes) is None
)
request = ScanRequest(
event_index=need_event_index,
pdg=need_vocab,
material=need_vocab or need_material_topn,
process=need_process,
pooled_pdg=need_pdg_pooled,
)
scan = scan_metadata(files, request) if request != ScanRequest() else MetadataScan()
if not need_event_index:
unique_ids, counts = cache.event_index
echo(f"event index: cache hit ({len(unique_ids):,} unique events)")
else:
echo("scanning event IDs …")
unique_ids, counts = setup_cache.compute_event_index_from_files(files)
assert scan.event_index is not None
unique_ids, counts = scan.event_index
if cache is not None:
cache.event_index = (unique_ids, counts)
@@ -141,55 +198,34 @@ def run_setup_stage(
n_train_steps = setup_cache.n_train_steps_for_split(unique_ids, counts, events_arr)
echo(f" {int(counts.sum()):,} steps | {len(train_events)} train events | {len(val_events)} val events")
if cache is not None and cache.vocab is not None:
if not need_vocab:
pdg_map, mat_map = cache.vocab
echo(f"vocabulary maps: cache hit ({len(pdg_map)} PDG codes, {len(mat_map)} materials)")
else:
echo("building vocabulary maps …")
pdg_map, mat_map = build_index_maps_from_files(files)
assert scan.pdg is not None and scan.material is not None
pdg_map = {v: i for i, v in enumerate(sorted(scan.pdg))}
mat_map = {v: i for i, v in enumerate(sorted(scan.material))}
echo(f" {len(pdg_map)} PDG codes | {len(mat_map)} materials")
if cache is not None:
cache.vocab = (pdg_map, mat_map)
# A process map is needed if either stage's router reads the physics
# process label (type="process"). Only one map is built even if both
# stages want one — see the module-level note in giant/cli.py's
# _router_total_experts for why composed-router n_experts isn't a plain
# int; process routers are never composed in practice, so this doesn't
# need that generality.
proc_map: dict[str, int] | None = None
process_router_cfg = next(
(r for r in (stage1_router, stage2_router) if r.get("enabled") and r.get("type") == "process"),
None,
)
if process_router_cfg is not None:
n_experts = process_router_cfg["n_experts"]
cached_proc_map = cache.proc_maps.get(n_experts) if cache is not None else None
if cached_proc_map is not None:
assert process_n_experts is not None
if not need_process:
assert cache is not None
cached_proc_map = cache.proc_maps.get(process_n_experts)
assert cached_proc_map is not None
proc_map = cached_proc_map
echo(f"process vocabulary: cache hit ({len(proc_map)} labels, {n_experts} experts)")
echo(f"process vocabulary: cache hit ({len(proc_map)} labels, {process_n_experts} experts)")
else:
echo("building process vocabulary …")
proc_map = build_process_map_from_files(files, n_experts=n_experts)
echo(f" {len(proc_map)} process labels mapped to {n_experts} experts")
assert scan.process is not None
proc_map, _, _ = _topn_plus_other_map(scan.process, process_n_experts)
echo(f" {len(proc_map)} process labels mapped to {process_n_experts} experts")
if cache is not None:
cache.proc_maps[n_experts] = proc_map
# Top-N-plus-other maps for onehot conditioning/type axes.
# The PDG axis is used independently by conditioning.particle.type="onehot"
# (cond_cat's onehot feature) and stage2_model.particle_type.target="onehot"
# (secondary-species decode) — their class counts can now differ (gitea
# #29: stage2_model.particle_type.n_classes, 0 = inherit
# conditioning.particle.emb_dim), so each is resolved and built
# independently via _pdg_topn below. cache.topn_maps is keyed by
# (axis, n_classes) (setup_cache.topn_key), so when the two resolve to
# the same N the second call is a cache hit against the first — no extra
# scan in the common case where they still match. The material axis is
# independent of both.
particle_cfg = cfg["conditioning"]["particle"]
material_cfg = cfg["conditioning"]["material"]
particle_type_cfg = config.ParticleTypeConfig.from_dict(cfg["stage2_model"].get("particle_type"))
particle_type_target = particle_type_cfg.target
cache.proc_maps[process_n_experts] = proc_map
def _pdg_topn(n_classes: int) -> TopNMap:
cache_key = setup_cache.topn_key("pdg", n_classes)
@@ -198,33 +234,36 @@ def run_setup_stage(
echo(f"pdg top-N map: cache hit ({len(cached.class_map)} codes, {n_classes} classes)")
return cached
echo("building pdg top-N map …")
topn_map = build_pdg_topn_map_from_files(files, n_classes=n_classes)
assert scan.pooled_pdg is not None
class_map, other_members, class_counts = _topn_plus_other_map(scan.pooled_pdg, n_classes)
topn_map = TopNMap(class_map=class_map, other_members=other_members, class_counts=class_counts)
echo(f" {len(topn_map.class_map)} pdg codes mapped to {n_classes} classes")
if cache is not None:
cache.topn_maps[cache_key] = topn_map
return topn_map
pdg_topn_map: TopNMap | None = None
if particle_cfg["type"] == "onehot":
pdg_topn_map = _pdg_topn(particle_cfg["emb_dim"])
pdg_topn_map: TopNMap | None = _pdg_topn(particle_cfg["emb_dim"]) if need_pdg_onehot else None
sec_type_topn_map: TopNMap | None = None
if particle_type_target == "onehot":
sec_type_n_classes = resolve_type_n_classes(particle_type_cfg, particle_cfg["emb_dim"])
if need_sec_type_onehot:
assert sec_type_n_classes is not None
sec_type_topn_map = _pdg_topn(sec_type_n_classes)
mat_topn_map: TopNMap | None = None
if material_cfg["type"] == "onehot":
n_classes = material_cfg["emb_dim"]
cache_key = setup_cache.topn_key("material", n_classes)
if need_material_onehot:
assert material_n_classes is not None
cache_key = setup_cache.topn_key("material", material_n_classes)
cached = cache.topn_maps.get(cache_key) if cache is not None else None
if cached is not None:
mat_topn_map = cached
echo(f"material top-N map: cache hit ({len(mat_topn_map.class_map)} materials, {n_classes} classes)")
echo(
f"material top-N map: cache hit ({len(mat_topn_map.class_map)} materials, {material_n_classes} classes)"
)
else:
echo("building material top-N map …")
mat_topn_map = build_topn_map_from_files(files, "material", n_classes=n_classes, cast=str)
echo(f" {len(mat_topn_map.class_map)} materials mapped to {n_classes} classes")
assert scan.material is not None
class_map, other_members, class_counts = _topn_plus_other_map(scan.material, material_n_classes)
mat_topn_map = TopNMap(class_map=class_map, other_members=other_members, class_counts=class_counts)
echo(f" {len(mat_topn_map.class_map)} materials mapped to {material_n_classes} classes")
if cache is not None:
cache.topn_maps[cache_key] = mat_topn_map
@@ -456,17 +495,30 @@ def run_train_job(
)
pin = device.type == "cuda"
# DataLoader worker subprocesses default to fork() on Linux, but by the
# time they're created this process has already run polars queries
# (run_setup_stage's fused metadata scan, above) — polars' native
# (rayon) thread pool doesn't survive a fork: a worker that inherits it
# mid-fork deadlocks the instant it touches polars itself, which
# StreamingStepsDataset's iter_file_chunks now does on every row group.
# "spawn" starts each worker as a fresh interpreter with no inherited
# thread-pool state, avoiding that hazard entirely. Only matters when
# workers actually exist — num_workers=0 runs the dataset in-process and
# never forks.
mp_context = "spawn" if num_workers > 0 else None
train_loader = DataLoader(
train_ds,
batch_size=None,
num_workers=num_workers,
pin_memory=pin,
multiprocessing_context=mp_context,
)
val_loader = DataLoader(
val_ds,
batch_size=None,
num_workers=num_workers,
pin_memory=pin,
multiprocessing_context=mp_context,
)
model_config = {
+162
View File
@@ -0,0 +1,162 @@
"""Benchmark `giant.pipeline.run_setup_stage`'s cold-cache scan against synthetic data.
Generates a schema-complete synthetic steps parquet (matching
`tests/test_pipeline.py`'s `_make_synthetic_steps`, but built with vectorized
numpy instead of a per-row Python loop so it scales to millions of rows) at a
few row counts, times `run_setup_stage` with `cache_setup=False` (so every
call is a genuine cold scan, never served from the sidecar), and prints a
before/after-style table. Run this on `master` before a change and again
after to see what a step actually bought see the "speed up dwarf
warm-cache" plan for the pass-by-pass breakdown this benchmark is meant to
attribute (giant/data/loader.py, giant/data/scan.py, giant/pipeline.py).
Usage: ``uv run python giant/tools/profile_setup_scan.py``
"""
from __future__ import annotations
import time
from pathlib import Path
from tempfile import TemporaryDirectory
import numpy as np
import polars as pl
from giant import config as gconfig
from giant.pipeline import run_setup_stage
ROW_COUNTS = [20_000, 100_000, 500_000, 2_000_000]
_MATERIALS = ["G4_AIR", "G4_Fe"]
_PDGS = [11, 22]
_PROCESSES = ["eIoni", "phot", "compt"]
def _unit_vectors(n: int, rng: np.random.Generator) -> np.ndarray:
v = rng.normal(size=(n, 3))
return v / np.linalg.norm(v, axis=1, keepdims=True)
def _ragged_lists(k: np.ndarray, rng: np.random.Generator, lo: float, hi: float) -> list[list[float]]:
total = int(k.sum())
flat = rng.uniform(lo, hi, size=total)
idx = np.cumsum(k)[:-1]
return [arr.tolist() for arr in np.split(flat, idx)]
def _make_synthetic_steps(n: int, seed: int = 0) -> pl.DataFrame:
"""Vectorized equivalent of tests/test_pipeline.py's `_make_synthetic_steps`.
event_id is assigned so each event gets 2-3 steps (matching that
fixture's structure), and pdg/material/process cycle deterministically
by row index rather than being drawn at random, same as the original.
"""
rng = np.random.default_rng(seed)
n_events = max(n // 3, 1)
pre_E = rng.uniform(50.0, 500.0, size=n)
n_sec = rng.integers(0, 3, size=n)
frac_dep = rng.uniform(0.05, 0.3, size=n)
frac_sec = np.where(n_sec > 0, rng.uniform(0.05, 0.2, size=n), 0.0)
frac_post = 1.0 - frac_dep - frac_sec
edep = pre_E * frac_dep
e_sec = pre_E * frac_sec
post_E = pre_E * frac_post
pre_pos = rng.uniform(-10, 10, size=(n, 3))
step_length = rng.uniform(0.1, 5.0, size=n)
pre_dir = np.zeros((n, 3))
pre_dir[:, 2] = 1.0
post_dir = _unit_vectors(n, rng)
post_pos = pre_pos + step_length[:, None] * pre_dir
row_idx = np.arange(n)
event_id = row_idx % n_events
sec_E = _ragged_lists(n_sec, rng, 0.1, 1.0) # placeholder magnitude, rescaled below
sec_dx = _ragged_lists(n_sec, rng, -1.0, 1.0)
sec_dy = _ragged_lists(n_sec, rng, -1.0, 1.0)
sec_dz = _ragged_lists(n_sec, rng, -1.0, 1.0)
total_sec = int(n_sec.sum())
flat_pdg = [_PDGS[(row_idx[i] + j) % 2] for i in range(n) for j in range(n_sec[i])]
idx = np.cumsum(n_sec)[:-1]
sec_pdg = (
[list(x) for x in np.split(np.array(flat_pdg, dtype=np.int64), idx)] if total_sec else [[] for _ in range(n)]
)
# Rescale each row's secondary energies to sum to that row's e_sec (a
# Dirichlet split, like the original fixture) rather than the raw
# uniform placeholder.
sec_E_scaled = []
for i in range(n):
vals = np.array(sec_E[i])
if vals.size:
sec_E_scaled.append((vals / vals.sum() * e_sec[i]).tolist())
else:
sec_E_scaled.append([])
return pl.DataFrame(
{
"event_id": event_id,
"pdg": np.array(_PDGS)[row_idx % 2],
"pre_x": pre_pos[:, 0],
"pre_y": pre_pos[:, 1],
"pre_z": pre_pos[:, 2],
"pre_E": pre_E,
"pre_dx": pre_dir[:, 0],
"pre_dy": pre_dir[:, 1],
"pre_dz": pre_dir[:, 2],
"material": np.array(_MATERIALS)[row_idx % 2],
"layer_id": row_idx % 5,
"child_track_ids": [list(range(int(k))) for k in n_sec],
"e_sec": e_sec,
"process": np.array(_PROCESSES)[row_idx % 3],
"step_length": step_length,
"post_E": post_E,
"edep": edep,
"post_dx": post_dir[:, 0],
"post_dy": post_dir[:, 1],
"post_dz": post_dir[:, 2],
"post_x": post_pos[:, 0],
"post_y": post_pos[:, 1],
"post_z": post_pos[:, 2],
"sec_E_list": sec_E_scaled,
"sec_pdg_list": sec_pdg,
"sec_dx_list": sec_dx,
"sec_dy_list": sec_dy,
"sec_dz_list": sec_dz,
}
)
def _time_setup_stage(data: Path) -> float:
cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, None, {})
gconfig.validate_config(cfg)
t0 = time.perf_counter()
run_setup_stage(
data,
val_fraction=cfg["train"]["val_fraction"],
seed=cfg["train"]["seed"],
cfg=cfg,
cache_setup=False,
echo=lambda *a, **k: None,
)
return time.perf_counter() - t0
def main() -> None:
with TemporaryDirectory(prefix="giant-setup-scan-profile-") as tmp:
tmp_path = Path(tmp)
print(f"{'n_rows':>10s} {'time (s)':>10s} {'rows/s':>12s}")
for n in ROW_COUNTS:
path = tmp_path / f"steps_{n}.parquet"
_make_synthetic_steps(n).write_parquet(path)
# warm the OS page cache so the timed pass measures compute, not
# the one-time cold read of a freshly-written file.
pl.scan_parquet(path).select(pl.len()).collect()
dt = _time_setup_stage(path)
print(f"{n:>10,d} {dt:>10.3f} {n / dt:>12,.0f}")
path.unlink()
if __name__ == "__main__":
main()
+5 -1
View File
@@ -6,7 +6,7 @@ readme = "README.md"
requires-python = ">=3.12"
dependencies = [
"numpy>=1.26,<3",
"pandas>=2.2,<4",
"polars>=1.0,<2",
"pyarrow>=16,<25",
"tqdm>=4.60,<5",
"typer>=0.12,<1",
@@ -28,6 +28,10 @@ dev = [
"ty>=0.0.50,<0.1",
"bump-my-version>=1.2,<2",
"git-cliff>=2,<3",
# Only used by test fixtures (writing small parquet files) — not a
# runtime dependency of giant itself since the pandas -> polars
# data-loading rewrite.
"pandas>=2.2,<4",
"giant[convert,analysis,geometry,wandb]",
]
geometry = [
+23 -6
View File
@@ -108,12 +108,12 @@ def test_build_process_map_from_files_spans_multiple_files(tmp_path):
def test_build_process_map_from_files_tie_breaking_pins_first_seen_order(tmp_path):
"""When two processes end up with equal total counts, ranking falls back
to whichever was accumulated first (`sorted(..., reverse=True)` is stable,
and `counts` is built in file/row-scan order) this is implementation-
defined, not a documented contract, so pin it explicitly: a future
rewrite (e.g. a polars-based single-scan) that ties differently would
silently reshuffle which processes get their own expert slot across a
retrain, and this test is what should catch that."""
to whichever was scanned first file order, then row order within a
file (`giant.data.scan`'s `first_seen` ordinal, ranked by
`giant.data.loader._topn_plus_other_map`'s `(-count, first_seen)` key).
This is an explicit, documented contract (not an accident of iteration
order), pinned here so a future change to the ranking can't silently
reshuffle which processes get their own expert slot across a retrain."""
path = tmp_path / "a.parquet"
pd.DataFrame({"process": ["compt", "phot", "compt", "phot"]}).to_parquet(path)
@@ -233,6 +233,23 @@ def test_build_pdg_topn_map_from_files_pools_primary_and_secondary_pdg(tmp_path)
assert m.class_counts == {0: 11, 1: 5}
def test_build_pdg_topn_map_from_files_pooled_tie_breaks_by_row_position(tmp_path):
"""Pooled pdg counting merges the primary `pdg` column and the exploded
`sec_pdg_list` column via one `group_by` over both (see
`giant.data.scan._pooled_pdg_lazy`), keyed by row position regardless of
which role (primary or secondary) a code was seen in not "all
primaries before all secondaries" the way a two-pass accumulation would.
11 (primary, row 0), 33 (primary, row 1) and 22 (secondary, row 1) all
end up with count 1; 11's strictly earlier row wins the tie over both,
whatever order 33/22 (tied with each other, same row) land in."""
path = tmp_path / "a.parquet"
pd.DataFrame({"pdg": [11, 33], "sec_pdg_list": [[], [22]]}).to_parquet(path)
m = build_pdg_topn_map_from_files([path], n_classes=4)
assert m.class_map[11] == 0
def test_build_pdg_topn_map_from_files_missing_sec_pdg_list_column(tmp_path):
"""Files predating the parent->child join have no sec_pdg_list column —
must not raise, just count the primary pdg column alone."""
+2 -2
View File
@@ -142,7 +142,7 @@ def test_run_train_job_second_run_hits_cache(tmp_path, data, monkeypatch):
def _forbidden(*a, **k):
raise AssertionError("should be served from cache, not recomputed")
monkeypatch.setattr("giant.pipeline.build_index_maps_from_files", _forbidden)
monkeypatch.setattr("giant.pipeline.scan_metadata", _forbidden)
monkeypatch.setattr("giant.pipeline.iter_file_chunks", _forbidden)
echo2 = _run(data, tmp_path / "out2")
@@ -283,7 +283,7 @@ def test_run_train_job_new_val_fraction_is_partial_hit(tmp_path, data, monkeypat
def _forbidden(*a, **k):
raise AssertionError("vocab should be served from cache")
monkeypatch.setattr("giant.pipeline.build_index_maps_from_files", _forbidden)
monkeypatch.setattr("giant.pipeline.scan_metadata", _forbidden)
echo2 = _run(data, tmp_path / "out2", cfg=_tiny_cfg(val_fraction=0.3))
joined = "\n".join(echo2)
Generated
+4 -2
View File
@@ -679,8 +679,8 @@ version = "0.3.16"
source = { editable = "." }
dependencies = [
{ name = "numpy" },
{ name = "pandas" },
{ name = "particle" },
{ name = "polars" },
{ name = "pyarrow" },
{ name = "pyyaml" },
{ name = "tqdm" },
@@ -712,6 +712,7 @@ dev = [
{ name = "git-cliff" },
{ name = "ipykernel" },
{ name = "matplotlib" },
{ name = "pandas" },
{ name = "plotstyle" },
{ name = "polars" },
{ name = "pytest" },
@@ -738,9 +739,10 @@ requires-dist = [
{ name = "ipykernel", marker = "extra == 'analysis'", specifier = ">=7.3.0" },
{ name = "matplotlib", marker = "extra == 'analysis'", specifier = ">=3.8,<4" },
{ name = "numpy", specifier = ">=1.26,<3" },
{ name = "pandas", specifier = ">=2.2,<4" },
{ name = "pandas", marker = "extra == 'dev'", specifier = ">=2.2,<4" },
{ name = "particle", specifier = ">=1.0,<2" },
{ name = "plotstyle", marker = "extra == 'analysis'", specifier = ">=1.0.0", index = "https://git.larsbogner.de/api/packages/lars/pypi/simple/" },
{ name = "polars", specifier = ">=1.0,<2" },
{ name = "polars", marker = "extra == 'analysis'", specifier = ">=1.0,<2" },
{ name = "polars", marker = "extra == 'convert'", specifier = ">=1.0,<2" },
{ name = "pyarrow", specifier = ">=16,<25" },