code_snapshots/WiSER_public_release_code/wiser/engine/matching.py

"""Hungarian / Sinkhorn matching between K path queries and GT merged paths.

**Round 5 revert**: the default backend is **scipy** (exact Hungarian via
`scipy.optimize.linear_sum_assignment` run in a `ThreadPoolExecutor` for
parallelism). Codex R4 finding 1 established that the Round-4 Sinkhorn
default silently dropped GT pairs from supervision in ~25 % of batches
(50 / 200 random cases), violating DEC-2 / AC-3 / AC-6's ratified pure-
Hungarian contract (plan.md:39, 282, 339). scipy Hungarian always returns
the full `m_q = min(K, gt_num_paths[q])` cardinality.

Two backends remain exposed:

- `scipy` (default, production): build the full `[Q, K, M_max]` cost tensor
  on the compute device, sync once to CPU, dispatch per-sample
  `linear_sum_assignment` through a `ThreadPoolExecutor`. scipy releases
  the GIL inside its C backend, so threads give real parallelism. The
  per-triple output is guaranteed to satisfy `m_q = min(K, gt_num_paths[q])`
  on every sample.
- `sinkhorn` (optional, DIAGNOSTIC ONLY): batched log-space Sinkhorn on the
  compute device with argmax + greedy tie-break. Does **not** preserve the
  plan's full-cardinality contract — on 200 random synthetic cases
  Sinkhorn differed from scipy in 124 / 200 and under-returned pairs in
  50 / 200. Retained behind `MatcherConfig(backend="sinkhorn")` for GPU-
  util experiments ONLY; must not be used for training runs that close
  AC-3 / AC-6.

The per-triple return contract is: a list of length `Q` where each entry is
`(query_rows [m_q], gt_cols [m_q])` with `m_q = min(K, gt_num_paths[q])` under
the scipy backend. The Sinkhorn backend may return `m_q' <= m_q`; a
regression test in `tests/test_matcher_cardinality.py` enforces full
cardinality on the scipy backend.
"""

from __future__ import annotations

from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
import os

import numpy as np
import torch
from scipy.optimize import linear_sum_assignment


@dataclass(slots=True)
class MatcherConfig:
    alpha_delay_ns: float = 1.0
    beta_peak_db: float = 0.1
    gamma_exists: float = 1.0
    backend: str = "scipy"           # production default — exact Hungarian, full cardinality. "sinkhorn" is diagnostic only.
    sinkhorn_iters: int = 30
    sinkhorn_epsilon: float = 0.05    # entropy regularisation strength


_MATCH_WORKERS = max(2, min(16, (os.cpu_count() or 4)))
_MATCH_POOL: ThreadPoolExecutor | None = None


def _get_pool() -> ThreadPoolExecutor:
    global _MATCH_POOL
    if _MATCH_POOL is None:
        _MATCH_POOL = ThreadPoolExecutor(
            max_workers=_MATCH_WORKERS,
            thread_name_prefix="hungarian",
        )
    return _MATCH_POOL


def _solve_one_scipy(cost_np: np.ndarray, m: int) -> tuple[np.ndarray, np.ndarray]:
    if m == 0:
        return np.empty((0,), dtype=np.int64), np.empty((0,), dtype=np.int64)
    sub = cost_np[:, :m]
    query_rows, gt_cols = linear_sum_assignment(sub)
    return query_rows.astype(np.int64), gt_cols.astype(np.int64)


