[MRG] New API ot.solve_sample

RELEASES.md

@@ -16,6 +16,7 @@
 + Add exact line-search for `gromov_wasserstein` and `fused_gromov_wasserstein` with KL loss (PR #556)
 + Add KL loss to all semi-relaxed (Fused) Gromov-Wasserstein solvers (PR #559)
 + Further upgraded unbalanced OT solvers for more flexibility and future use (PR #551)
++ New API function `ot.solve_sample` for solving OT problems from empirical samples (PR #563)
 
 #### Closed issues
 - Fix line search evaluating cost outside of the interpolation range (Issue #502, PR #504)

ot/__init__.py

@@ -36,6 +36,7 @@
 from . import solvers
 from . import gaussian
 
+
 # OT functions
 from .lp import (emd, emd2, emd_1d, emd2_1d, wasserstein_1d,
                  binary_search_circle, wasserstein_circle,
@@ -50,7 +51,7 @@
                      gromov_barycenters, fused_gromov_wasserstein, fused_gromov_wasserstein2)
 from .weak import weak_optimal_transport
 from .factored import factored_optimal_transport
-from .solvers import solve, solve_gromov
+from .solvers import solve, solve_gromov, solve_sample
 
 # utils functions
 from .utils import dist, unif, tic, toc, toq
@@ -65,7 +66,7 @@
            'sinkhorn_unbalanced2', 'sliced_wasserstein_distance', 'sliced_wasserstein_sphere',
            'gromov_wasserstein', 'gromov_wasserstein2', 'gromov_barycenters', 'fused_gromov_wasserstein',
            'fused_gromov_wasserstein2', 'max_sliced_wasserstein_distance', 'weak_optimal_transport',
-           'factored_optimal_transport', 'solve', 'solve_gromov',
+           'factored_optimal_transport', 'solve', 'solve_gromov','solve_sample', 
            'smooth', 'stochastic', 'unbalanced', 'partial', 'regpath', 'solvers',
            'binary_search_circle', 'wasserstein_circle',
            'semidiscrete_wasserstein2_unif_circle', 'sliced_wasserstein_sphere_unif']

ot/bregman/_empirical.py

@@ -11,12 +11,56 @@
 
 import warnings
 
-from ..utils import dist, list_to_array, unif
+from ..utils import dist, list_to_array, unif, LazyTensor
 from ..backend import get_backend
 
 from ._sinkhorn import sinkhorn, sinkhorn2
 
 
+def get_sinkhorn_lazytensor(X_a, X_b, f, g, metric='sqeuclidean', reg=1e-1, nx=None):
+    r""" Get a LazyTensor of Sinkhorn solution from the dual potentials
+
+    The returned LazyTensor is
+    :math:`\mathbf{T} = exp(  \mathbf{f} \mathbf{1}_b^\top + \mathbf{1}_a \mathbf{g}^\top - \mathbf{C}/reg)`, where :math:`\mathbf{C}` is the pairwise metric matrix between samples :math:`\mathbf{X}_a` and :math:`\mathbf{X}_b`.
+
+    Parameters
+    ----------
+    X_a : array-like, shape (n_samples_a, dim)
+        samples in the source domain
+    X_b : array-like, shape (n_samples_b, dim)
+        samples in the target domain
+    f : array-like, shape (n_samples_a,)
+        First dual potentials (log space)
+    g : array-like, shape (n_samples_b,)
+        Second dual potentials (log space)
+    metric : str, default='sqeuclidean'
+        Metric used for the cost matrix computation
+    reg : float, default=1e-1
+        Regularization term >0
+    nx : Backend(), default=None
+        Numerical backend used
+
+
+    Returns
+    -------
+    T : LazyTensor
+        Sinkhorn solution tensor
+    """
+
+    if nx is None:
+        nx = get_backend(X_a, X_b, f, g)
+
+    shape = (X_a.shape[0], X_b.shape[0])
+
+    def func(i, j, X_a, X_b, f, g, metric, reg):
+        C = dist(X_a[i], X_b[j], metric=metric)
+        return nx.exp(f[i, None] + g[None, j] - C / reg)
+
+    T = LazyTensor(shape, func, X_a=X_a, X_b=X_b, f=f, g=g, metric=metric, reg=reg)
+
+    return T
+
+
 def empirical_sinkhorn(X_s, X_t, reg, a=None, b=None, metric='sqeuclidean',
                        numIterMax=10000, stopThr=1e-9, isLazy=False, batchSize=100, verbose=False,
                        log=False, warn=True, warmstart=None, **kwargs):
@@ -198,6 +242,8 @@ def empirical_sinkhorn(X_s, X_t, reg, a=None, b=None, metric='sqeuclidean',
         if log:
             dict_log["u"] = f
             dict_log["v"] = g
+            dict_log["niter"] = i_ot
+            dict_log["lazy_plan"] = get_sinkhorn_lazytensor(X_s, X_t, f, g, metric, reg)
             return (f, g, dict_log)
         else:
             return (f, g)

ot/gaussian.py

@@ -249,7 +249,7 @@ def bures_wasserstein_distance(ms, mt, Cs, Ct, log=False):
     Cs12 = nx.sqrtm(Cs)
 
     B = nx.trace(Cs + Ct - 2 * nx.sqrtm(dots(Cs12, Ct, Cs12)))
-    W = nx.sqrt(nx.norm(ms - mt)**2 + B)
+    W = nx.sqrt(nx.maximum(nx.norm(ms - mt)**2 + B, 0))
 
     if log:
         log = {}