1 files changed, 201 insertions, 0 deletions
diff --git a/worldalign/match_battery.py b/worldalign/match_battery.py
index 45fa26e..9035e62 100644
--- a/worldalign/match_battery.py
+++ b/worldalign/match_battery.py
@@ -123,12 +123,213 @@ def solve_bapg(visual, shuffled, **_):
     return coupling_to_permutation(coupling)
 
 
+def solve_faq(visual, shuffled, **_):
+    """Frank-Wolfe on the Birkhoff polytope: the canonical QAP relaxation.
+
+    This is the convex-concave family the register listed and never reached.
+    It optimises over doubly stochastic matrices and projects at the end,
+    which is a different relaxation from both the spectral solvers and GW.
+    """
+    from scipy.optimize import quadratic_assignment
+
+    result = quadratic_assignment(
+        visual, shuffled, method="faq", options={"maximize": True, "n_init": 1}
+    )
+    return result.col_ind
+
+
+def solve_faq_multi(visual, shuffled, **_):
+    """FAQ from many random doubly stochastic starts, best by its own objective."""
+    from scipy.optimize import quadratic_assignment
+
+    result = quadratic_assignment(
+        visual, shuffled, method="faq",
+        options={"maximize": True, "n_init": 30, "rng": 0},
+    )
+    return result.col_ind
+
+
+def solve_2opt(visual, shuffled, **_):
+    from scipy.optimize import quadratic_assignment
+
+    result = quadratic_assignment(
+        visual, shuffled, method="2opt", options={"maximize": True, "rng": 0}
+    )
+    return result.col_ind
+
+
+def solve_gw_multi(visual, shuffled, **_):
+    """GW is non-convex; restart it and keep the best by the GW objective."""
+    import ot
+
+    size = len(visual)
+    weights = ot.unif(size)
+    best, best_value = None, np.inf
+    generator = np.random.default_rng(0)
+    for attempt in range(12):
+        init = None
+        if attempt:
+            noise = generator.random((size, size)) + 1e-3
+            init = noise / noise.sum()
+        try:
+            coupling, log = ot.gromov.gromov_wasserstein(
+                visual, shuffled, weights, weights, "square_loss",
+                G0=init, max_iter=200, tol_rel=1e-9, log=True,
+            )
+        except Exception:
+            continue
+        value = float(log["gw_dist"])
+        if value < best_value:
+            best, best_value = coupling, value
+    if best is None:
+        raise RuntimeError("all GW restarts failed")
+    return coupling_to_permutation(best)
+
+
+def solve_gw_kl(visual, shuffled, **_):
+    import ot
+
+    size = len(visual)
+    weights = ot.unif(size)
+    shift = min(visual.min(), shuffled.min())
+    coupling = ot.gromov.gromov_wasserstein(
+        visual - shift + 1e-3, shuffled - shift + 1e-3, weights, weights,
+        "kl_loss", max_iter=200, tol_rel=1e-9,
+    )
+    return coupling_to_permutation(coupling)
+
+
+def solve_gw_deep_anneal(visual, shuffled, **_):
+    """The winning solver, annealed further and started from the smooth end."""
+    import ot
+
+    size = len(visual)
+    weights = ot.unif(size)
+    coupling = None
+    schedule = (2e-1, 1e-1, 5e-2, 3e-2, 2e-2, 1e-2, 7e-3, 5e-3, 3e-3, 2e-3, 1e-3)
+    for epsilon in schedule:
+        try:
+            candidate = ot.gromov.entropic_gromov_wasserstein(
+                visual, shuffled, weights, weights, "square_loss",
+                epsilon=epsilon, G0=coupling, max_iter=300, tol=1e-9, solver="PGD",
+            )
+        except Exception:
+            break
+        if not np.isfinite(candidate).all():
+            break
+        coupling = candidate
+    if coupling is None:
+        raise RuntimeError("deep anneal diverged immediately")
+    return coupling_to_permutation(coupling)
+
+
+def solve_semirelaxed(visual, shuffled, **_):
+    import ot
+
+    size = len(visual)
+    coupling = ot.gromov.semirelaxed_gromov_wasserstein(
+        visual, shuffled, ot.unif(size), "square_loss", max_iter=200,
+    )
+    return coupling_to_permutation(coupling)
+
+
+def _birkhoff_frank_wolfe(visual, shuffled, mu, coupling, steps):
+    """Frank-Wolfe on the Birkhoff polytope for one concavity setting."""
+    from scipy.optimize import linear_sum_assignment
+
+    size = len(visual)
+    for step in range(steps):
+        gradient = 2.0 * (visual @ coupling @ shuffled) + 2.0 * mu * coupling
+        _, columns = linear_sum_assignment(-gradient)
+        vertex = np.zeros_like(coupling)
+        vertex[np.arange(size), columns] = 1.0
+        direction = vertex - coupling
+        # exact line search on the quadratic along the segment
+        quad = float((direction * (visual @ direction @ shuffled)).sum()) + mu * float(
+            (direction * direction).sum()
+        )
+        linear = float((direction * gradient).sum())
+        if abs(quad) < 1e-12:
+            gamma = 1.0 if linear > 0 else 0.0
+        else:
+            gamma = float(np.clip(-linear / (2.0 * quad), 0.0, 1.0)) if quad < 0 else (
+                1.0 if linear > 0 else 0.0
+            )
+        if gamma <= 1e-12:
+            break
+        coupling = coupling + gamma * direction
+    return coupling
+
+
+def solve_path(visual, shuffled, **_):
+    """Convex-to-concave path following over the Birkhoff polytope.
+
+    The relaxation the register named and never ran. Maximising a concave
+    surrogate over the polytope has one smooth interior optimum; maximising a
+    convex one attains its maximum at a vertex, which is a permutation. Adding
+    mu*||P||^2 and sweeping mu from negative to positive walks continuously
+    between the two, and the argument for it is exactly our situation: the
+    smooth end has no local optima to be trapped by, and the solution is
+    tracked into the sharp end rather than searched for there.
+    """
+    size = len(visual)
+    scale = float(np.abs(np.linalg.eigvalsh(visual)).max()
+                  * np.abs(np.linalg.eigvalsh(shuffled)).max())
+    coupling = np.full((size, size), 1.0 / size)
+    for fraction in np.linspace(-1.0, 1.0, 21):
+        coupling = _birkhoff_frank_wolfe(
+            visual, shuffled, fraction * scale, coupling, 30
+        )
+    return coupling_to_permutation(coupling)
+
+
+def solve_moment_ladder(visual, shuffled, **_):
+    """Match the eigenvalue moments of M, not just its sum.
+
+    The energy is a moment functional in disguise: the pairwise term is the
+    first moment of M = permuted-text * visual and the triangle term its
+    third. The sequence {tr(M^k)} is what pins down M's spectrum, and we have
+    only ever used two of its entries. This optimises a weighted combination
+    over several k by Frank-Wolfe, which is the cheapest way to ask whether a
+    longer moment sequence has a wider basin than a single term.
+    """
+    from scipy.optimize import linear_sum_assignment
+
+    size = len(visual)
+    coupling = np.full((size, size), 1.0 / size)
+    for step in range(120):
+        permuted = coupling @ shuffled @ coupling.T
+        product = permuted * visual
+        # d/dP of tr(M^k) terms, accumulated over the ladder
+        gradient = 2.0 * (visual @ coupling @ shuffled)
+        squared = product @ product
+        gradient = gradient + 1.0 * (
+            (visual * product) @ coupling @ shuffled
+            + (visual * squared) @ coupling @ shuffled
+        )
+        _, columns = linear_sum_assignment(-gradient)
+        vertex = np.zeros_like(coupling)
+        vertex[np.arange(size), columns] = 1.0
+        gamma = 2.0 / (step + 2.0)
+        coupling = coupling + gamma * (vertex - coupling)
+    return coupling_to_permutation(coupling)
+
+
 SOLVERS = {
     "grampa (current)": solve_grampa,
     "umeyama": solve_umeyama,
     "gromov-wasserstein": solve_gw,
     "entropic GW annealed": solve_gw_annealed,
     "BAPG GW": solve_bapg,
+    "FAQ (Frank-Wolfe)": solve_faq,
+    "FAQ x30 restarts": solve_faq_multi,
+    "2-opt": solve_2opt,
+    "GW x12 restarts": solve_gw_multi,
+    "GW kl-loss": solve_gw_kl,
+    "entropic GW deep anneal": solve_gw_deep_anneal,
+    "semirelaxed GW": solve_semirelaxed,
+    "PATH convex-concave": solve_path,
+    "moment ladder tr(M^k)": solve_moment_ladder,
 }