1 files changed, 79 insertions, 0 deletions
diff --git a/worldalign/gw.py b/worldalign/gw.py
new file mode 100644
index 0000000..62c3f5c
--- /dev/null
+++ b/worldalign/gw.py
@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+import numpy as np
+import ot
+import torch
+from sklearn.cluster import MiniBatchKMeans
+
+from .common import pairwise_cosine_distance
+
+
+def gw_pseudo_targets(
+    vision: torch.Tensor,
+    text: torch.Tensor,
+    clusters: int,
+    seed: int,
+    max_iter: int = 100,
+) -> dict[str, torch.Tensor | float | int]:
+    """Build structure-only pseudo-correspondences between modality prototypes."""
+    x = vision.float().numpy()
+    y = text.float().numpy()
+    k = min(clusters, len(x), len(y))
+    vx = MiniBatchKMeans(
+        n_clusters=k,
+        random_state=seed,
+        batch_size=min(4096, len(x)),
+        n_init=3,
+        max_iter=200,
+    ).fit(x)
+    ty = MiniBatchKMeans(
+        n_clusters=k,
+        random_state=seed + 1,
+        batch_size=min(4096, len(y)),
+        n_init=3,
+        max_iter=200,
+    ).fit(y)
+
+    cx = pairwise_cosine_distance(vx.cluster_centers_)
+    cy = pairwise_cosine_distance(ty.cluster_centers_)
+    p = np.bincount(vx.labels_, minlength=k).astype(np.float64)
+    q = np.bincount(ty.labels_, minlength=k).astype(np.float64)
+    p /= p.sum()
+    q /= q.sum()
+
+    coupling, log = ot.gromov.gromov_wasserstein(
+        cx,
+        cy,
+        p,
+        q,
+        loss_fun="square_loss",
+        armijo=False,
+        log=True,
+        max_iter=max_iter,
+        tol_rel=1e-8,
+        tol_abs=1e-8,
+    )
+    target_centers = coupling @ ty.cluster_centers_
+    target_centers /= p[:, None].clip(min=1e-12)
+    target_centers /= np.linalg.norm(target_centers, axis=1, keepdims=True).clip(
+        min=1e-12
+    )
+    row_entropy = -np.sum(
+        (coupling / p[:, None].clip(min=1e-12))
+        * np.log(
+            (coupling / p[:, None].clip(min=1e-12)).clip(min=1e-12)
+        ),
+        axis=1,
+    ).mean()
+
+    return {
+        "vision_centers": torch.from_numpy(vx.cluster_centers_).float(),
+        "text_centers": torch.from_numpy(ty.cluster_centers_).float(),
+        "target_centers": torch.from_numpy(target_centers).float(),
+        "vision_assignments": torch.from_numpy(vx.labels_).long(),
+        "coupling": torch.from_numpy(coupling).float(),
+        "gw_distance": float(log["gw_dist"]),
+        "coupling_row_entropy": float(row_entropy),
+        "clusters": k,
+    }
+