lars
a5683517be
Brings in the rollout-validation fixes developed alongside Phase 2 (exact e_sec budget rescaling in decode_secondaries, filtering synthetic termination rows out of load_rollout_vs_truth, Tier 4 truth overlay, --energy-gev support in dwarf make-root) and reconciles them with this branch's mixture-of-experts routing work: build_features/ build_models/dataset plumbing keep the ProcessRouter's proc_map/ proc_idx threading, and create_root_files.py's job_seed folds in both the per-job seed derivation and the new energy_gev component.
595 lines
24 KiB
Python
595 lines
24 KiB
Python
import warnings
|
||
|
||
import numpy as np
|
||
|
||
_EPS = 1e-8
|
||
|
||
# Floor added to each energy fraction before taking log-ratios so the simplex
|
||
# coordinates stay finite on the heavily populated boundary: e_sec is 0 in ~half
|
||
# of all steps and post_E is 0 at every track end. Kept tiny (1e-5 of pre_E) so
|
||
# the conservation it slightly softens is physically negligible (~0.001%).
|
||
_SIMPLEX_FLOOR = 1e-5
|
||
|
||
|
||
def log_transform(x: np.ndarray, eps: float = _EPS) -> np.ndarray:
|
||
return np.log(np.asarray(x, dtype=np.float32) + eps)
|
||
|
||
|
||
def inv_log_transform(y: np.ndarray, eps: float = _EPS) -> np.ndarray:
|
||
return np.exp(np.asarray(y, dtype=np.float32)) - eps
|
||
|
||
|
||
def energy_simplex_encode(
|
||
edep: np.ndarray,
|
||
e_sec: np.ndarray,
|
||
post_E: np.ndarray,
|
||
pre_E: np.ndarray,
|
||
) -> np.ndarray:
|
||
"""Encode (edep, e_sec, post_E) as 2 additive-log-ratio (ALR) coordinates.
|
||
|
||
The three energies are first expressed as fractions of pre_E that sum to 1:
|
||
post_E is kept exact (so the particle's retained energy — the next step's
|
||
input during a rollout — is preserved), and the energy actually lost,
|
||
`delta_e = pre_E - post_E`, is split between local deposit and secondaries
|
||
in the recorded edep:e_sec ratio (re-attributing any sub-threshold / rest-mass
|
||
leakage proportionally so the three fractions sum to exactly 1). A small floor
|
||
(`_SIMPLEX_FLOOR`) keeps the log-ratios finite where a fraction is 0.
|
||
|
||
Returns an (N, 2) array of ALR coordinates referenced to the post fraction;
|
||
`energy_simplex_decode` is its inverse (up to the floor softening).
|
||
"""
|
||
pre_E = np.maximum(np.asarray(pre_E, dtype=np.float32), _EPS)
|
||
raw_post_E = np.asarray(post_E, dtype=np.float32)
|
||
post_E = np.clip(raw_post_E, 0.0, pre_E)
|
||
delta_e = pre_E - post_E
|
||
edep = np.asarray(edep, dtype=np.float32)
|
||
e_sec = np.asarray(e_sec, dtype=np.float32)
|
||
lost = edep + e_sec
|
||
has_loss = lost > _EPS
|
||
# Clipping post_E down to pre_E forces delta_e (and thus the rescaled
|
||
# edep/e_sec below) to 0 even on rows where edep/e_sec were genuinely
|
||
# recorded as nonzero — warn so this doesn't silently discard real data.
|
||
discarded = (raw_post_E > pre_E) & has_loss
|
||
if np.any(discarded):
|
||
n = int(np.sum(discarded))
|
||
warnings.warn(
|
||
f"energy_simplex_encode: {n}/{len(discarded)} step(s) had "
|
||
"post_E > pre_E (clipped) while recording nonzero edep/e_sec; "
|
||
"that recorded energy deposit is discarded to keep delta_e "
|
||
"consistent with the clip.",
|
||
stacklevel=2,
|
||
)
|
||
scale = np.where(has_loss, delta_e / np.maximum(lost, _EPS), 0.0)
|
||
# Where nothing was recorded as deposited/secondary but energy was lost,
|
||
# attribute all of delta_e to local deposit.
|
||
edep_r = np.where(has_loss, edep * scale, delta_e)
|
||
e_sec_r = np.where(has_loss, e_sec * scale, 0.0)
|
||
|
||
f = np.stack([edep_r, e_sec_r, post_E], axis=-1) / pre_E[:, None] # (N,3), sums≈1
|
||
f = (f + _SIMPLEX_FLOOR) / (1.0 + 3.0 * _SIMPLEX_FLOOR)
|
||
log_f = np.log(f)
|
||
z = log_f[:, :2] - log_f[:, 2:3] # ALR vs post reference
|
||
return z.astype(np.float32)
|
||
|
||
|
||
def energy_simplex_decode(
|
||
z: np.ndarray, pre_E: np.ndarray
|
||
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
|
||
"""Inverse of `energy_simplex_encode`: ALR coords + pre_E → physical energies.
|
||
|
||
A softmax over `[z_edep, z_sec, 0]` recovers the three simplex fractions, so
|
||
`edep + e_sec + post_E == pre_E` holds by construction (the energy-conservation
|
||
inductive bias). Returns `(edep, e_sec, post_E, delta_e)` in physical units.
|
||
"""
|
||
z = np.asarray(z, dtype=np.float32)
|
||
pre_E = np.asarray(pre_E, dtype=np.float32)
|
||
logits = np.concatenate([z, np.zeros((len(z), 1), dtype=np.float32)], axis=1)
|
||
logits = logits - logits.max(axis=1, keepdims=True)
|
||
f = np.exp(logits)
|
||
f /= f.sum(axis=1, keepdims=True) # sums to 1 exactly → exact conservation
|
||
E = f * pre_E[:, None]
|
||
edep, e_sec, post_E = E[:, 0], E[:, 1], E[:, 2]
|
||
delta_e = pre_E - post_E
|
||
return (
|
||
edep.astype(np.float32),
|
||
e_sec.astype(np.float32),
|
||
post_E.astype(np.float32),
|
||
delta_e.astype(np.float32),
|
||
)
|
||
|
||
|
||
def _rodrigues_axis(pre_dir: np.ndarray) -> np.ndarray:
|
||
"""Unit rotation axis `pre_dir × ẑ`, closed-form since ẑ = [0, 0, 1] is constant.
|
||
|
||
`cross(a, [0,0,1]) = [a_y, -a_x, 0]` — substituting the constant operand
|
||
avoids a generic `np.cross` call (shape/broadcast handling for an
|
||
arbitrary second operand) on every row; profiling on a 114M-row file
|
||
showed `np.cross` as the single hottest call inside this rotation.
|
||
"""
|
||
axis = np.stack(
|
||
[pre_dir[:, 1], -pre_dir[:, 0], np.zeros_like(pre_dir[:, 0])], axis=1
|
||
)
|
||
axis_norm = np.linalg.norm(axis, axis=1, keepdims=True) # (N,1)
|
||
# Replace zero-norm axes with x̂ (the Rodrigues terms that involve the axis
|
||
# are multiplied by sin_t≈0 and (1-cos_t)≈0, so the choice is irrelevant).
|
||
safe_norm = np.where(axis_norm < 1e-7, 1.0, axis_norm)
|
||
return np.where(axis_norm < 1e-7, np.array([[1.0, 0.0, 0.0]]), axis / safe_norm)
|
||
|
||
|
||
def _cross_with_z_axis(axis: np.ndarray, v: np.ndarray) -> np.ndarray:
|
||
"""`axis × v`, closed-form since `axis` from `_rodrigues_axis` always has z = 0.
|
||
|
||
`cross([ax,ay,0], [bx,by,bz]) = [ay*bz, -ax*bz, ax*by - ay*bx]`.
|
||
"""
|
||
ax, ay = axis[:, 0:1], axis[:, 1:2]
|
||
bx, by, bz = v[:, 0:1], v[:, 1:2], v[:, 2:3]
|
||
return np.concatenate([ay * bz, -ax * bz, ax * by - ay * bx], axis=1)
|
||
|
||
|
||
def _validate_unit_pre_dir(pre_dir: np.ndarray) -> np.ndarray:
|
||
"""Normalize pre_dir and raise if any row is too degenerate to define a frame.
|
||
|
||
`local_frame_rotation`/`inv_local_frame_rotation` treat pre_dir[:, 2] as
|
||
cos(angle to ẑ), which is only correct for a unit vector. Small float32
|
||
drift is corrected silently; a near-zero-norm row has no well-defined
|
||
direction, so it's raised loudly instead of producing a meaningless
|
||
rotation (previously it fell through to an arbitrary axis with no error).
|
||
"""
|
||
pre_dir = np.asarray(pre_dir, dtype=np.float32)
|
||
norm = np.linalg.norm(pre_dir, axis=1, keepdims=True)
|
||
if np.any(norm < 1e-6):
|
||
raise ValueError(
|
||
f"pre_dir has {int(np.sum(norm < 1e-6))} row(s) with near-zero norm "
|
||
"(< 1e-6); local/inv_local_frame_rotation require a well-defined "
|
||
"incoming direction for every row."
|
||
)
|
||
return pre_dir / norm
|
||
|
||
|
||
def local_frame_rotation(pre_dir: np.ndarray, post_dir: np.ndarray) -> np.ndarray:
|
||
"""Rotate post_dir into the local frame where pre_dir maps to ẑ (Rodrigues).
|
||
|
||
Preserves the angle between pre_dir and post_dir; the result has post_dir
|
||
expressed relative to a coordinate system in which the incoming particle
|
||
travels along +z.
|
||
"""
|
||
pre_dir = _validate_unit_pre_dir(pre_dir)
|
||
cos_t = np.clip(pre_dir[:, 2:3], -1.0, 1.0) # (N,1); dot with ẑ = z-component
|
||
sin_t = np.sqrt(np.maximum(0.0, 1.0 - cos_t**2)) # (N,1)
|
||
|
||
axis = _rodrigues_axis(pre_dir) # (N,3); zero-z, zero-norm when pre_dir ∥ ẑ
|
||
kxv = _cross_with_z_axis(axis, post_dir) # (N,3)
|
||
kdv = (axis * post_dir).sum(axis=1, keepdims=True) # (N,1)
|
||
|
||
return (post_dir * cos_t + kxv * sin_t + axis * kdv * (1.0 - cos_t)).astype(
|
||
np.float32
|
||
)
|
||
|
||
|
||
class Normalizer:
|
||
def __init__(self) -> None:
|
||
self.mean: np.ndarray | None = None
|
||
self.std: np.ndarray | None = None
|
||
|
||
def fit(self, X: np.ndarray) -> "Normalizer":
|
||
self.mean = X.mean(axis=0).astype(np.float32)
|
||
self.std = X.std(axis=0).astype(np.float32)
|
||
self.std = np.where(self.std < _EPS, 1.0, self.std).astype(np.float32)
|
||
return self
|
||
|
||
def transform(self, X: np.ndarray) -> np.ndarray:
|
||
return ((X - self.mean) / self.std).astype(np.float32)
|
||
|
||
def inverse_transform(self, X: np.ndarray) -> np.ndarray:
|
||
return (X * self.std + self.mean).astype(np.float32)
|
||
|
||
def to_dict(self) -> dict:
|
||
assert self.mean is not None and self.std is not None, "Normalizer not fitted"
|
||
return {"mean": self.mean.tolist(), "std": self.std.tolist()}
|
||
|
||
@classmethod
|
||
def from_dict(cls, d: dict) -> "Normalizer":
|
||
obj = cls()
|
||
obj.mean = np.array(d["mean"], dtype=np.float32)
|
||
obj.std = np.array(d["std"], dtype=np.float32)
|
||
return obj
|
||
|
||
|
||
class _WelfordAccumulator:
|
||
"""Streaming mean/variance (Welford's online algorithm, batch update).
|
||
|
||
Use to fit a Normalizer over data that doesn't fit in memory:
|
||
acc = _WelfordAccumulator(n_features)
|
||
for chunk in data:
|
||
acc.update(chunk)
|
||
normalizer = acc.to_normalizer()
|
||
"""
|
||
|
||
def __init__(self, n_features: int) -> None:
|
||
self.n = 0
|
||
self._mean = np.zeros(n_features, dtype=np.float64)
|
||
self._M2 = np.zeros(n_features, dtype=np.float64)
|
||
|
||
def update(self, X: np.ndarray) -> None:
|
||
X = np.asarray(X, dtype=np.float64)
|
||
B = X.shape[0]
|
||
new_n = self.n + B
|
||
delta = X - self._mean
|
||
self._mean += delta.sum(0) / new_n
|
||
delta2 = X - self._mean
|
||
self._M2 += (delta * delta2).sum(0)
|
||
self.n = new_n
|
||
|
||
def to_normalizer(self) -> "Normalizer":
|
||
norm = Normalizer()
|
||
norm.mean = self._mean.astype(np.float32)
|
||
std = np.sqrt(self._M2 / max(self.n, 1)).astype(np.float32)
|
||
norm.std = np.where(std < _EPS, 1.0, std).astype(np.float32)
|
||
return norm
|
||
|
||
|
||
def travel_direction(pre_pos: np.ndarray, post_pos: np.ndarray) -> np.ndarray:
|
||
"""World-frame unit vector pointing from pre_pos to post_pos.
|
||
|
||
Kept independent of `step_length`: that scalar already encodes the
|
||
magnitude of this displacement, so this function only ever returns
|
||
direction (norm-guarded the same way as `local_frame_rotation`'s axis).
|
||
"""
|
||
disp = post_pos - pre_pos
|
||
norm = np.linalg.norm(disp, axis=1, keepdims=True)
|
||
safe_norm = np.where(norm < 1e-7, 1.0, norm)
|
||
return np.where(norm < 1e-7, np.array([[0.0, 0.0, 1.0]]), disp / safe_norm).astype(
|
||
np.float32
|
||
)
|
||
|
||
|
||
def reconstruct_post_pos(
|
||
pre_pos: np.ndarray,
|
||
pre_dir: np.ndarray,
|
||
step_length: np.ndarray,
|
||
travel_dir_local: np.ndarray,
|
||
) -> np.ndarray:
|
||
"""Inverse of the travel_direction/local_frame_rotation encoding.
|
||
|
||
Single source of truth for combining the magnitude (`step_length`) and
|
||
direction (`travel_dir_local`) back into a world-frame post_pos, so
|
||
`step_length` and post_pos stay consistent by construction.
|
||
"""
|
||
travel_dir_world = inv_local_frame_rotation(pre_dir, travel_dir_local)
|
||
return (pre_pos + step_length.reshape(-1, 1) * travel_dir_world).astype(np.float32)
|
||
|
||
|
||
def inv_local_frame_rotation(
|
||
pre_dir: np.ndarray, post_dir_local: np.ndarray
|
||
) -> np.ndarray:
|
||
"""Inverse of local_frame_rotation: rotate from local frame back to world frame.
|
||
|
||
Applies R^T (same axis, negative angle) to post_dir_local.
|
||
"""
|
||
pre_dir = _validate_unit_pre_dir(pre_dir)
|
||
cos_t = np.clip(pre_dir[:, 2:3], -1.0, 1.0) # dot with ẑ = z-component
|
||
sin_t = np.sqrt(np.maximum(0.0, 1.0 - cos_t**2))
|
||
|
||
axis = _rodrigues_axis(pre_dir)
|
||
kxv = _cross_with_z_axis(axis, post_dir_local)
|
||
kdv = (axis * post_dir_local).sum(axis=1, keepdims=True)
|
||
|
||
# Negative angle: sin_t → -sin_t
|
||
return (post_dir_local * cos_t - kxv * sin_t + axis * kdv * (1.0 - cos_t)).astype(
|
||
np.float32
|
||
)
|
||
|
||
|
||
_STICK_LOGIT_CLIP = 10.0 # logit value used for the last valid secondary slot
|
||
|
||
|
||
def encode_secondaries(
|
||
sec_E_list: np.ndarray,
|
||
sec_dir_list: np.ndarray,
|
||
sec_valid: np.ndarray,
|
||
e_sec: np.ndarray,
|
||
pre_dir: np.ndarray,
|
||
) -> np.ndarray:
|
||
"""Encode per-secondary attributes into continuous per-slot targets.
|
||
|
||
Secondaries must already be sorted descending by energy (as stored in the
|
||
parquet). Returns sec_cont of shape (N, K_MAX, 4):
|
||
slot[i] = [stick_break_logit, local_dir_x, local_dir_y, local_dir_z]
|
||
|
||
Stick-breaking logit: for slot i, f_i = E_i / remaining_budget, where
|
||
remaining_budget = e_sec - sum(E_0..E_{i-1}). The logit is log(f/(1-f)),
|
||
clipped to ±_STICK_LOGIT_CLIP. The last valid slot gets +_STICK_LOGIT_CLIP
|
||
(takes the full remaining budget). Padding slots get 0.
|
||
|
||
sec_pdg_idx (integer) is not processed here — kept separate so the loss
|
||
function can look up the embedding table at training time.
|
||
"""
|
||
N, K = sec_E_list.shape
|
||
e_sec = np.asarray(e_sec, dtype=np.float64)
|
||
|
||
stick_logits = np.zeros((N, K), dtype=np.float32)
|
||
for i in range(K):
|
||
if i == 0:
|
||
remaining = np.maximum(e_sec, _EPS)
|
||
else:
|
||
remaining = np.maximum(e_sec - sec_E_list[:, :i].sum(axis=1), _EPS)
|
||
f = np.clip(sec_E_list[:, i].astype(np.float64) / remaining, _EPS, 1.0 - _EPS)
|
||
logit = np.log(f / (1.0 - f)).astype(np.float32)
|
||
# Last valid slot: give it the full remaining budget
|
||
is_last = sec_valid[:, i] & ~(
|
||
sec_valid[:, i + 1] if i + 1 < K else np.zeros(N, dtype=bool)
|
||
)
|
||
logit = np.where(is_last, _STICK_LOGIT_CLIP, logit)
|
||
logit = np.where(
|
||
sec_valid[:, i], np.clip(logit, -_STICK_LOGIT_CLIP, _STICK_LOGIT_CLIP), 0.0
|
||
)
|
||
stick_logits[:, i] = logit.astype(np.float32)
|
||
|
||
# Rotate each slot's direction into the local frame of the primary.
|
||
# pre_dir is broadcast across all K slots.
|
||
dir_local = np.zeros((N, K, 3), dtype=np.float32)
|
||
for i in range(K):
|
||
# Only rotate valid slots; leave padded slots as (0,0,1) or whatever.
|
||
valid_mask = sec_valid[:, i]
|
||
if valid_mask.any():
|
||
dir_local[valid_mask, i] = local_frame_rotation(
|
||
pre_dir[valid_mask], sec_dir_list[valid_mask, i]
|
||
)
|
||
|
||
sec_cont = np.concatenate(
|
||
[stick_logits[:, :, None], dir_local], axis=-1
|
||
) # (N, K, 4)
|
||
return sec_cont.astype(np.float32)
|
||
|
||
|
||
def decode_secondaries(
|
||
sec_cont: np.ndarray,
|
||
sec_pdg_pred: np.ndarray,
|
||
n_sec: np.ndarray,
|
||
e_sec: np.ndarray,
|
||
pre_dir: np.ndarray,
|
||
pdg_map_inv: dict[int, int],
|
||
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
|
||
"""Inverse of encode_secondaries: continuous targets → physical secondary attrs.
|
||
|
||
sec_cont: (N, K_MAX, 4) — [stick_logit, local_dir_x, local_dir_y, local_dir_z]
|
||
sec_pdg_pred: (N, K_MAX) integer PDG indices (from nearest-neighbor snap)
|
||
n_sec: (N,) integer secondary counts
|
||
e_sec: (N,) total secondary energy budget [MeV]
|
||
pre_dir: (N, 3) pre-step world-frame direction
|
||
pdg_map_inv: maps model index → PDG code
|
||
|
||
Returns (sec_E, sec_dir_world, sec_pdg_code, sec_valid) each shape (N, K_MAX).
|
||
The valid slots' energies (`sec_E[sec_valid]`, per row) always sum to
|
||
exactly `e_sec` — see the rescaling below.
|
||
"""
|
||
N, K, _ = sec_cont.shape
|
||
stick_logits = sec_cont[:, :, 0] # (N, K)
|
||
dir_local = sec_cont[:, :, 1:].copy() # (N, K, 3)
|
||
|
||
# Flow-matching output isn't guaranteed unit norm; normalise before the
|
||
# rotation below, which preserves magnitude rather than fixing it up.
|
||
norms = np.linalg.norm(dir_local, axis=-1, keepdims=True)
|
||
dir_local /= np.where(norms < 1e-8, 1.0, norms)
|
||
|
||
fractions = 1.0 / (1.0 + np.exp(-stick_logits.astype(np.float64)))
|
||
|
||
sec_E = np.zeros((N, K), dtype=np.float64)
|
||
e_sec = np.asarray(e_sec, dtype=np.float64)
|
||
remaining = e_sec.copy()
|
||
for i in range(K):
|
||
sec_E[:, i] = fractions[:, i] * remaining
|
||
remaining = np.maximum(remaining - sec_E[:, i], 0.0)
|
||
|
||
sec_valid = np.arange(K)[None, :] < n_sec[:, None] # (N, K)
|
||
|
||
# Stick-breaking guarantees sum(sec_E[valid]) <= e_sec (each fraction is in
|
||
# [0,1] of an already-shrinking remainder) but rarely hits it exactly, so
|
||
# rescale the valid slots by one common per-row factor to close that gap —
|
||
# rather than dumping the shortfall into whichever slot happens to be last
|
||
# by energy rank, which would let one low-energy secondary balloon and
|
||
# distort the shower's topology. This preserves each row's relative split
|
||
# across its secondaries and only ever scales up (valid_sum <= e_sec).
|
||
# Rows where every valid slot decoded to ~zero (scale undefined) fall back
|
||
# to an even split of e_sec across the n_sec valid slots.
|
||
sec_E = sec_E * sec_valid
|
||
valid_sum = sec_E.sum(axis=1)
|
||
degenerate = (valid_sum <= _EPS) & (n_sec > 0)
|
||
scale = np.where(valid_sum > _EPS, e_sec / np.maximum(valid_sum, _EPS), 0.0)
|
||
sec_E = sec_E * scale[:, None]
|
||
even_share = e_sec / np.maximum(n_sec, 1).astype(np.float64)
|
||
sec_E = np.where(degenerate[:, None] & sec_valid, even_share[:, None], sec_E)
|
||
sec_E = sec_E.astype(np.float32)
|
||
|
||
sec_dir_world = np.zeros((N, K, 3), dtype=np.float32)
|
||
for i in range(K):
|
||
valid = sec_valid[:, i]
|
||
if valid.any():
|
||
sec_dir_world[valid, i] = inv_local_frame_rotation(
|
||
pre_dir[valid], dir_local[valid, i]
|
||
)
|
||
|
||
sec_pdg_code = np.array(
|
||
[
|
||
[pdg_map_inv.get(int(sec_pdg_pred[n, i]), 0) for i in range(K)]
|
||
for n in range(N)
|
||
],
|
||
dtype=np.int32,
|
||
)
|
||
|
||
return sec_E, sec_dir_world, sec_pdg_code, sec_valid
|
||
|
||
|
||
def build_cond_features(
|
||
data: dict[str, np.ndarray],
|
||
pdg_map: dict[int, int],
|
||
mat_map: dict[str, int],
|
||
cond_normalizer: "Normalizer | None" = None,
|
||
) -> tuple[np.ndarray, np.ndarray]:
|
||
"""Build conditioning arrays only — no target, no post-step variables."""
|
||
cond_cont = np.column_stack(
|
||
[
|
||
data["pre_pos"],
|
||
log_transform(data["pre_E"]),
|
||
data["pre_dir"],
|
||
data["layer_id"].astype(np.float32),
|
||
]
|
||
).astype(np.float32)
|
||
|
||
pdg_idx = np.array([pdg_map[int(p)] for p in data["pdg"]], dtype=np.int64)
|
||
mat_idx = np.array([mat_map[str(m)] for m in data["material"]], dtype=np.int64)
|
||
cond_cat = np.column_stack([pdg_idx, mat_idx])
|
||
|
||
if cond_normalizer is not None:
|
||
cond_cont = cond_normalizer.transform(cond_cont)
|
||
|
||
return cond_cont, cond_cat
|
||
|
||
|
||
def build_features(
|
||
data: dict[str, np.ndarray],
|
||
pdg_map: dict[int, int],
|
||
mat_map: dict[str, int],
|
||
cond_normalizer: Normalizer | None = None,
|
||
target_normalizer: Normalizer | None = None,
|
||
fit: bool = False,
|
||
proc_map: dict[str, int] | None = None,
|
||
require_secondaries: bool = False,
|
||
) -> tuple[
|
||
np.ndarray,
|
||
np.ndarray,
|
||
np.ndarray,
|
||
np.ndarray,
|
||
np.ndarray,
|
||
np.ndarray,
|
||
np.ndarray,
|
||
Normalizer | None,
|
||
Normalizer | None,
|
||
]:
|
||
"""Assemble (cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx, proc_idx) arrays.
|
||
|
||
target_s1: (N, 9) Stage-1 primary post-step target (unchanged from Phase 1)
|
||
n_sec: (N,) integer secondary counts (target for n_sec head)
|
||
sec_cont: (N, K_MAX, 4) continuous secondary targets [stick_logit, dir_local]
|
||
sec_pdg_idx: (N, K_MAX) integer PDG model-indices; used to look up embedding
|
||
targets in the training loop
|
||
proc_idx: (N,) integer process-class label (ProcessRouter supervision only —
|
||
never conditioning). Zeros when `proc_map` is None or the loaded
|
||
data has no "process" column (e.g. pre-conversion parquet files).
|
||
|
||
require_secondaries: when True, raise if any step has n_sec > 0 but the
|
||
per-secondary list columns are absent (a mis-converted file that would
|
||
otherwise silently zero all Stage-2 targets). Training paths set this;
|
||
Stage-1-only callers (e.g. `giant predict`) leave it False.
|
||
"""
|
||
from giant.constants import K_MAX
|
||
|
||
post_dir_local = local_frame_rotation(data["pre_dir"], data["post_dir"])
|
||
travel_dir_local = local_frame_rotation(
|
||
data["pre_dir"], travel_direction(data["pre_pos"], data["post_pos"])
|
||
)
|
||
|
||
energy_z = energy_simplex_encode(
|
||
data["edep"], data["e_sec"], data["post_E"], data["pre_E"]
|
||
) # (N, 2)
|
||
|
||
target_s1 = np.column_stack(
|
||
[
|
||
log_transform(data["step_length"]),
|
||
energy_z,
|
||
post_dir_local,
|
||
travel_dir_local,
|
||
]
|
||
).astype(np.float32) # (N, 9)
|
||
|
||
# Phase 2: conditioning drops n_sec and log(e_sec)
|
||
cond_cont = np.column_stack(
|
||
[
|
||
data["pre_pos"],
|
||
log_transform(data["pre_E"]),
|
||
data["pre_dir"],
|
||
data["layer_id"].astype(np.float32),
|
||
]
|
||
).astype(np.float32) # (N, COND_DIM=8)
|
||
|
||
pdg_idx = np.array([pdg_map[int(p)] for p in data["pdg"]], dtype=np.int64)
|
||
mat_idx = np.array([mat_map[str(m)] for m in data["material"]], dtype=np.int64)
|
||
cond_cat = np.column_stack([pdg_idx, mat_idx]) # (N, 2)
|
||
|
||
n_sec_raw = data["n_sec"].astype(
|
||
np.int64
|
||
) # (N,) unclamped, for the valid-slot mask
|
||
# Clamp the classification label to K_MAX: the head only has K_MAX+1 classes
|
||
# (0..K_MAX), and truncating here mirrors the K_MAX-slot truncation already
|
||
# applied to sec_cont/sec_pdg_idx by the loader's list padding. Without this,
|
||
# a rare high-multiplicity step (real data goes up to ~37) hands
|
||
# cross_entropy an out-of-range target and CUDA asserts.
|
||
n_sec = np.minimum(n_sec_raw, K_MAX).astype(np.int64) # (N,)
|
||
|
||
# Secondary continuous targets
|
||
sec_E_list = data.get("sec_E_list")
|
||
sec_dir_list = data.get("sec_dir_list")
|
||
sec_pdg_list = data.get("sec_pdg_list")
|
||
|
||
if sec_E_list is not None and sec_dir_list is not None and sec_pdg_list is not None:
|
||
sec_valid = np.arange(K_MAX)[None, :] < n_sec_raw[:, None] # (N, K_MAX)
|
||
sec_cont = encode_secondaries(
|
||
sec_E_list, sec_dir_list, sec_valid, data["e_sec"], data["pre_dir"]
|
||
) # (N, K_MAX, 4)
|
||
# Padding slots carry sentinel pdg 0 (see loader._pad_list_col_int),
|
||
# which is never a real PDG code, so `.get(..., 0)` naturally maps
|
||
# both real unknown codes and padding to the same masked-out index.
|
||
sec_pdg_idx = np.vectorize(lambda p: pdg_map.get(int(p), 0))(
|
||
sec_pdg_list
|
||
).astype(np.int64)
|
||
else:
|
||
# Guard against silently training Stage 2 on zeroed targets: if any step
|
||
# actually spawned secondaries (n_sec > 0, from child_track_ids) but the
|
||
# per-secondary columns are absent, the file was never run through the
|
||
# parent->child join (steps_to_parquet._add_secondary_attributes /
|
||
# `dwarf convert`). Zero-filling here would collapse every secondary to
|
||
# PDG index 0 and a constant energy fraction — a broken Stage 2 with no
|
||
# error. Callers that only need Stage-1 (e.g. `giant predict`) keep the
|
||
# default require_secondaries=False.
|
||
if require_secondaries and n_sec_raw.max(initial=0) > 0:
|
||
n_with_sec = int((n_sec_raw > 0).sum())
|
||
raise ValueError(
|
||
f"{n_with_sec} step(s) have secondaries (n_sec > 0) but the "
|
||
"per-secondary columns (sec_E_list / sec_pdg_list / sec_dx_list "
|
||
"…) are missing. This parquet was not run through the "
|
||
"parent->child join (steps_to_parquet._add_secondary_attributes "
|
||
"/ `dwarf convert`); training on it would silently zero all "
|
||
"Stage-2 targets. Re-convert the file, or pass "
|
||
"require_secondaries=False for Stage-1-only use."
|
||
)
|
||
N = len(n_sec)
|
||
sec_cont = np.zeros((N, K_MAX, 4), dtype=np.float32)
|
||
sec_pdg_idx = np.zeros((N, K_MAX), dtype=np.int64)
|
||
|
||
if fit:
|
||
cond_normalizer = Normalizer().fit(cond_cont)
|
||
target_normalizer = Normalizer().fit(target_s1)
|
||
|
||
if cond_normalizer is not None:
|
||
cond_cont = cond_normalizer.transform(cond_cont)
|
||
if target_normalizer is not None:
|
||
target_s1 = target_normalizer.transform(target_s1)
|
||
|
||
process = data.get("process")
|
||
if proc_map is not None and process is not None:
|
||
proc_idx = np.array([proc_map[str(p)] for p in process], dtype=np.int64)
|
||
else:
|
||
proc_idx = np.zeros(len(cond_cat), dtype=np.int64)
|
||
|
||
return (
|
||
cond_cont,
|
||
cond_cat,
|
||
target_s1,
|
||
n_sec,
|
||
sec_cont,
|
||
sec_pdg_idx,
|
||
proc_idx,
|
||
cond_normalizer,
|
||
target_normalizer,
|
||
)
|