"""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"]