From 0c96ddf765d6778aa7a467418a2b99891851f6a9 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Mon, 6 Jul 2026 14:06:07 +0200 Subject: [PATCH] Wire up predict CLI to load and run the Stage-2 sec_decoder MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `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 --- giant/cli.py | 97 +++++++++++++++++++++++++++++++++++----- giant/data/transforms.py | 22 ++++++--- 2 files changed, 104 insertions(+), 15 deletions(-) diff --git a/giant/cli.py b/giant/cli.py index e74b868..d9f8208 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -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: diff --git a/giant/data/transforms.py b/giant/data/transforms.py index 8e5aa25..416eca4 100644 --- a/giant/data/transforms.py +++ b/giant/data/transforms.py @@ -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, )