code/scripts/make_goldens.py

#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Goldens of frnet-p150 from the independent fp32 CPU reference (research venv: torch, onnx, onnxruntime; CPU only).

    cd bundles/frnet-p150 && OMP_NUM_THREADS=4 PYTHONPATH=code \\
        /home/ubuntu/experiments/tt-models/tools/research-venv/bin/python code/scripts/make_goldens.py --set all

Sets (output under ``research/frnet/goldens`` unless ``--out``; never shipped, except ``host``):

- ``taps``: the PandaSet 019 f40 sample (staged) in both variants through the bundle's host pre-processing and the
  torch fp32 reference with every module tap (``reference/weights.py`` ``MODULE_TAPS`` + the per-voxel max of every
  point branch, ``backbone.segmax0-4``) -> ``<variant>/pandaset_019_40.taps.npz`` (+ the four network inputs as
  ``input.*`` and ONNX Runtime's ``pred_probs`` as ``ort.pred_probs``) and a ``.json`` with each tap's shape,
  sha256 and its PCC / max |diff| against the same tensor of the deployed graph under ONNX Runtime (intermediates);
- ``frames``: every catalogued frame (``tests/_data.py`` ``FRAMES``: 16 OT128 + 4 QT128 public goldens, the VLS-128
  x > 0 halves) through the torch reference vs the stored ORT ``pred_probs``: PCC, max |diff|, argmax agreement (raw
  rows and all rows) and the ORT top-2 margin of every disagreeing point -> ``manifest.json``;
- ``host``: ``code/tt_frnet/tests/goldens/preprocess.json`` -- sha256 of the four network inputs that the RESEARCH
  pre-processing reference (``research/frnet/scripts/frnet_preproc_ref.py``) produces for the PandaSet 019 f40 sample (with
  its crop transform) and for the research synthetic scans, so the host tests can check the bundle's
  pre-processing where the research tree is absent.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import sys
import time
from pathlib import Path

import numpy as np

HERE = Path(__file__).resolve()
CODE = HERE.parents[1]
sys.path.insert(0, str(CODE))

WS = Path("/home/ubuntu/experiments/tt-models")
RES = WS / "research" / "frnet"
STEM = "pandaset_019_40"


def sha(a: np.ndarray) -> str:
    return hashlib.sha256(np.ascontiguousarray(a).tobytes()).hexdigest()


def pcc64(a: np.ndarray, b: np.ndarray) -> float:
    from tt_frnet.ttaw.metrics import pcc

    return float(pcc(a, b))


def make_taps(out: Path, weights_dir: str, threads: int) -> None:
    import torch

    from tt_frnet.reference.ort import OrtFRNet
    from tt_frnet.reference.pipeline import FRNetReference
    from tt_frnet.reference.weights import MODULE_TAPS
    from tt_frnet.ttaw.golden import TapRegistry

    torch.set_num_threads(threads)
    patterns = list(MODULE_TAPS) + ["backbone.segmax*"]
    for variant in ("ot128", "qt128"):
        t0 = time.time()
        ref = FRNetReference(weights_dir, variant, backend="torch", threads=threads)
        taps = TapRegistry(include=patterns)
        staged = CODE.parent / "staging_samples_pandaset"              # git-ignored (tests/_data.py)
        T = json.loads((staged / "calib" / f"{STEM}.json").read_text())["T_base_link_from_sensor"]
        run = ref.run(staged / "samples" / f"{STEM}.npz", calibration={"T_base_link_from_sensor": T}, taps=taps)
        t_torch = time.time() - t0
        inp = run.prepared.inputs
        ort_taps = TapRegistry(include=patterns)
        ort = OrtFRNet(weights_dir, variant, threads=threads, intermediates=True, tap_patterns=patterns,
                       weights=ref.weights)
        probs_ort = ort(inp.points, inp.coors, inp.voxel_coors, inp.inverse_map, taps=ort_taps)
        del ort
        gold = np.load(RES / "public_data" / "golden" / f"{STEM}_{variant}.npz", allow_pickle=False)
        if not np.array_equal(probs_ort, gold["autoware_pred_probs"]):
            sys.exit(f"{variant}: ORT pred_probs differ from the research golden (onnxruntime version?)")
        meta = {"frame": STEM, "variant": variant, "reference": "tt_frnet.reference.model (torch fp32)",
                "onnx_sha256": ref.weights.onnx_sha256, "torch": torch.__version__, "threads": threads,
                "N": inp.num_points, "N_raw": inp.num_points_raw, "M": inp.num_voxels,
                "torch_s": round(t_torch, 1), "taps": {}}
        arrays = {f"input.{k}": v for k, v in inp.onnx_feeds().items()}
        arrays["ort.pred_probs"] = probs_ort
        for name in taps.names():
            v = np.ascontiguousarray(taps[name], np.float32)
            arrays[name] = v
            entry = {"shape": list(v.shape), "sha256": sha(v)}
            if name in ort_taps:
                o = np.asarray(ort_taps[name], np.float32)
                entry.update(pcc_vs_ort=pcc64(v, o), max_abs_vs_ort=float(np.abs(v - o).max()))
            meta["taps"][name] = entry
        worst = min((e["pcc_vs_ort"], k) for k, e in meta["taps"].items() if "pcc_vs_ort" in e)
        meta["worst_tap_vs_ort"] = {"tap": worst[1], "pcc": worst[0]}
        d = out / variant
        d.mkdir(parents=True, exist_ok=True)
        path = d / f"{STEM}.taps.npz"
        np.savez(path, **arrays)
        meta["file_sha256"] = hashlib.sha256(path.read_bytes()).hexdigest()
        meta["bytes"] = path.stat().st_size
        (d / f"{STEM}.taps.json").write_text(json.dumps(meta, indent=1) + "\n")
        print(f"taps {variant}: {len(taps.names())} taps, {path.stat().st_size / 1e9:.2f} GB, worst vs ORT "
              f"{worst[1]} {worst[0]:.12f}, {time.time() - t0:.0f} s", flush=True)
        del taps, ort_taps, arrays


def make_frames(out: Path, weights_dir: str, threads: int) -> None:
    import torch

    from tt_frnet.reference.model import FRNetTorch
    from tt_frnet.reference.weights import load_weights
    from tt_frnet.tests import _data

    torch.set_num_threads(threads)
    nets = {}
    manifest = {"what": "torch fp32 reference vs the ORT pred_probs stored in each research golden",
                "torch": torch.__version__, "threads": threads, "frames": {}}
    for name, (rel, prefix, variant) in _data.FRAMES.items():
        path = RES / rel
        if not path.is_file():
            print(f"{name}: missing {rel}", flush=True)
            continue
        if variant not in nets:
            nets[variant] = FRNetTorch(load_weights(weights_dir, variant))
        inputs, ort_probs, g = _data.load_frame(name)
        t0 = time.time()
        probs = nets[variant](inputs.points, inputs.coors, inputs.voxel_coors, inputs.inverse_map)
        dt = time.time() - t0
        n_raw = inputs.num_points_raw
        a, b = probs.argmax(1), ort_probs.argmax(1)
        diff = np.flatnonzero(a != b)
        top2 = np.sort(ort_probs[diff], axis=1)[:, -2:] if len(diff) else np.zeros((0, 2), np.float32)
        counts = np.bincount(inputs.inverse_map)
        manifest["frames"][name] = {
            "golden": f"research/frnet/{rel}", "variant": variant, "N": inputs.num_points, "N_raw": n_raw,
            "M": inputs.num_voxels, "max_points_per_cell": int(counts.max()),
            "pcc": pcc64(probs, ort_probs), "max_abs": float(np.abs(probs - ort_probs).max()),
            "argmax_agreement": float((a == b).mean()), "argmax_agreement_raw": float((a[:n_raw] == b[:n_raw]).mean()),
            "disagreements": int(len(diff)),
            "disagreement_ort_margin_max": float((top2[:, 1] - top2[:, 0]).max()) if len(diff) else 0.0,
            "torch_s": round(dt, 1)}
        print(f"{name}: N {inputs.num_points} M {inputs.num_voxels} pcc {manifest['frames'][name]['pcc']:.12f} "
              f"max|d| {manifest['frames'][name]['max_abs']:.2e} argmax {manifest['frames'][name]['argmax_agreement']:.6f}"
              f" ({len(diff)} flips) {dt:.1f} s", flush=True)
    out.mkdir(parents=True, exist_ok=True)
    (out / "manifest.json").write_text(json.dumps(manifest, indent=1) + "\n")


def make_host(weights_dir: str) -> None:
    from tt_frnet.reference.config import FRNetConfig

    sys.path.insert(0, str(RES / "scripts"))
    import frnet_preproc_ref as R  # noqa: E402  (workspace research code, numpy only)

    out = {"what": "sha256 of the four network inputs produced by research/frnet/scripts/frnet_preproc_ref.py "
                   "(variant 'autoware'); the bundle's host pre-processing must reproduce them",
           "cases": {}}
    pub = RES / "public_data"
    for variant in ("ot128", "qt128"):
        cfg = FRNetConfig.from_model_dir(weights_dir, variant)
        z = np.load(pub / "inputs" / f"{STEM}.npz", allow_pickle=False)
        T = z["T_base_link_sensor"]
        crop12 = np.concatenate([T[:3, :3].reshape(-1), T[:3, 3]]).astype(np.float32)
        cases = {f"{STEM}.{variant}": (z["xyzi"], crop12, cfg.crop_bounds)}
        pts, _ = R.synthetic_scan(R.CFGS[variant], n_cols=600 if variant == "qt128" else 1200, seed=0)
        cases[f"synthetic_seed0.{variant}"] = (pts, None, None)
        for name, (xyzi, crop, bounds) in cases.items():
            pre = R.autoware_preprocess(xyzi, R.CFGS[variant], sensor_to_ref_12=crop, crop_bounds=bounds,
                                        variant="autoware")
            out["cases"][name] = {"N0": int(len(xyzi)), "N_raw": int(pre["num_points_raw"]),
                                  "N": int(pre["num_points"]), "M": int(pre["num_unique"]),
                                  "input_sha256": sha(np.ascontiguousarray(xyzi, np.float32)),
                                  **{k: sha(pre[k]) for k in ("points", "coors", "voxel_coors", "inverse_map")}}
            print(f"host {name}: N {pre['num_points']} M {pre['num_unique']}", flush=True)
    path = CODE / "tt_frnet" / "tests" / "goldens" / "preprocess.json"
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(out, indent=1) + "\n")


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("--set", choices=("taps", "frames", "host", "all"), default="all")
    ap.add_argument("--out", type=Path, default=RES / "goldens")
    ap.add_argument("--weights-dir", default=str(WS / "assets" / "frnet" / "hf_lidar_frnet"))
    ap.add_argument("--threads", type=int, default=4)
    a = ap.parse_args()
    if a.set in ("host", "all"):
        make_host(a.weights_dir)
    if a.set in ("frames", "all"):
        make_frames(a.out, a.weights_dir, a.threads)
    if a.set in ("taps", "all"):
        make_taps(a.out, a.weights_dir, a.threads)
    return 0


if __name__ == "__main__":
    sys.exit(main())