code/tt_ptv3/tests/test_reference_cpu.py

# SPDX-License-Identifier: Apache-2.0
"""The CPU reference (``tt_ptv3.reference``) against its oracles (PLAN.md 0.3 item 1; research/ptv3/SPEC.md 6).
Runs in the research venv (onnx, onnxruntime, torch CPU); skipped where they are absent.

    OMP_NUM_THREADS=2 PYTHONPATH=code tools/research-venv/bin/python -m pytest -q code/tt_ptv3/tests/test_reference_cpu.py

Oracles:

- encoder + seg head: the **unmodified deployed graphs** under ``onnx.reference.ReferenceEvaluator`` with numpy
  versions of the four ``autoware::`` ops (``reference/ort.py``; ORT cannot load them), fed the identical host tables
  -- on cropped frames (the numpy graph run is slow) in both neighbour modes; and, on full frames, the independent
  research torch reference ``research/ptv3/scripts/ptv3_torch_ref.py`` (raw initializers, unfolded BNs; verified
  against the graph at PCC 1.0 in SPEC 6) and the research public goldens;
- det head: ONNX Runtime on the deployed ``ptv3_det3d_head.onnx``: every tap before the data-dependent TopK at PCC
  >= 0.9999, the top-500 **set-based**, then teacher-forced (the reference consumes ORT's TopK indices) PCC >= 0.9999
  and identical post-processed boxes (PLAN.md 0.3 item 1).
"""
from __future__ import annotations

import os

import numpy as np
import pytest

from tt_ptv3.tests import _data as D
from tt_ptv3.ttaw.metrics import pcc

PCC = 0.9999
os.environ.setdefault("OMP_NUM_THREADS", "2")


@pytest.fixture(scope="module")
def ref():
    D.need("torch")
    D.need("onnx")
    from tt_ptv3.reference.pipeline import Reference, set_threads
    set_threads(2)
    return Reference(D.weights_dir())


def _taps_ok(a, b, name, thr=PCC):
    v = pcc(np.asarray(a, np.float64), np.asarray(b, np.float64))
    assert v >= thr, f"{name}: PCC {v:.7f} < {thr}"
    return v


# ---------------------------------------------------------------------------------------------- weights
def test_weights_structure_and_folds(ref):
    from tt_ptv3.ttaw.weights import OnnxWeights
    W = ref.weights
    assert len(W.encoder.stages) == 5 and [len(s.blocks) for s in W.encoder.stages] == [2] * 5
    assert [b.heads for s in W.encoder.stages for b in s.blocks] == [2, 2, 4, 4, 8, 8, 16, 16, 32, 32]
    assert [b.order for s in W.encoder.stages for b in s.blocks] == [0, 1] * 5
    assert sorted(W.seg.dec) == [0, 1, 2, 3] and W.seg.dec[0].block is None and W.seg.dec[3].block is not None
    assert W.det.bev_pos.shape == (65536, 2) and W.det.bev_pos[257].tolist() == [1.5, 1.5]
    assert W.det.bev_pos[1].tolist() == [1.5, 0.5]                                  # channel 0 along x
    # 34.05M parameters (SPEC 4: 28.15M + 3.55M + 2.35M) minus the folded BN statistics
    assert 33_900_000 < W.num_parameters() < 34_100_000
    # the stem fold equals Linear then BN (float64)
    w = OnnxWeights(D.weights_dir() / "ptv3_encoder.onnx")
    g = w.gemm("/embedding/stem/stem.0/Gemm")
    bn = w.batchnorm("/embedding/stem/stem.1/BatchNormalization")
    x = np.random.default_rng(0).normal(size=(64, 4))
    lit = bn.apply(x @ g.weight.T.astype(np.float64) + g.bias, axis=1)
    fold = x @ W.encoder.embed.w.T.astype(np.float64) + W.encoder.embed.b
    assert np.abs(lit - fold).max() < 1e-5


# ---------------------------------------------------------------------------------------------- graph oracle
@pytest.mark.parametrize("frame,radius", [("sample", 12.0), ("t4_scene10_00795", 10.0), ("ns0103_k20", 15.0)])
@pytest.mark.parametrize("mode", ["alias-first", "exact"])
def test_encoder_seg_vs_unmodified_graph(ref, frame, radius, mode):
    from tt_ptv3.reference.ort import ENCODER_GRAPH_TAPS, GraphReference, crop_points
    pts = crop_points(D.load_frame(frame), radius)
    taps: dict = {}
    res = ref.run(pts, neighbour_mode=mode, det=False, taps=taps)
    graph = GraphReference(D.weights_dir())
    pf, kept = graph.encoder(res.frame, keep=list(ENCODER_GRAPH_TAPS.values()))
    for s in range(5):
        _taps_ok(res.point_feat[s], pf[s], f"point_feat_{s}")
    for name, tensor in ENCODER_GRAPH_TAPS.items():
        _taps_ok(taps[name], kept[tensor], name)
    labels, probs = graph.seg_head(res.frame, pf)
    _taps_ok(res.pred_probs, probs, "pred_probs")
    assert np.mean(res.pred_labels == labels) >= PCC


# ---------------------------------------------------------------------------------------------- research reference
def _research_stage_tables(frame):
    import torch
    st = []
    for s, lv in enumerate(frame.levels):
        d = dict(grid=lv.grid, order=torch.from_numpy(lv.order), inverse=torch.from_numpy(lv.inverse),
                 pair=torch.from_numpy(np.ascontiguousarray(lv.nmap.T.astype(np.int64))))
        if s > 0:
            p = frame.pools[s - 1]
            d.update(indptr=torch.from_numpy(p.indptr), indices=torch.from_numpy(p.indices))
        st.append(d)
    return st


