From 6962d1b730722afe44cff0b868828808ed6d2384 Mon Sep 17 00:00:00 2001 From: vt2211 Date: Fri, 3 Jul 2026 11:29:04 -0500 Subject: [PATCH] Add MST-based component alignment for GGMP local GMMs --- fvgp/ggmp.py | 203 +++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 203 insertions(+) diff --git a/fvgp/ggmp.py b/fvgp/ggmp.py index c09a482..a07c03f 100644 --- a/fvgp/ggmp.py +++ b/fvgp/ggmp.py @@ -5,6 +5,7 @@ from scipy.stats import norm, multivariate_normal, wasserstein_distance from scipy.optimize import minimize, linear_sum_assignment from scipy.linalg import LinAlgError +from scipy.sparse.csgraph import minimum_spanning_tree from contextlib import contextmanager, nullcontext from concurrent.futures import ThreadPoolExecutor, as_completed import hashlib @@ -1350,6 +1351,208 @@ def align_local_gmms_sequence( } +def _choose_mst_root(x_data: np.ndarray) -> int: + x_data = np.asarray(x_data, dtype=float) + if x_data.ndim == 1: + x_data = x_data.reshape(-1, 1) + centroid = np.mean(x_data, axis=0, keepdims=True) + distances = np.linalg.norm(x_data - centroid, axis=1) + return int(np.argmin(distances)) + + +def build_input_mst( + x_data: np.ndarray, + *, + root: int | None = None, +) -> dict: + """ + Build a Euclidean MST over the inputs and return a deterministic BFS traversal. + + The paper describes label propagation along MST edges. This helper exposes the + tree explicitly so experiments can log the chosen traversal and verify that the + alignment matches the manuscript. + """ + x_data = np.asarray(x_data, dtype=float) + if x_data.ndim == 1: + x_data = x_data.reshape(-1, 1) + if x_data.ndim != 2: + raise ValueError("x_data must be a 2-D array") + + n = int(x_data.shape[0]) + if n == 0: + raise ValueError("x_data is empty") + + if root is None: + root = _choose_mst_root(x_data) + root = int(root) + if root < 0 or root >= n: + raise ValueError("root index out of range") + + if n == 1: + return { + "root": root, + "order": [root], + "parents": np.array([-1], dtype=int), + "adjacency": np.zeros((1, 1), dtype=float), + "edges": [], + } + + diffs = x_data[:, None, :] - x_data[None, :, :] + dist = np.linalg.norm(diffs, axis=-1) + mst = minimum_spanning_tree(dist).toarray().astype(float) + adjacency = mst + mst.T + + parents = np.full(n, -2, dtype=int) + parents[root] = -1 + order = [] + queue = [root] + while queue: + u = queue.pop(0) + order.append(int(u)) + nbrs = np.flatnonzero(adjacency[u] > 0) + nbrs = sorted( + (int(v) for v in nbrs if parents[int(v)] == -2), + key=lambda v: (float(adjacency[u, v]), int(v)), + ) + for v in nbrs: + parents[v] = int(u) + queue.append(int(v)) + + edges = [] + for child in order[1:]: + parent = int(parents[child]) + weight = float(adjacency[parent, child]) + edges.append((parent, int(child), weight)) + + return { + "root": root, + "order": order, + "parents": parents, + "adjacency": adjacency, + "edges": edges, + } + + +def align_local_gmms_mst( + x_data: np.ndarray, + weights_list: Sequence[np.ndarray], + means_list: Sequence[np.ndarray], + covs_list: Sequence[np.ndarray], + *, + metric: str = "w2", + root: int | None = None, +) -> dict: + """ + Align local GMM components by propagating labels along the input MST. + + This implements the geometry-aware matching procedure described in the paper: + build a Euclidean MST on the inputs, then solve one Hungarian assignment per + tree edge using squared Gaussian W2 cost. + """ + if not (len(weights_list) == len(means_list) == len(covs_list)): + raise ValueError("weights_list, means_list, covs_list must have equal length") + + n = int(len(means_list)) + if n == 0: + raise ValueError("Empty sequence") + + x_data = np.asarray(x_data, dtype=float) + if x_data.ndim == 1: + x_data = x_data.reshape(-1, 1) + if x_data.shape[0] != n: + raise ValueError("x_data length must match number of local GMMs") + + mst = build_input_mst(x_data, root=root) + order = [int(v) for v in mst["order"]] + parents = np.asarray(mst["parents"], dtype=int) + + aligned_w = [None] * n + aligned_m = [None] * n + aligned_c = [None] * n + perms = [None] * n + costs = [None] * n + + root_idx = int(mst["root"]) + aligned_w[root_idx] = np.asarray(weights_list[root_idx], dtype=float).reshape(-1).copy() + aligned_m[root_idx] = np.asarray(means_list[root_idx], dtype=float).copy() + aligned_c[root_idx] = np.asarray(covs_list[root_idx], dtype=float).copy() + perms[root_idx] = np.arange(aligned_m[root_idx].shape[0], dtype=int) + + for child in order[1:]: + parent = int(parents[child]) + m_ref = np.asarray(aligned_m[parent], dtype=float) + c_ref = np.asarray(aligned_c[parent], dtype=float) + m_cur = np.asarray(means_list[child], dtype=float) + c_cur = np.asarray(covs_list[child], dtype=float) + w_cur = np.asarray(weights_list[child], dtype=float).reshape(-1) + + perm, cost = align_gmm_components_hungarian( + m_ref, + c_ref, + m_cur, + c_cur, + metric=str(metric), + return_cost=True, + ) + aligned_w[child] = w_cur[perm].copy() + aligned_m[child] = m_cur[perm].copy() + aligned_c[child] = c_cur[perm].copy() + perms[child] = perm.copy() + costs[child] = cost.copy() + + return { + "weights": aligned_w, + "means": aligned_m, + "covs": aligned_c, + "perms": perms, + "costs": costs, + "metric": str(metric), + "method": "mst", + "root": root_idx, + "mst_order": order, + "mst_parents": parents, + "mst_edges": mst["edges"], + "mst_adjacency": mst["adjacency"], + } + + +def align_local_gmms( + weights_list: Sequence[np.ndarray], + means_list: Sequence[np.ndarray], + covs_list: Sequence[np.ndarray], + *, + x_data: np.ndarray | None = None, + metric: str = "w2", + method: str = "sequence", + reference: str = "previous", + root: int | None = None, +) -> dict: + """ + Dispatch between the legacy sequence aligner and the manuscript MST aligner. + """ + method_key = str(method).lower() + if method_key == "sequence": + return align_local_gmms_sequence( + weights_list, + means_list, + covs_list, + metric=str(metric), + reference=str(reference), + ) + if method_key == "mst": + if x_data is None: + raise ValueError("x_data is required for method='mst'") + return align_local_gmms_mst( + x_data, + weights_list, + means_list, + covs_list, + metric=str(metric), + root=root, + ) + raise ValueError("method must be 'sequence' or 'mst'") + + def _log_mvn_density( y: np.ndarray, mean: np.ndarray,