"""WiSER Final dual-task validation. Loads the final phase checkpoint and runs both: - Radiomap val (formal100 held-out 80 samples) - CIR val (tiny10 triples, capped to val_max_triples) Writes a summary JSON with both sets of metrics + verdict. """ from __future__ import annotations import argparse import json import sys from pathlib import Path import torch import torch.distributed as dist from torch.utils.data import DataLoader from torch.utils.data.distributed import DistributedSampler ROOT = Path(__file__).resolve().parent.parent sys.path.insert(0, str(ROOT)) from wiser.config import ModelConfig as SharedModelConfig, TrainConfig from wiser.data.collate import triple_collate from wiser.data.csi_path_targets import MergedPathTargetConfig from wiser.data.dataset import MultiSceneTripleDataset from wiser.data.radiomap_dataset import RadiomapDataset, radiomap_collate from wiser.engine.matching import MatcherConfig from wiser.engine.trainer import CsiSetTrainer, TrainerConfig from wiser.utils.metrics import compute_metric_bundle from wiser.alt_models import JointRadiomapCIRModel from wiser.alt_engine.checkpoint import ( detect_csi_head_arch_from_ckpt, load_phase_ckpt_into_model, ) def parse_args(): p = argparse.ArgumentParser() p.add_argument("--ckpt", required=True, help="Final phase ckpt path") p.add_argument("--radiomap-manifest", required=True) p.add_argument("--csi-manifest", "--cir-manifest", required=True) p.add_argument("--d22-ckpt", required=True, help="For CIR head arch detection") p.add_argument("--out-json", required=True) p.add_argument("--device", default="cuda:0") p.add_argument("--wireless-root", default=None, help="Root containing /wireless assets; overrides manifest paths for portable releases.") p.add_argument("--scene3d-root", default=None, help="Root containing /voxel_10cm/mesh_voxel_cache_100mm.pt.") p.add_argument("--cir-dataset-tag", default="voxel_original_csi_path_10cm_1e6") p.add_argument("--radiomap-bs", type=int, default=4) p.add_argument("--csi-bs", "--cir-bs", type=int, default=2048) p.add_argument("--csi-max-triples", "--cir-max-triples", type=int, default=5000) p.add_argument("--channels", type=int, default=512) p.add_argument("--grid-h", type=int, default=36) p.add_argument("--grid-w", type=int, default=36) p.add_argument("--corridor-budget", type=int, default=192) p.add_argument("--num-cross-layers", type=int, default=4) p.add_argument("--num-self-layers", type=int, default=2) p.add_argument("--huber-beta", type=float, default=2.0) p.add_argument("--gradient-weight", type=float, default=1.0) return p.parse_args() def main(): args = parse_args() device = torch.device(args.device) torch.manual_seed(0) # Build model matching the ckpt's CIR head arch info = detect_csi_head_arch_from_ckpt(args.d22_ckpt) csi_arch_kwargs = {k: v for k, v in info.items() if k.startswith("csi_")} shared = SharedModelConfig( backbone_kind="trellis2", backbone_channels=args.channels, backbone_downsample_stages=info.get("backbone_downsample_stages", 3), backbone_blocks_per_stage=info.get("backbone_blocks_per_stage", 3), ) model = JointRadiomapCIRModel( shared=shared, channels=args.channels, spatial_channels=256, stem_depth=4, output_depth=4, grid_h=args.grid_h, grid_w=args.grid_w, corridor_budget=args.corridor_budget, num_cross_layers=args.num_cross_layers, num_self_layers=args.num_self_layers, **csi_arch_kwargs, ).to(device) model.eval() # Load ckpt meta = load_phase_ckpt_into_model(args.ckpt, model) print(f"[load] ckpt from phase={meta.get('phase_name')} ep={meta.get('epoch')}") # ===== Radiomap val ===== print("\n========== Radiomap validation (held-out 80) ==========") rm_manifest = json.loads(Path(args.radiomap_manifest).read_text()) rm_kwargs = {} rm_scene3d_root = args.scene3d_root or rm_manifest.get("scene3d_root") if rm_scene3d_root: rm_kwargs["scene3d_root"] = rm_scene3d_root rm_val = RadiomapDataset( rm_manifest["val_heldout"], channels=args.channels, grid_h=args.grid_h, grid_w=args.grid_w, db_floor=-300.0, dataset_kind="sionna_radiomap", **rm_kwargs, ) rm_loader = DataLoader(rm_val, batch_size=args.radiomap_bs, shuffle=False, collate_fn=radiomap_collate, num_workers=2, pin_memory=True) import torch.nn.functional as F from wiser.alt_models.joint_model import _extract_level_tensors rm_mae_sum = 0.0; rm_rmse_sum = 0.0; rm_bias_sum = 0.0; rm_n = 0 with torch.no_grad(): for batch in rm_loader: voxel_feats_list = [t.to(device, non_blocking=True) for t in batch["voxel_feats_list"]] voxel_coords_list = [t.to(device, non_blocking=True) for t in batch["voxel_coords_list"]] tx_xyz = batch["tx_xyz_norm"].to(device, non_blocking=True) rx_grid = batch["rx_grid_xyz_norm"].to(device, non_blocking=True) gt = batch["gt_radiomap_db"].to(device, non_blocking=True) mask = batch["extent_mask"].to(device, non_blocking=True) tx_xyz_m = batch["tx_xyz_metric"].to(device, non_blocking=True) rx_grid_m = batch["rx_grid_xyz_metric"].to(device, non_blocking=True) scene_ext_m = batch["scene_extent_xyz_m"].to(device, non_blocking=True) origin_m_list = [t.to(device, non_blocking=True) for t in batch["voxel_origin_m_list"]] cell_size_m = float(batch["voxel_cell_size_m"]) with torch.autocast(device_type="cuda", dtype=torch.bfloat16): pred = model.forward_radiomap( voxel_feats_list, voxel_coords_list, tx_xyz_norm=tx_xyz, tx_xyz_metric=tx_xyz_m, rx_grid_xyz_norm=rx_grid, rx_grid_xyz_metric=rx_grid_m, scene_extent_xyz_m=scene_ext_m, voxel_origin_m_list=origin_m_list, voxel_cell_size_m=cell_size_m, extent_mask=mask, ) mh = mask.squeeze(1).bool() if mh.sum() > 0: diff = (pred[mh].float() - gt[mh].float()) B = gt.shape[0] rm_mae_sum += float(diff.abs().mean().item()) * B rm_rmse_sum += float(diff.pow(2).mean().sqrt().item()) * B rm_bias_sum += float(diff.mean().item()) * B rm_n += B rm_mae = rm_mae_sum / max(rm_n, 1) rm_rmse = rm_rmse_sum / max(rm_n, 1) rm_bias = rm_bias_sum / max(rm_n, 1) # PSNR: dynamic range 120 dB import math as _math rm_psnr = 20.0 * _math.log10(120.0) - 20.0 * _math.log10(max(rm_rmse, 1e-6)) print(f"Radiomap val: n={rm_n} MAE={rm_mae:.3f} RMSE={rm_rmse:.3f} PSNR={rm_psnr:.2f} bias={rm_bias:+.3f}") # ===== CIR val ===== print("\n========== CIR validation ==========") from wiser.config import reject_radiomap_artifacts reject_radiomap_artifacts(args.csi_manifest) csi_manifest_data = json.loads(Path(args.csi_manifest).read_text()) target_cfg = MergedPathTargetConfig() csi_val = MultiSceneTripleDataset( csi_manifest_data, target_config=target_cfg, voxel_channels=args.channels, wireless_root=args.wireless_root or csi_manifest_data.get("wireless_root", None), scene3d_root=args.scene3d_root or csi_manifest_data.get("scene3d_root", None), dataset_tag=args.cir_dataset_tag or csi_manifest_data.get("dataset_tag", "voxel_original_csi_path_10cm_1e6"), ) if args.csi_max_triples > 0 and len(csi_val.triples) > args.csi_max_triples: csi_val.triples = csi_val.triples[:args.csi_max_triples] csi_loader = DataLoader(csi_val, batch_size=args.csi_bs, shuffle=False, collate_fn=triple_collate, num_workers=2, pin_memory=True) # Trainer for CIR (does Hungarian matching + metric compute) trainer_cfg = TrainerConfig() trainer_cfg.matcher = MatcherConfig(backend="scipy") from wiser.engine.losses import LossWeights trainer_cfg.loss_weights = LossWeights(no_object_exists=5.0) trainer_cfg.model = model.config trainer = CsiSetTrainer(model, trainer_cfg) csi_sum = {"pdp_cosine_nonzero": 0.0, "peak_db_mae_matched": 0.0, "peak_db_bias_matched": 0.0, "delay_mae_matched_ns": 0.0, "count_acc": 0.0, "nonzero_path_count_mae": 0.0, "zero_path_fpr": 0.0, "match_rate_over_all_samples": 0.0} csi_n = 0 with torch.no_grad(): for batch in csi_loader: for k, v in batch.items(): if isinstance(v, torch.Tensor): batch[k] = v.to(device, non_blocking=True) if "scene_voxel_levels" in batch: batch["scene_voxel_levels"] = [ {kk: (tt.to(device, non_blocking=True) if torch.is_tensor(tt) else tt) for kk, tt in sv.items()} if isinstance(sv, dict) else sv for sv in batch.get("scene_voxel_levels", []) ] gt = {k: batch[k] for k in ["gt_num_paths", "gt_delay_ns", "gt_peak_db", "gt_path_mask", "gt_truncated"]} out = trainer.step(batch, gt) mets = compute_metric_bundle(out["predictions"], gt, out["matchings"], query_budget=8) B = gt["gt_num_paths"].shape[0] for k in csi_sum: csi_sum[k] += float(mets.get(k, 0.0)) * B csi_n += B csi_out = {k: csi_sum[k] / max(csi_n, 1) for k in csi_sum} print(f"CIR val: n={csi_n} peak_mae={csi_out['peak_db_mae_matched']:.2f}dB delay_mae={csi_out['delay_mae_matched_ns']:.3f}ns count_acc={csi_out['count_acc']:.3f} fpr={csi_out['zero_path_fpr']:.3f} pdp_cosine={csi_out['pdp_cosine_nonzero']:.4f}") # ===== Verdict ===== # Baselines for context V16_MAE = 3.40; V22A_MAE = 3.86 D22_PEAK = 7.50; D22_COUNT = 0.495 if rm_mae <= 3.50 and csi_out["peak_db_mae_matched"] <= 12.0: verdict = "pipeline_validated_both_tasks_work" elif rm_mae <= 3.80 and csi_out["peak_db_mae_matched"] <= 15.0: verdict = "pipeline_works_but_suboptimal_consider_full" elif rm_mae > V22A_MAE + 0.3: verdict = "radiomap_regressed_investigate_P0_P2" elif csi_out["peak_db_mae_matched"] > 15.0: verdict = "csi_failed_to_restore_investigate_P1" else: verdict = "mixed_result_review_per_phase_logs" summary = { "final_ckpt": str(args.ckpt), "radiomap": { "val_mae_db": float(rm_mae), "val_rmse_db": float(rm_rmse), "val_psnr_db": float(rm_psnr), "val_bias_db": float(rm_bias), "val_n_samples": int(rm_n), "baseline_v16_best": V16_MAE, "baseline_v22a_best": V22A_MAE, "delta_vs_v16": float(rm_mae - V16_MAE), "delta_vs_v22a": float(rm_mae - V22A_MAE), }, "csi": { **{k: float(v) for k, v in csi_out.items()}, "val_n_triples": int(csi_n), "baseline_d22_best": {"peak_mae_db": D22_PEAK, "count_acc": D22_COUNT}, "delta_peak_mae_vs_d22": float(csi_out["peak_db_mae_matched"] - D22_PEAK), "delta_count_acc_vs_d22": float(csi_out["count_acc"] - D22_COUNT), }, "verdict": verdict, } out_path = Path(args.out_json) out_path.parent.mkdir(parents=True, exist_ok=True) with open(out_path, "w") as f: json.dump(summary, f, indent=2) print(f"\n✓ Dual-task val summary written: {out_path}") print(f"Verdict: {verdict}") if __name__ == "__main__": main()