@pytest.mark.parametrize("mode", ["alias-first", "exact"])
def test_full_frame_vs_research_torch_reference(ref, mode):
    import torch
    D.research_modules()
    import ptv3_torch_ref as R
    pts = D.load_frame("sample")
    res = ref.run(pts, neighbour_mode=mode, det=False)
    W = R.load_weights(str(D.weights_dir() / "ptv3_encoder.onnx"))
    W.update(R.load_weights(str(D.weights_dir() / "ptv3_seg3d_head.onnx")))
    st = _research_stage_tables(res.frame)
    with torch.no_grad():
        pf = R.encoder(W, res.frame.encoder_inputs(), st)
        labels, probs = R.seg_head(W, pf, [torch.from_numpy(c) for c in res.frame.clusters], st)
    for s in range(5):
        _taps_ok(res.point_feat[s], pf[s].numpy(), f"point_feat_{s}")
    _taps_ok(res.pred_probs, probs.numpy(), "pred_probs")
    assert np.mean(res.pred_labels == labels.numpy()) >= PCC


@pytest.mark.parametrize("frame", ["ps019_f40", "ns0103_k20"])
@pytest.mark.parametrize("mode", ["exact", "alias-first"])
def test_vs_research_public_goldens(ref, frame, mode):
    """The research goldens (research/ptv3/public_data, same faithful tables) are reproduced: features, voxel and
    point labels, det outputs and the final boxes."""
    gold = D.public_golden(frame, mode)
    res = ref.run(D.load_frame(frame), neighbour_mode=mode)
    for s in range(5):
        _taps_ok(res.point_feat[s], gold[f"point_feat_{s}"], f"point_feat_{s}")
    assert np.mean(res.pred_labels == gold["pred_labels"]) >= PCC
    assert np.mean(res.seg.labels == gold["point_label15"]) >= PCC
    _taps_ok(res.det["dense_heatmap"], gold["det_dense_heatmap"], "dense_heatmap")
    from tt_ptv3.ttaw.metrics import match_detections
    b, g = res.boxes.array(), gold["boxes"]
    m = match_detections(b[:, 2:4], b[:, 0], b[:, 1], g[:, 2:4], g[:, 0], g[:, 1], max_dist=0.05)
    assert m.recall == 1.0 and m.precision == 1.0, m.to_dict()
    assert np.abs(np.sort(b[:, 1]) - np.sort(g[:, 1])).max() < 1e-4


# ---------------------------------------------------------------------------------------------- det head vs ORT
@pytest.fixture(scope="module")
def det_case(ref):
    D.need("onnxruntime")
    from tt_ptv3.reference.ort import DetHeadOrt
    res = ref.run(D.load_frame("sample"), neighbour_mode="alias-first", seg=False, det=False)
    ort = DetHeadOrt(D.weights_dir(), intermediates=True)
    outs, otaps = ort.run(res.frame.det_head_inputs(res.point_feat))
    return res, outs, otaps


def _det_run(ref, res, topk=None):
    import torch
    from tt_ptv3.host.preprocess import bev_index
    from tt_ptv3.reference import model as M
    taps: dict = {}
    with torch.no_grad():
        pf = [torch.from_numpy(p) for p in res.point_feat]
        bev = bev_index(res.frame.pools[2].grid_coord, sentinel=res.frame.counts[3])
        out = M.det_head_forward(ref.weights.det, pf, res.frame, bev, topk_indices=topk, taps=taps)
    return out, taps


def test_det_head_before_selection_vs_ort(ref, det_case):
    res, outs, otaps = det_case
    out, taps = _det_run(ref, res)
    for name in ("det.fusion", "det.point_proj", "det.bev", "det.bev_refine", "det.stem", "det.block0", "det.block1",
                 "det.block2", "det.block3", "det.shared", "det.heat_nms"):
        a, b = taps[name], otaps[name]
        _taps_ok(a.reshape(b.shape) if a.size == b.size else a, b, name)
    _taps_ok(out["dense_heatmap"], outs["dense_heatmap"], "dense_heatmap")
    _taps_ok(taps["det.key_pos_embed"], otaps["det.key_pos_embed"].reshape(taps["det.key_pos_embed"].shape),
             "key_pos_embed")
    # data-dependent selection: set-based (two fp32 implementations disagree on tail ties)
    from tt_ptv3.ttaw.ops.topk import selection_metrics
    ref_idx = otaps["det.topk_indices"].reshape(-1)
    m = selection_metrics(taps["det.topk_indices"], ref_idx,
                          reference_scores=otaps["det.heat_nms"].reshape(-1)[ref_idx])
    assert m["shared"] >= 0.98 * 500, m
    assert not m["strong_missed"], m


def test_det_head_teacher_forced_vs_ort(ref, det_case):
    from tt_ptv3.host.postprocess import postprocess_detections
    res, outs, otaps = det_case
    out, taps = _det_run(ref, res, topk=otaps["det.topk_indices"].reshape(-1))
    assert np.array_equal(out["query_labels"], outs["query_labels"])
    _taps_ok(taps["det.query_feat"].T, otaps["det.query_feat"][0], "query_feat")
    _taps_ok(taps["det.query_pos_embed"], otaps["det.query_pos_embed"][0], "query_pos_embed")
    for n in ("det.norm1", "det.norm2", "det.norm3"):
        _taps_ok(taps[n], otaps[n][0], n)
    for k in ("query_heatmap_score", "heatmap", "center", "height", "dim", "rot", "vel"):
        _taps_ok(out[k], outs[k], k)
    a, b = postprocess_detections(out), postprocess_detections(outs)
    assert len(a) == len(b) and len(a) > 0
    assert np.array_equal(a.label, b.label) and np.abs(a.array() - b.array()).max() < 1e-3