SimurgArray/scripts/validate_p3_large.py

#!/usr/bin/env python3
"""Phase P3 validation on a large set of vendor GTCs: refit every sample's normalization transforms from the GTC's
raw X/Y (normalization 1.2.0 and 1.1.2) and compare them with the transforms stored in the GTC.

Usage: validate_p3_large.py --root DIR --bpm F --out summary.json [--workers N] [--limit N]
"""

from __future__ import annotations

import argparse
import json
import multiprocessing as mp
import os
import sys
import time
import traceback
from concurrent.futures import ProcessPoolExecutor
from pathlib import Path
from typing import Any

import numpy as np

FIELDS = ("offset_x", "offset_y", "scale_x", "scale_y", "shear", "theta")
VERSIONS = ("1.2.0", "1.1.2")
_S: dict[str, Any] = {}


def _init(bins: np.ndarray) -> None:  # type: ignore[type-arg]
    _S["bins"] = bins


def _ulps(a: np.ndarray, b: np.ndarray) -> np.ndarray:  # type: ignore[type-arg]
    ai = a.astype(np.float32).view(np.int32).astype(np.int64)
    bi = b.astype(np.float32).view(np.int32).astype(np.int64)
    # map sign-magnitude to a monotonic integer line so ULP distance works across zero
    ai = np.where(ai < 0, -(ai & 0x7FFFFFFF), ai)
    bi = np.where(bi < 0, -(bi & 0x7FFFFFFF), bi)
    return np.abs(ai - bi)


def check(path: str) -> dict[str, Any]:
    from simurg.formats.gtc import read_gtc
    from simurg.signal.normalization.fit import fit_transforms

    res: dict[str, Any] = {"gtc": path}
    try:
        g = read_gtc(path, check_write_complete=False)
        ref = g.normalization_transforms
        refm = np.array([[r[f] for f in FIELDS] for r in ref], dtype=np.float32)
        refr = np.array([r["reserved"] for r in ref], dtype=np.float32)
        res["autocall_version"] = g.autocall_version
        for v in VERSIONS:
            xf = fit_transforms(g.raw_x, g.raw_y, _S["bins"], v)
            mine = np.array([[getattr(x, f) for f in FIELDS] for x in xf], dtype=np.float32)
            mres = np.array([[x.cvx, x.cvy, x.nn12, x.rr12, x.taa, x.tbb] for x in xf], dtype=np.float32)
            if mine.shape != refm.shape:
                res[v] = {"error": f"{len(xf)} bins vs {len(ref)} in GTC"}
                continue
            u = _ulps(mine, refm)
            ur = _ulps(mres, refr)
            res[v] = {
                "bins_exact": int((u.max(axis=1) == 0).sum()),
                "n_bins": len(ref),
                "max_ulp": int(u.max()),
                "max_ulp_per_field": dict(zip(FIELDS, u.max(axis=0).tolist(), strict=True)),
                "reserved_exact": bool((ur == 0).all()),
                "reserved_ref_all_zero": bool((refr == 0).all()),
            }
    except Exception as e:
        res["error"] = f"{type(e).__name__}: {e}"
        res["trace"] = traceback.format_exc(limit=3)
    return res


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", required=True, type=Path)
    ap.add_argument("--bpm", required=True, type=Path)
    ap.add_argument("--out", required=True, type=Path)
    ap.add_argument("--workers", type=int, default=os.cpu_count() or 1)
    ap.add_argument("--limit", type=int, default=0)
    a = ap.parse_args()

    from simurg.formats.bpm import read_bpm
    from simurg.formats.discovery import walk_files

    t0 = time.time()
    gtcs = [str(p) for p in walk_files(a.root, (".gtc",))]
    if a.limit:
        gtcs = gtcs[: a.limit]
    bins = read_bpm(a.bpm).norm_bin_ids()
    with ProcessPoolExecutor(a.workers, mp_context=mp.get_context("spawn"), initializer=_init, initargs=(bins,)) as ex:
        results = list(ex.map(check, gtcs, chunksize=2))

    summary: dict[str, Any] = {"n_gtc": len(gtcs), "errors": [r for r in results if "error" in r][:20]}
    for v in VERSIONS:
        rs = [r[v] for r in results if isinstance(r.get(v), dict) and "bins_exact" in r[v]]
        summary[v] = {
            "samples_all_bins_exact": sum(r["bins_exact"] == r["n_bins"] for r in rs),
            "bins_exact": sum(r["bins_exact"] for r in rs),
            "bins_total": sum(r["n_bins"] for r in rs),
            "max_ulp": max((r["max_ulp"] for r in rs), default=None),
            "ulp_histogram": dict(
                sorted(
                    {
                        str(k): int(n)
                        for k, n in zip(*np.unique([r["max_ulp"] for r in rs], return_counts=True), strict=True)
                    }.items(),
                    key=lambda kv: int(kv[0]),
                )
            ),
            "max_ulp_per_field": {f: max((r["max_ulp_per_field"][f] for r in rs), default=None) for f in FIELDS},
            "reserved_exact_samples": sum(r["reserved_exact"] for r in rs),
            "reserved_ref_all_zero_samples": sum(r["reserved_ref_all_zero"] for r in rs),
        }
    summary["autocall_versions"] = sorted({r.get("autocall_version", "") for r in results})
    worst = sorted(
        (r for r in results if isinstance(r.get("1.1.2"), dict) and "max_ulp" in r["1.1.2"]),
        key=lambda r: -min(r["1.1.2"]["max_ulp"], r["1.2.0"].get("max_ulp", 1 << 40)),
    )[:10]
    summary["worst_samples"] = worst
    summary["seconds"] = round(time.time() - t0, 1)
    best = max(VERSIONS, key=lambda v: summary[v]["bins_exact"])
    summary["best_matching_version"] = best
    summary["passed"] = summary[best]["bins_exact"] == summary[best]["bins_total"] > 0
    a.out.parent.mkdir(parents=True, exist_ok=True)
    a.out.write_text(json.dumps(summary, indent=2, default=str) + "\n")
    print(
        json.dumps(
            {k: summary[k] for k in ("n_gtc", "best_matching_version", "passed", "seconds")}
            | {
                v: {k: summary[v][k] for k in ("samples_all_bins_exact", "bins_exact", "bins_total", "max_ulp")}
                for v in VERSIONS
            }
        )
    )
    return 0 if summary["passed"] else 1


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