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

Wire up predict CLI to load and run the Stage-2 sec_decoder

Browse Source 
`giant predict` only ever ran Stage 1, echoing ground-truth n_sec instead
of predicting it — Phase 2 training already produced a joint checkpoint
but nothing consumed the sec_decoder half of it. Loads sec_decoder
alongside the Stage-1 model (filtering model_config per-model, since
splatting it whole into either constructor breaks on the other's
sec_slot_dim/k_max-only keys), runs sample_secondaries + PDG snapping in
--coord global mode, and appends predicted n_sec/species/energy/direction
columns to the output parquet.

Also fixes decode_secondaries rotating raw (non-unit) flow output straight
into world frame without normalizing first — a rotation preserves
magnitude, so un-normalized ODE output produced non-unit secondary
directions, caught via an end-to-end smoke test.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
Lars Bogner
2026-07-06 14:06:07 +02:00
parent 4ba419ebe4
commit 0c96ddf765
2 changed files with 104 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
giant/cli.py
+87 -10
View File
@@ -30,15 +30,34 @@ from giant.data.loader import (
from giant.data.transforms import (
build_features,
build_cond_features,
decode_secondaries,
energy_simplex_decode,
inv_local_frame_rotation,
inv_log_transform,
reconstruct_post_pos,
Normalizer,
)
from giant.model.network import DenoisingMLP
from giant.model.network import DenoisingMLP, SecondaryDecoder
from giant.pipeline import run_train_job
from giant.sample import sample_flow
from giant.sample import sample_flow, sample_secondaries, snap_type_to_pdg_idx
_STAGE1_MODEL_KEYS = {
"pdg_vocab",
"mat_vocab",
"hidden_dim",
"n_blocks",
"emb_dim",
"dropout",
"k_max",
}
_SEC_DECODER_MODEL_KEYS = {
"pdg_vocab",
"mat_vocab",
"hidden_dim",
"n_blocks",
"emb_dim",
"dropout",
}
app = typer.Typer(no_args_is_help=True)
@@ -354,6 +373,13 @@ def predict(
)
raise typer.Exit(1)
if "sec_decoder" not in ckpt:
typer.echo(
"error: checkpoint has no sec_decoder — retrain with the current code",
err=True,
)
raise typer.Exit(1)
model_cfg = ckpt["model_config"]
if batch_size_auto:
@@ -370,16 +396,27 @@ def predict(
typer.echo(
f"batch_size: {batch_size_value} (auto-estimated from free GPU memory)"
)
assert batch_size_value is not None
bs = batch_size_value
pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()}
pdg_map_inv = {v: k for k, v in pdg_map.items()}
mat_map = {str(k): v for k, v in ckpt["mat_map"].items()}
cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"])
tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"])
model = DenoisingMLP(**model_cfg)
model = DenoisingMLP(
**{k: v for k, v in model_cfg.items() if k in _STAGE1_MODEL_KEYS}
)
model.load_state_dict(ckpt["model"])
model.to(_device).eval()
sec_decoder = SecondaryDecoder(
**{k: v for k, v in model_cfg.items() if k in _SEC_DECODER_MODEL_KEYS}
)
sec_decoder.load_state_dict(ckpt["sec_decoder"])
sec_decoder.to(_device).eval()
typer.echo(f"loaded checkpoint: {checkpoint}")
gconfig.warn_if_checkpoint_config_mismatch(checkpoint)
@@ -419,8 +456,20 @@ def predict(
cc = torch.from_numpy(cond_cont).float().to(_device)
ck = torch.from_numpy(cond_cat).long().to(_device)
pred, _n_sec = sample_flow(model, cc, ck, steps=steps)
pred = pred.cpu().numpy() # normalised
stage1_norm, n_sec_pred = sample_flow(model, cc, ck, steps=steps)
if coord == Coord.global_:
sec_cont, sec_type_emb, _sec_valid_pred = sample_secondaries(
sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps=steps
)
sec_pdg_idx = snap_type_to_pdg_idx(
sec_type_emb, model.pdg_embedding_weight()
)
sec_cont_np = sec_cont.cpu().numpy()
sec_pdg_idx_np = sec_pdg_idx.cpu().numpy()
n_sec_pred_np = n_sec_pred.cpu().numpy()
pred = stage1_norm.cpu().numpy() # normalised
# Inverse-normalise → local frame, log-scaled scalars
raw = tgt_norm.inverse_transform(pred)
@@ -454,8 +503,10 @@ def predict(
step_length = inv_log_transform(raw[:, 0])
# Columns 1:3 are ALR coords of the deposit/secondary/post energy
# simplex; decode them against pre_E so edep + e_sec + post_E == pre_E
# (hence delta_e == edep + e_sec) holds by construction.
edep, _e_sec, _post_E, delta_e = energy_simplex_decode(
# (hence delta_e == edep + e_sec) holds by construction. e_sec_pred
# doubles as the stick-breaking energy budget for the Stage-2 decode
# below, since the model has no other source for it at inference.
edep, e_sec_pred, _post_E, delta_e = energy_simplex_decode(
raw[:, 1:3], piece["pre_E"]
)
@@ -474,6 +525,28 @@ def predict(
piece["pre_pos"], piece["pre_dir"], step_length, travel_dir_local
)
sec_E, sec_dir_world, sec_pdg_code, _sec_valid = decode_secondaries(
sec_cont_np,
sec_pdg_idx_np,
n_sec_pred_np,
e_sec_pred,
piece["pre_dir"],
pdg_map_inv,
)
sec_pdg_list = [
sec_pdg_code[i, :n].tolist() for i, n in enumerate(n_sec_pred_np)
]
sec_E_list = [sec_E[i, :n].tolist() for i, n in enumerate(n_sec_pred_np)]
sec_dx_list = [
sec_dir_world[i, :n, 0].tolist() for i, n in enumerate(n_sec_pred_np)
]
sec_dy_list = [
sec_dir_world[i, :n, 1].tolist() for i, n in enumerate(n_sec_pred_np)
]
sec_dz_list = [
sec_dir_world[i, :n, 2].tolist() for i, n in enumerate(n_sec_pred_np)
]
table = pa.table(
{
"event_id": piece["event_id"],
@@ -488,6 +561,7 @@ def predict(
"material": piece["material"],
"layer_id": piece["layer_id"],
"n_sec": piece["n_sec"],
"n_sec_pred": n_sec_pred_np,
"step_length": step_length,
"delta_e": delta_e,
"edep": edep,
@@ -497,6 +571,11 @@ def predict(
"post_x": post_pos_world[:, 0],
"post_y": post_pos_world[:, 1],
"post_z": post_pos_world[:, 2],
"sec_pdg_list": sec_pdg_list,
"sec_E_list": sec_E_list,
"sec_dx_list": sec_dx_list,
"sec_dy_list": sec_dy_list,
"sec_dz_list": sec_dz_list,
}
)
@@ -544,9 +623,7 @@ def predict(
if writer is not None:
writer.close()
ref_path = _write_prediction_ref(
checkpoint, pred_uuid, out, dataset_path, comment
)
ref_path = _write_prediction_ref(checkpoint, pred_uuid, out, dataset_path, comment)
typer.echo(f"reference: {ref_path}")
if skipped:
giant/data/transforms.py
+17 -5
View File
@@ -316,9 +316,13 @@ def encode_secondaries(
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))
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)
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.
@@ -358,8 +362,13 @@ def decode_secondaries(
Returns (sec_E, sec_dir_world, sec_pdg_code, sec_valid) each shape (N, K_MAX).
"""
N, K, _ = sec_cont.shape
stick_logits = sec_cont[:, :, 0] # (N, K)
dir_local = sec_cont[:, :, 1:] # (N, K, 3)
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)))
@@ -381,7 +390,10 @@ def decode_secondaries(
)
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)],
[
[pdg_map_inv.get(int(sec_pdg_pred[n, i]), 0) for i in range(K)]
for n in range(N)
],
dtype=np.int32,
)