def _sinkhorn_assign_batch(
    cost: torch.Tensor,       # [Q, K, M_max]
    gt_num: torch.Tensor,     # [Q] long
    iters: int,
    epsilon: float,
) -> list[tuple[np.ndarray, np.ndarray]]:
    """Batched Sinkhorn soft-assignment → hard 1-to-1 pairing via argmax.

    For each sample q:
      - Build a square cost matrix by zero-padding to `K × K` (columns ≥ m_q
        are dummy "no-object" columns with a large cost so they are never
        selected by argmax when m_q < K).
      - Run `iters` iterations of log-space Sinkhorn on `-cost / epsilon`,
        yielding a doubly-stochastic assignment matrix `P [Q, K, K]`.
      - Take `argmax` over the GT axis; slice `[:m_q]` to get the matched
        `(query_row, gt_col)` pairs per sample. The Sinkhorn-argmax 1-to-1
        recovery is exact when query/GT costs are well-separated (the unit
        test case) and approximate when they are close, which is the
        behaviour DETR training already accepts from exact Hungarian under
        noise.
    """
    Q, K, M_max = cost.shape
    device = cost.device
    dtype = torch.float32

    # Pad / truncate to [Q, K, K] so the Sinkhorn is over a balanced bipartite
    # graph. Dummy columns (index >= m_q) carry a large cost so argmax avoids
    # them whenever a real GT path exists for that row.
    big = torch.tensor(1e6, device=device, dtype=dtype)
    if M_max < K:
        pad = big.expand(Q, K, K - M_max)
        square = torch.cat([cost.to(dtype), pad], dim=2)
    elif M_max > K:
        square = cost[:, :, :K].to(dtype)
    else:
        square = cost.to(dtype)

    # Dummy columns for samples with gt_num < K: set cost to `big` on those
    # columns so no valid query picks them.
    gt_mask = (torch.arange(K, device=device).unsqueeze(0) < gt_num.unsqueeze(1)).to(dtype)  # [Q, K]
    col_mask = gt_mask.unsqueeze(1).expand(Q, K, K)
    square = torch.where(col_mask.bool(), square, big.expand_as(square))

    # Log-space Sinkhorn on -cost / epsilon
    log_a = torch.zeros(Q, K, device=device, dtype=dtype)   # uniform row priors
    log_b = torch.zeros(Q, K, device=device, dtype=dtype)   # uniform col priors
    log_K = -square / epsilon
    log_u = torch.zeros(Q, K, device=device, dtype=dtype)
    log_v = torch.zeros(Q, K, device=device, dtype=dtype)
    for _ in range(iters):
        log_u = log_a - torch.logsumexp(log_K + log_v.unsqueeze(1), dim=-1)
        log_v = log_b - torch.logsumexp(log_K + log_u.unsqueeze(2), dim=-2)
    log_P = log_u.unsqueeze(2) + log_K + log_v.unsqueeze(1)  # [Q, K, K]
    # Argmax over GT columns → query → gt assignment
    gt_argmax = log_P.argmax(dim=-1)  # [Q, K]

    # Only a single CPU transfer here, and it's just two small int64 tensors.
    gt_argmax_cpu = gt_argmax.detach().to(torch.int64).cpu().numpy()
    gt_num_cpu = gt_num.detach().to(torch.int64).cpu().numpy()

    out: list[tuple[np.ndarray, np.ndarray]] = []
    for q in range(Q):
        m = int(gt_num_cpu[q])
        if m == 0:
            out.append((np.empty((0,), dtype=np.int64), np.empty((0,), dtype=np.int64)))
            continue
        # For each sample we need a 1-to-1 mapping between m query rows and m
        # GT columns. Take every query's best GT column, then drop queries
        # whose best GT column is a dummy (>= m) or collides with another
        # already-chosen GT column. This greedy tie-break is only used as a
        # final clean-up — the Sinkhorn matrix is already near-permutation
        # for the relevant rows.
        used_cols: set[int] = set()
        pairs: list[tuple[int, int]] = []
        # Sort queries by confidence (highest log_P entry on row) so the most
        # confident assignment goes first.
        row_max_vals = log_P[q].max(dim=-1).values
        order = torch.argsort(row_max_vals, descending=True).tolist()
        for k_idx in order:
            gt = int(gt_argmax_cpu[q, k_idx])
            if gt >= m:
                continue
            if gt in used_cols:
                continue
            used_cols.add(gt)
            pairs.append((k_idx, gt))
            if len(pairs) >= m:
                break
        if pairs:
            qr = np.array([p[0] for p in pairs], dtype=np.int64)
            gc = np.array([p[1] for p in pairs], dtype=np.int64)
        else:
            qr = np.empty((0,), dtype=np.int64)
            gc = np.empty((0,), dtype=np.int64)
        out.append((qr, gc))
    return out


def hungarian_match_batch(
    predictions: dict[str, torch.Tensor],
    gt: dict[str, torch.Tensor],
    cfg: MatcherConfig | None = None,
) -> list[tuple[np.ndarray, np.ndarray]]:
    """Batched matcher between `K` path queries and GT paths.

    Backend is picked by `cfg.backend`:
      - `"scipy"` (default, production, Round 5): exact Hungarian via
        `linear_sum_assignment` dispatched through a `ThreadPoolExecutor`.
        Guarantees the full `m_q = min(K, gt_num_paths[q])` cardinality
        contract on every sample (AC-3 / AC-6).
      - `"sinkhorn"` (DIAGNOSTIC ONLY): batched log-space Sinkhorn with
        argmax + greedy tie-break, fully on GPU. Does NOT preserve the
        full-cardinality contract — retained for GPU-util experiments only.
    """
    cfg = cfg or MatcherConfig()
    pred_exists = predictions["csi_exists_logits"]
    pred_delay = predictions["csi_delay_ns"]
    pred_peak = predictions["csi_peak_db"]
    gt_num = gt["gt_num_paths"]
    gt_delay = gt["gt_delay_ns"]
    gt_peak = gt["gt_peak_db"]

    Q, K = pred_delay.shape
    M_max = gt_delay.shape[1] if gt_delay.ndim == 2 else 0

    if M_max == 0 or Q == 0:
        return [(np.empty((0,), dtype=np.int64), np.empty((0,), dtype=np.int64)) for _ in range(Q)]

    # Shared cost tensor (fp32 on compute device).
    delay_diff = (pred_delay.unsqueeze(2) - gt_delay.unsqueeze(1)).abs()
    peak_diff = (pred_peak.unsqueeze(2) - gt_peak.unsqueeze(1)).abs()
    exists_cost = torch.nn.functional.softplus(-pred_exists).unsqueeze(2).expand(Q, K, M_max)
    cost = (
        cfg.alpha_delay_ns * delay_diff
        + cfg.beta_peak_db * peak_diff
        + cfg.gamma_exists * exists_cost
    )

    if cfg.backend == "sinkhorn":
        with torch.no_grad():
            return _sinkhorn_assign_batch(
                cost, gt_num,
                iters=cfg.sinkhorn_iters,
                epsilon=cfg.sinkhorn_epsilon,
            )

    # scipy fallback (diagnostic / regression)
    cost_np = cost.detach().to(dtype=torch.float32).cpu().numpy()
    gt_num_np = gt_num.detach().cpu().numpy().astype(np.int64)
    pool = _get_pool()
    futures = [
        pool.submit(_solve_one_scipy, cost_np[q], int(gt_num_np[q]))
        for q in range(Q)
    ]
    return [f.result() for f in futures]


__all__ = ["MatcherConfig", "hungarian_match_batch"]