added fixed-point barycenter function to ot.lp._barycenter_solvers_

Author: eloitanguy

CONTRIBUTORS.md

@@ -44,7 +44,7 @@ The contributors to this library are:
 * [Cédric Vincent-Cuaz](https://github.com/cedricvincentcuaz) (Graph Dictionary Learning, FGW,
   semi-relaxed FGW, quantized FGW, partial FGW)
 * [Eloi Tanguy](https://github.com/eloitanguy) (Generalized Wasserstein
-  Barycenters, GMMOT)
+  Barycenters, GMMOT, Barycenters for General Transport Costs)
 * [Camille Le Coz](https://www.linkedin.com/in/camille-le-coz-8593b91a1/) (EMD2 debug)
 * [Eduardo Fernandes Montesuma](https://eddardd.github.io/my-personal-blog/) (Free support sinkhorn barycenter)
 * [Theo Gnassounou](https://github.com/tgnassou) (OT between Gaussian distributions)

README.md

@@ -391,3 +391,7 @@ Artificial Intelligence.
 [72] Thibault Séjourné, François-Xavier Vialard, and Gabriel Peyré (2021). [The Unbalanced Gromov Wasserstein Distance: Conic Formulation and Relaxation](https://proceedings.neurips.cc/paper/2021/file/4990974d150d0de5e6e15a1454fe6b0f-Paper.pdf). Neural Information Processing Systems (NeurIPS).
 
 [73] Séjourné, T., Vialard, F. X., & Peyré, G. (2022). [Faster Unbalanced Optimal Transport: Translation Invariant Sinkhorn and 1-D Frank-Wolfe](https://proceedings.mlr.press/v151/sejourne22a.html). In International Conference on Artificial Intelligence and Statistics (pp. 4995-5021). PMLR.
+
+[74] Tanguy, Eloi and Delon, Julie and Gozlan, Nathaël (2024). [Computing
+Barycentres of Measures for Generic Transport
+Costs](https://arxiv.org/abs/2501.04016). arXiv preprint 2501.04016 (2024)

RELEASES.md

@@ -7,6 +7,8 @@
 - Added feature `grad=last_step` for `ot.solvers.solve` (PR #693)
 - Automatic PR labeling and release file update check (PR #704)
 - Reorganize sub-module `ot/lp/__init__.py` into separate files (PR #714)
+- Implement fixed-point solver for OT barycenters with generic cost functions
+  (generalizes `ot.lp.free_support_barycenter`). (PR #715)
 
 #### Closed issues
 - Fixed `ot.mapping` solvers which depended on deprecated `cvxpy` `ECOS` solver (PR #692, Issue #668)

ot/lp/_barycenter_solvers.py

@@ -422,3 +422,90 @@ def generalized_free_support_barycenter(
         return Y, log_dict
     else:
         return Y
+
+
+class StoppingCriterionReached(Exception):
+    pass
+
+
+def solve_OT_barycenter_fixed_point(
+    X_init,
+    Y_list,
+    b_list,
+    cost_list,
+    B,
+    max_its=5,
+    stop_threshold=1e-5,
+    log=False,
+):
+    """
+    Solves the OT barycenter problem using the fixed point algorithm, iterating
+    the function B on plans between the current barycentre and the measures.
+
+    Parameters
+    ----------
+    X_init : array-like
+        Array of shape (n, d) representing initial barycentre points.
+    Y_list : list of array-like
+        List of K arrays of measure positions, each of shape (m_k, d_k).
+    b_list : list of array-like
+        List of K arrays of measure weights, each of shape (m_k).
+    cost_list : list of callable
+        List of K cost functions R^(n, d) x R^(m_k, d_k) -> R_+^(n, m_k).
+    B : callable
+        Function from R^d_1 x ... x R^d_K to R^d accepting a list of K arrays of shape (n, d_K), computing the ground barycentre.
+    max_its : int, optional
+        Maximum number of iterations (default is 5).
+    stop_threshold : float, optional
+        If the iterations move less than this, terminate (default is 1e-5).
+    log : bool, optional
+        Whether to return the log dictionary (default is False).
+
+    Returns
+    -------
+    X : array-like
+        Array of shape (n, d) representing barycentre points.
+    log_dict : list of array-like, optional
+        log containing the exit status, list of iterations and list of
+        displacements if log is True.
+    """
+    nx = get_backend(X_init, Y_list[0])
+    K = len(Y_list)
+    n = X_init.shape[0]
+    a = nx.ones(n) / n
+    X_list = [X_init] if log else []  # store the iterations
+    X = X_init
+    dX_list = []  # store the displacement squared norms
+    exit_status = "Unknown"
+
+    try:
+        for _ in range(max_its):
+            pi_list = [  # compute the pairwise transport plans
+                emd(a, b_list[k], cost_list[k](X, Y_list[k])) for k in range(K)
+            ]
+            Y_perm = []
+            for k in range(K):  # compute barycentric projections
+                Y_perm.append(n * pi_list[k] @ Y_list[k])
+            X_next = B(Y_perm)
+
+            if log:
+                X_list.append(X_next)
+
+            # stationary criterion: move less than the threshold
+            dX = nx.sum((X - X_next) ** 2)
+            X = X_next
+
+            if log:
+                dX_list.append(dX)
+
+            if dX < stop_threshold:
+                exit_status = "Stationary Point"
+                raise StoppingCriterionReached
+
+        exit_status = "Max iterations reached"
+        raise StoppingCriterionReached
+
+    except StoppingCriterionReached:
+        if log:
+            return X, {"X_list": X_list, "exit_status": exit_status, "dX_list": dX_list}
+        return X