code_snapshots/WiSER_public_release_code/scripts/evaluate_dual.py

"""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 <scene_id>/wireless assets; overrides manifest paths for portable releases.")
    p.add_argument("--scene3d-root", default=None,
                   help="Root containing <scene_id>/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()