[MRG] Efficient Discrete Multi Marginal Optimal Transport

CONTRIBUTORS.md

@@ -42,6 +42,8 @@ The contributors to this library are:
 * [Eduardo Fernandes Montesuma](https://eddardd.github.io/my-personal-blog/) (Free support sinkhorn barycenter)
 * [Theo Gnassounou](https://github.com/tgnassou) (OT between Gaussian distributions)
 * [Clément Bonet](https://clbonet.github.io) (Wassertstein on circle, Spherical Sliced-Wasserstein)
+* [Ronak Mehta](https://ronakrm.github.io) (Efficient Discrete Multi Marginal Optimal Transport Regularization)
+* [Xizheng Yu](https://github.com/x12hengyu) (Efficient Discrete Multi Marginal Optimal Transport Regularization)
 
 ## Acknowledgments
 

README.md

@@ -43,6 +43,7 @@ POT provides the following generic OT solvers (links to examples):
 * [Spherical Sliced Wasserstein](https://pythonot.github.io/auto_examples/sliced-wasserstein/plot_variance_ssw.html) [46]
 * [Graph Dictionary Learning solvers](https://pythonot.github.io/auto_examples/gromov/plot_gromov_wasserstein_dictionary_learning.html) [38].
 * [Semi-relaxed (Fused) Gromov-Wasserstein divergences](https://pythonot.github.io/auto_examples/gromov/plot_semirelaxed_fgw.html) (exact and regularized [48]).
+* [Efficient Discrete Multi Marginal Optimal Transport Regularization](https://pythonot.github.io/auto_examples/others/plot_demd_gradient_minimize.html) [50].
 * [Several backends](https://pythonot.github.io/quickstart.html#solving-ot-with-multiple-backends) for easy use of POT with  [Pytorch](https://pytorch.org/)/[jax](https://github.com/google/jax)/[Numpy](https://numpy.org/)/[Cupy](https://cupy.dev/)/[Tensorflow](https://www.tensorflow.org/) arrays.
 
 POT provides the following Machine Learning related solvers:
@@ -312,3 +313,8 @@ Dictionary Learning](https://arxiv.org/pdf/2102.06555.pdf), International Confer
 [50] Liu, T., Puigcerver, J., & Blondel, M. (2023). [Sparsity-constrained optimal transport](https://openreview.net/forum?id=yHY9NbQJ5BP). Proceedings of the Eleventh International Conference on Learning Representations (ICLR).
 
 [51] Xu, H., Luo, D., Zha, H., & Duke, L. C. (2019). [Gromov-wasserstein learning for graph matching and node embedding](http://proceedings.mlr.press/v97/xu19b.html). In International Conference on Machine Learning (ICML), 2019.
+
+[52] Ronak Mehta, Jeffery Kline, Vishnu Suresh Lokhande, Glenn Fung, & Vikas Singh (2023). [Efficient Discrete Multi Marginal Optimal Transport Regularization](https://openreview.net/forum?id=R98ZfMt-jE). In The Eleventh International Conference on Learning Representations (ICLR).
+
+[53] Jeffery Kline. [Properties of the d-dimensional earth mover’s problem](https://www.sciencedirect.com/science/article/pii/S0166218X19301441). Discrete Applied Mathematics, 265: 128–141, 2019.
+

RELEASES.md

@@ -16,6 +16,8 @@
 - Added entropic semi-relaxed (Fused) Gromov-Wasserstein solvers in `ot.gromov` + examples (PR #455)
 - Make marginal parameters optional for (F)GW solvers in `._gw`, `._bregman` and `._semirelaxed` (PR #455)
 
+- Added feature Efficient Discrete Multi Marginal Optimal Transport Regularization + examples (PR #454)
+
 #### Closed issues
 - Fix circleci-redirector action and codecov (PR #460)
 - Fix issues with cuda for ot.binary_search_circle and with gradients for ot.sliced_wasserstein_sphere (PR #457)

examples/others/plot_d-mmot.py

@@ -0,0 +1,150 @@
+# -*- coding: utf-8 -*-
+r"""
+===============================================================================
+Computing d-dimensional Barycenters via d-MMOT
+===============================================================================
+
+When the cost is discretized (Monge), the d-MMOT solver can more quickly
+compute and minimize the distance between many distributions without the need
+for intermediate barycenter computations. This example compares the time to
+identify, and the quality of, solutions for the d-MMOT problem using a
+primal/dual algorithm and classical LP barycenter approaches.
+"""
+
+# Author: Ronak Mehta <[email protected]>
+#         Xizheng Yu <[email protected]>
+#
+# License: MIT License
+
+# %%
+# Generating 2 distributions
+# -----
+import numpy as np
+import matplotlib.pyplot as pl
+import ot
+
+np.random.seed(0)
+
+n = 100
+d = 2
+# Gaussian distributions
+a1 = ot.datasets.make_1D_gauss(n, m=20, s=5)  # m=mean, s=std
+a2 = ot.datasets.make_1D_gauss(n, m=60, s=8)
+A = np.vstack((a1, a2)).T
+x = np.arange(n, dtype=np.float64)
+M = ot.utils.dist(x.reshape((n, 1)), metric='minkowski')
+
+pl.figure(1, figsize=(6.4, 3))
+pl.plot(x, a1, 'b', label='Source distribution')
+pl.plot(x, a2, 'r', label='Target distribution')
+pl.legend()
+
+# %%
+# Minimize the distances among distributions, identify the Barycenter
+# -----
+# The objective being minimized is different for both methods, so the objective
+# values cannot be compared.
+
+print('LP Iterations:')
+ot.tic()
+alpha = 1  # /d  # 0<=alpha<=1
+weights = np.array(d * [alpha])
+lp_bary, lp_log = ot.lp.barycenter(
+    A, M, weights, solver='interior-point', verbose=False, log=True)
+print('Time\t: ', ot.toc(''))
+print('Obj\t: ', lp_log['fun'])
+
+print('')
+print('Discrete MMOT Algorithm:')
+ot.tic()
+# dmmot_obj, log = ot.lp.discrete_mmot(A.T, n, d)
+barys, log = ot.lp.discrete_mmot_converge(
+    A, niters=3000, lr=0.000002, log=True)
+dmmot_obj = log['primal objective']
+print('Time\t: ', ot.toc(''))
+print('Obj\t: ', dmmot_obj)
+
+
+# %%
+# Compare Barycenters in both methods
+# ---------
+pl.figure(1, figsize=(6.4, 3))
+for i in range(len(barys)):
+    if i == 0:
+        pl.plot(x, barys[i], 'g-*', label='Discrete MMOT')
+    else:
+        continue
+        # pl.plot(x, barys[i], 'g-*')
+pl.plot(x, lp_bary, 'k-', label='LP Barycenter')
+pl.plot(x, a1, 'b', label='Source distribution')
+pl.plot(x, a2, 'r', label='Target distribution')
+pl.title('Barycenters')
+pl.legend()
+
+# %%
+# More than 2 distributions
+# --------------------------------------------------
+# Generate 7 pseudorandom gaussian distributions with 50 bins.
+n = 50  # nb bins
+d = 7
+vecsize = n * d
+
+data = []
+for i in range(d):
+    m = n * (0.5 * np.random.rand(1)) * float(np.random.randint(2) + 1)
+    a = ot.datasets.make_1D_gauss(n, m=m, s=5)
+    data.append(a)
+
+x = np.arange(n, dtype=np.float64)
+M = ot.utils.dist(x.reshape((n, 1)), metric='minkowski')
+A = np.vstack(data).T
+
+pl.figure(1, figsize=(6.4, 3))
+for i in range(len(data)):
+    pl.plot(x, data[i])
+
+pl.title('Distributions')
+pl.legend()
+
+# %%
+# Minimizing Distances Among Many Distributions
+# ---------------
+# The objective being minimized is different for both methods, so the objective
+# values cannot be compared.
+
+# Perform gradient descent optimization using the d-MMOT method.
+barys = ot.lp.discrete_mmot_converge(A, niters=9000, lr=0.00001)
+
+# after minimization, any distribution can be used as a estimate of barycenter.
+bary = barys[0]
+
+# Compute 1D Wasserstein barycenter using the LP method
+weights = ot.unif(d)
+lp_bary, bary_log = ot.lp.barycenter(A, M, weights, solver='interior-point',
+                                     verbose=True, log=True)
+
+# %%
+# Compare Barycenters in both methods
+# ---------
+pl.figure(1, figsize=(6.4, 3))
+pl.plot(x, bary, 'g-*', label='Discrete MMOT')
+pl.plot(x, lp_bary, 'k-', label='LP Wasserstein')
+pl.title('Barycenters')
+pl.legend()
+
+# %%
+# Compare with original distributions
+# ---------
+pl.figure(1, figsize=(6.4, 3))
+for i in range(len(data)):
+    pl.plot(x, data[i])
+for i in range(len(barys)):
+    if i == 0:
+        pl.plot(x, barys[i], 'g-*', label='Discrete MMOT')
+    else:
+        continue
+        # pl.plot(x, barys[i], 'g')
+pl.plot(x, lp_bary, 'k-', label='LP Wasserstein')
+# pl.plot(x, bary, 'g', label='Discrete MMOT')
+pl.title('Barycenters')
+pl.legend()

ot/lp/__init__.py

@@ -17,6 +17,7 @@
 
 from . import cvx
 from .cvx import barycenter
+from .dmmot import *
 
 # import compiled emd
 from .emd_wrap import emd_c, check_result, emd_1d_sorted
@@ -30,7 +31,8 @@
 
 __all__ = ['emd', 'emd2', 'barycenter', 'free_support_barycenter', 'cvx', ' emd_1d_sorted',
            'emd_1d', 'emd2_1d', 'wasserstein_1d', 'generalized_free_support_barycenter',
-           'binary_search_circle', 'wasserstein_circle', 'semidiscrete_wasserstein2_unif_circle']
+           'binary_search_circle', 'wasserstein_circle', 'semidiscrete_wasserstein2_unif_circle',
+           'discrete_mmot', 'discrete_mmot_converge']
 
 
 def check_number_threads(numThreads):