[WIP] OT barycenters for generic transport costs

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

@@ -54,6 +54,8 @@ POT provides the following generic OT solvers (links to examples):
 * [Co-Optimal Transport](https://pythonot.github.io/auto_examples/others/plot_COOT.html) [49] and
 [unbalanced Co-Optimal Transport](https://pythonot.github.io/auto_examples/others/plot_learning_weights_with_COOT.html) [71].
 * Fused unbalanced Gromov-Wasserstein [70].
+* [Optimal Transport Barycenters for Generic Costs](https://pythonot.github.io/auto_examples/barycenters/plot_free_support_barycenter_generic_cost.html) [76]
+* [Barycenters between Gaussian Mixture Models](https://pythonot.github.io/auto_examples/barycenters/plot_gmm_barycenter.html) [69, 76]
 
 POT provides the following Machine Learning related solvers:
 
@@ -394,3 +396,5 @@ Artificial Intelligence.
 [74] Chewi, S., Maunu, T., Rigollet, P., & Stromme, A. J. (2020). [Gradient descent algorithms for Bures-Wasserstein barycenters](https://proceedings.mlr.press/v125/chewi20a.html). In Conference on Learning Theory (pp. 1276-1304). PMLR.
 
 [75] Altschuler, J., Chewi, S., Gerber, P. R., & Stromme, A. (2021). [Averaging on the Bures-Wasserstein manifold: dimension-free convergence of gradient descent](https://papers.neurips.cc/paper_files/paper/2021/hash/b9acb4ae6121c941324b2b1d3fac5c30-Abstract.html). Advances in Neural Information Processing Systems, 34, 22132-22145.
+
+[76] 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,9 @@
 - 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`), with example. (PR #715)
+- Implement fixed-point solver for barycenters between GMMs (PR #715), with example.
 - Fix warning raise when import the library (PR #716)
 - Implement projected gradient descent solvers for entropic partial FGW (PR #702)
 - Fix documentation in the module `ot.gaussian` (PR #718)

examples/barycenters/plot_free_support_barycenter_generic_cost.py

@@ -0,0 +1,177 @@
+# -*- coding: utf-8 -*-
+"""
+=====================================
+OT Barycenter with Generic Costs Demo
+=====================================
+
+This example illustrates the computation of an Optimal Transport Barycenter for
+a ground cost that is not a power of a norm. We take the example of ground costs
+:math:`c_k(x, y) = \|P_k(x)-y\|_2^2`, where :math:`P_k` is the (non-linear)
+projection onto a circle k. This is an example of the fixed-point barycenter
+solver introduced in [76] which generalises [20] and [43].
+
+The ground barycenter function :math:`B(y_1, ..., y_K) = \mathrm{argmin}_{x \in
+\mathbb{R}^2} \sum_k \lambda_k c_k(x, y_k)` is computed by gradient descent over
+:math:`x` with Pytorch.
+
+[76] Tanguy, Eloi and Delon, Julie and Gozlan, Nathaël (2024). Computing
+Barycentres of Measures for Generic Transport Costs. arXiv preprint 2501.04016
+(2024)
+
+[20] Cuturi, M. and Doucet, A. (2014) Fast Computation of Wasserstein
+Barycenters. InternationalConference in Machine Learning
+
+[43] Álvarez-Esteban, Pedro C., et al. A fixed-point approach to barycenters in
+Wasserstein space. Journal of Mathematical Analysis and Applications 441.2
+(2016): 744-762.
+
+"""
+
+# Author: Eloi Tanguy <[email protected]>
+#
+# License: MIT License
+
+# sphinx_gallery_thumbnail_number = 1
+
+# %%
+# Generate data
+import torch
+from torch.optim import Adam
+from ot.utils import dist
+import numpy as np
+from ot.lp import free_support_barycenter_generic_costs
+import matplotlib.pyplot as plt
+
+
+torch.manual_seed(42)
+
+n = 200  # number of points of the of the barycentre
+d = 2  # dimensions of the original measure
+K = 4  # number of measures to barycentre
+m = 50  # number of points of the measures
+b_list = [torch.ones(m) / m] * K  # weights of the 4 measures
+weights = torch.ones(K) / K  # weights for the barycentre
+stop_threshold = 1e-20  # stop threshold for B and for fixed-point algo
+
+
+# map R^2 -> R^2 projection onto circle
+def proj_circle(X, origin, radius):
+    diffs = X - origin[None, :]
+    norms = torch.norm(diffs, dim=1)
+    return origin[None, :] + radius * diffs / norms[:, None]
+
+
+# circles on which to project
+origin1 = torch.tensor([-1.0, -1.0])
+origin2 = torch.tensor([-1.0, 2.0])
+origin3 = torch.tensor([2.0, 2.0])
+origin4 = torch.tensor([2.0, -1.0])
+r = np.sqrt(2)
+P_list = [
+    lambda X: proj_circle(X, origin1, r),
+    lambda X: proj_circle(X, origin2, r),
+    lambda X: proj_circle(X, origin3, r),
+    lambda X: proj_circle(X, origin4, r),
+]
+
+# measures to barycentre are projections of different random circles
+# onto the K circles
+Y_list = []
+for k in range(K):
+    t = torch.rand(m) * 2 * np.pi
+    X_temp = 0.5 * torch.stack([torch.cos(t), torch.sin(t)], axis=1)
+    X_temp = X_temp + torch.tensor([0.5, 0.5])[None, :]
+    Y_list.append(P_list[k](X_temp))
+
+
+# %%
+# Define costs and ground barycenter function
+# cost_list[k] is a function taking x (n, d) and y (n_k, d_k) and returning a
+# (n, n_k) matrix of costs
+def c1(x, y):
+    return dist(P_list[0](x), y)
+
+
+def c2(x, y):
+    return dist(P_list[1](x), y)
+
+
+def c3(x, y):
+    return dist(P_list[2](x), y)
+
+
+def c4(x, y):
+    return dist(P_list[3](x), y)
+
+
+cost_list = [c1, c2, c3, c4]
+
+
+# batched total ground cost function for candidate points x (n, d)
+# for computation of the ground barycenter B with gradient descent
+def C(x, y):
+    """
+    Computes the barycenter cost for candidate points x (n, d) and
+    measure supports y: List(n, d_k).
+    """
+    n = x.shape[0]
+    K = len(y)
+    out = torch.zeros(n)
+    for k in range(K):
+        out += (1 / K) * torch.sum((P_list[k](x) - y[k]) ** 2, axis=1)
+    return out
+
+
+# ground barycenter function
+def B(y, its=150, lr=1, stop_threshold=stop_threshold):
+    """
+    Computes the ground barycenter for measure supports y: List(n, d_k).
+    Output: (n, d) array
+    """
+    x = torch.randn(n, d)
+    x.requires_grad_(True)
+    opt = Adam([x], lr=lr)
+    for _ in range(its):
+        x_prev = x.data.clone()
+        opt.zero_grad()
+        loss = torch.sum(C(x, y))
+        loss.backward()
+        opt.step()
+        diff = torch.sum((x.data - x_prev) ** 2)
+        if diff < stop_threshold:
+            break
+    return x
+
+
+# %%
+# Compute the barycenter measure
+fixed_point_its = 3
+X_init = torch.rand(n, d)
+X_bar = free_support_barycenter_generic_costs(
+    Y_list,
+    b_list,
+    X_init,
+    cost_list,
+    B,
+    numItermax=fixed_point_its,
+    stopThr=stop_threshold,
+)
+
+# %%
+# Plot Barycenter (Iteration 3)
+alpha = 0.4
+s = 80
+labels = ["circle 1", "circle 2", "circle 3", "circle 4"]
+for Y, label in zip(Y_list, labels):
+    plt.scatter(*(Y.numpy()).T, alpha=alpha, label=label, s=s)
+plt.scatter(
+    *(X_bar.detach().numpy()).T, label="Barycenter", c="black", alpha=alpha, s=s
+)
+plt.axis("equal")
+plt.xlim(-0.3, 1.3)
+plt.ylim(-0.3, 1.3)
+plt.axis("off")
+plt.legend()
+plt.tight_layout()
+
+# %%

examples/barycenters/plot_generalized_free_support_barycenter.py

@@ -14,7 +14,7 @@
 
 """
 
-# Author: Eloi Tanguy <eloi.tanguy@polytechnique.edu>
+# Author: Eloi Tanguy <eloi.tanguy@math.cnrs.fr>
 #
 # License: MIT License
 

examples/barycenters/plot_gmm_barycenter.py

@@ -0,0 +1,142 @@
+# -*- coding: utf-8 -*-
+"""
+=====================================
+Gaussian Mixture Model OT Barycenters
+=====================================
+
+This example illustrates the computation of a barycenter between Gaussian
+Mixtures in the sense of GMM-OT [69]. This computation is done using the
+fixed-point method for OT barycenters with generic costs [76], for which POT
+provides a general solver, and a specific GMM solver. Note that this is a
+'free-support' method, implying that the number of components of the barycenter
+GMM and their weights are fixed.
+
+The idea behind GMM-OT barycenters is to see the GMMs as discrete measures over
+the space of Gaussian distributions :math:`\mathcal{N}` (or equivalently the
+Bures-Wasserstein manifold), and to compute barycenters with respect to the
+2-Wasserstein distance between measures in :math:`\mathcal{P}(\mathcal{N})`: a
+gaussian mixture is a finite combination of Diracs on specific gaussians, and
+two mixtures are compared with the 2-Wasserstein distance on this space, where
+ground cost the squared Bures distance between gaussians.
+
+[69] Delon, J., & Desolneux, A. (2020). A Wasserstein-type distance in the space
+of Gaussian mixture models. SIAM Journal on Imaging Sciences, 13(2), 936-970.
+
+[76] Tanguy, Eloi and Delon, Julie and Gozlan, Nathaël (2024). Computing
+Barycentres of Measures for Generic Transport Costs. arXiv preprint 2501.04016
+(2024)
+
+"""
+
+# Author: Eloi Tanguy <[email protected]>
+#
+# License: MIT License
+
+# sphinx_gallery_thumbnail_number = 1
+
+# %%
+# Generate data
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.patches import Ellipse
+import ot
+from ot.gmm import gmm_barycenter_fixed_point
+
+
+K = 3  # number of GMMs
+d = 2  # dimension
+n = 6  # number of components of the desired barycenter
+
+
+def get_random_gmm(K, d, seed=0, min_cov_eig=1, cov_scale=1e-2):
+    rng = np.random.RandomState(seed=seed)
+    means = rng.randn(K, d)
+    P = rng.randn(K, d, d) * cov_scale
+    # C[k] = P[k] @ P[k]^T + min_cov_eig * I
+    covariances = np.einsum("kab,kcb->kac", P, P)
+    covariances += min_cov_eig * np.array([np.eye(d) for _ in range(K)])
+    weights = rng.random(K)
+    weights /= np.sum(weights)
+    return means, covariances, weights
+
+
+m_list = [5, 6, 7]  # number of components in each GMM
+offsets = [np.array([-3, 0]), np.array([2, 0]), np.array([0, 4])]
+means_list = []  # list of means for each GMM
+covs_list = []  # list of covariances for each GMM
+w_list = []  # list of weights for each GMM
+
+# generate GMMs
+for k in range(K):
+    means, covs, b = get_random_gmm(
+        m_list[k], d, seed=k, min_cov_eig=0.25, cov_scale=0.5
+    )
+    means = means / 2 + offsets[k][None, :]
+    means_list.append(means)
+    covs_list.append(covs)
+    w_list.append(b)
+
+# %%
+# Compute the barycenter using the fixed-point method
+init_means, init_covs, _ = get_random_gmm(n, d, seed=0)
+weights = ot.unif(K)  # barycenter coefficients
+means_bar, covs_bar, log = gmm_barycenter_fixed_point(
+    means_list,
+    covs_list,
+    w_list,
+    init_means,
+    init_covs,
+    weights,
+    iterations=3,
+    log=True,
+)
+
+
+# %%
+# Define plotting functions
+
+
+# draw a covariance ellipse
+def draw_cov(mu, C, color=None, label=None, nstd=1, alpha=0.5, ax=None):
+    def eigsorted(cov):
+        vals, vecs = np.linalg.eigh(cov)
+        order = vals.argsort()[::-1].copy()
+        return vals[order], vecs[:, order]
+
+    vals, vecs = eigsorted(C)
+    theta = np.degrees(np.arctan2(*vecs[:, 0][::-1]))
+    w, h = 2 * nstd * np.sqrt(vals)
+    ell = Ellipse(
+        xy=(mu[0], mu[1]),
+        width=w,
+        height=h,
+        alpha=alpha,
+        angle=theta,
+        facecolor=color,
+        edgecolor=color,
+        label=label,
+        fill=True,
+    )
+    if ax is None:
+        ax = plt.gca()
+    ax.add_artist(ell)
+
+
+# draw a gmm as a set of ellipses with weights shown in alpha value
+def draw_gmm(ms, Cs, ws, color=None, nstd=0.5, alpha=1, label=None, ax=None):
+    for k in range(ms.shape[0]):
+        draw_cov(
+            ms[k], Cs[k], color, label if k == 0 else None, nstd, alpha * ws[k], ax=ax
+        )
+
+
+# %%
+# Plot the results
+fig, ax = plt.subplots(figsize=(6, 6))
+axis = [-4, 4, -2, 6]
+ax.set_title("Fixed Point Barycenter (3 Iterations)", fontsize=16)
+for k in range(K):
+    draw_gmm(means_list[k], covs_list[k], w_list[k], color="C0", ax=ax)
+draw_gmm(means_bar, covs_bar, ot.unif(n), color="C1", ax=ax)
+ax.axis(axis)
+ax.axis("off")

examples/others/plot_GMMOT_plan.py

@@ -16,7 +16,7 @@
 
 """
 
-# Author: Eloi Tanguy <eloi.tanguy@u-paris>
+# Author: Eloi Tanguy <eloi.tanguy@math.cnrs.fr>
 #         Remi Flamary <[email protected]>
 #         Julie Delon <[email protected]>
 #

examples/others/plot_GMM_flow.py

@@ -10,7 +10,7 @@
 
 """
 
-# Author: Eloi Tanguy <eloi.tanguy@u-paris>
+# Author: Eloi Tanguy <eloi.tanguy@math.cnrs.fr>
 #         Remi Flamary <[email protected]>
 #         Julie Delon <[email protected]>
 #

examples/others/plot_SSNB.py

@@ -38,7 +38,7 @@
         2017.
 """
 
-# Author: Eloi Tanguy <eloi.tanguy@u-paris.fr>
+# Author: Eloi Tanguy <eloi.tanguy@math.cnrs.fr>
 # License: MIT License
 
 # sphinx_gallery_thumbnail_number = 3