[MRG] New API for OT solver (with pre-computed ground cost matrix) (#388)

* new API for OT solver

* use itertools for product of parameters

* add tests for result class

* add tests for result class

* add tests for result class last time?

* add sinkhorn

* make partial OT bckend compatible

* add TV as unbalanced flavor

* better tests

* make smoth backend compatible and add l2 tregularizatio to solve

* add reularizedd unbalanced

* add test for more complex attibutes

* add test for more complex attibutes

* add generic unbalaned solver and implement it for ot.solve

* add entropy to possible regularization

* star of documentation for ot.solv

* weird new pep8

* documenttaion for function ot.solve done

* pep8

* Update ot/solvers.py

Co-authored-by: Alexandre Gramfort <alexandre.gramfort@m4x.org>

* update release file

* Apply suggestions from code review

Co-authored-by: Alexandre Gramfort <alexandre.gramfort@m4x.org>

* add test NotImplemented

* pep8

* pep8gcmp pep8!

* compute kl in backend

* debug tensorflow kl backend

Co-authored-by: Alexandre Gramfort <alexandre.gramfort@m4x.org>

Author: rflamary

RELEASES.md

@@ -6,6 +6,9 @@
 
 - Added Generalized Wasserstein Barycenter solver + example (PR #372), fixed graphical details on the example (PR #376)
 - Added Free Support Sinkhorn Barycenter + example (PR #387)
+- New API for OT solver using function `ot.solve` (PR #388)
+- Backend version of `ot.partial` and `ot.smooth`  (PR #388)
+
 
 #### Closed issues
 

ot/__init__.py

@@ -34,6 +34,7 @@
 from . import regpath
 from . import weak
 from . import factored
+from . import solvers
 
 # OT functions
 from .lp import emd, emd2, emd_1d, emd2_1d, wasserstein_1d
@@ -46,7 +47,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
 
 # utils functions
 from .utils import dist, unif, tic, toc, toq
@@ -61,5 +62,5 @@
            'sinkhorn_unbalanced2', 'sliced_wasserstein_distance',
            'gromov_wasserstein', 'gromov_wasserstein2', 'gromov_barycenters', 'fused_gromov_wasserstein', 'fused_gromov_wasserstein2',
             'max_sliced_wasserstein_distance', 'weak_optimal_transport',
-            'factored_optimal_transport', 
-           'smooth', 'stochastic', 'unbalanced', 'partial', 'regpath']
+            'factored_optimal_transport', 'solve',
+           'smooth', 'stochastic', 'unbalanced', 'partial', 'regpath', 'solvers']

ot/backend.py

@@ -854,6 +854,21 @@ def sqrtm(self, a):
         """
         raise NotImplementedError()
 
+    def kl_div(self, p, q, eps=1e-16):
+        r"""
+        Computes the Kullback-Leibler divergence.
+
+        This function follows the api from :any:`scipy.stats.entropy`.
+
+        Parameter eps is used to avoid numerical errors and is added in the log.
+
+        .. math::
+             KL(p,q) = \sum_i p(i) \log (\frac{p(i)}{q(i)}+\epsilon)
+
+        See: https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.entropy.html
+        """
+        raise NotImplementedError()
+
     def isfinite(self, a):
         r"""
         Tests element-wise for finiteness (not infinity and not Not a Number).
@@ -1158,6 +1173,9 @@ def inv(self, a):
     def sqrtm(self, a):
         return scipy.linalg.sqrtm(a)
 
+    def kl_div(self, p, q, eps=1e-16):
+        return np.sum(p * np.log(p / q + eps))
+
     def isfinite(self, a):
         return np.isfinite(a)
 
@@ -1481,6 +1499,9 @@ def sqrtm(self, a):
         L, V = jnp.linalg.eigh(a)
         return (V * jnp.sqrt(L)[None, :]) @ V.T
 
+    def kl_div(self, p, q, eps=1e-16):
+        return jnp.sum(p * jnp.log(p / q + eps))
+
     def isfinite(self, a):
         return jnp.isfinite(a)
 
@@ -1901,6 +1922,9 @@ def sqrtm(self, a):
         L, V = torch.linalg.eigh(a)
         return (V * torch.sqrt(L)[None, :]) @ V.T
 
+    def kl_div(self, p, q, eps=1e-16):
+        return torch.sum(p * torch.log(p / q + eps))
+
     def isfinite(self, a):
         return torch.isfinite(a)
 
@@ -2248,6 +2272,9 @@ def sqrtm(self, a):
         L, V = cp.linalg.eigh(a)
         return (V * self.sqrt(L)[None, :]) @ V.T
 
+    def kl_div(self, p, q, eps=1e-16):
+        return cp.sum(p * cp.log(p / q + eps))
+
     def isfinite(self, a):
         return cp.isfinite(a)
 
@@ -2608,6 +2635,9 @@ def inv(self, a):
     def sqrtm(self, a):
         return tf.linalg.sqrtm(a)
 
+    def kl_div(self, p, q, eps=1e-16):
+        return tnp.sum(p * tnp.log(p / q + eps))
+
     def isfinite(self, a):
         return tnp.isfinite(a)
 

ot/partial.py

@@ -8,6 +8,8 @@
 
 import numpy as np
 from .lp import emd
+from .backend import get_backend
+from .utils import list_to_array
 
 
 def partial_wasserstein_lagrange(a, b, M, reg_m=None, nb_dummies=1, log=False,
@@ -114,14 +116,22 @@ def partial_wasserstein_lagrange(a, b, M, reg_m=None, nb_dummies=1, log=False,
     ot.partial.partial_wasserstein : Partial Wasserstein with fixed mass
     """
 
-    if np.sum(a) > 1 or np.sum(b) > 1:
+    a, b, M = list_to_array(a, b, M)
+
+    nx = get_backend(a, b, M)
+
+    if nx.sum(a) > 1 or nx.sum(b) > 1:
         raise ValueError("Problem infeasible. Check that a and b are in the "
                          "simplex")
 
     if reg_m is None:
-        reg_m = np.max(M) + 1
-    if reg_m < -np.max(M):
-        return np.zeros((len(a), len(b)))
+        reg_m = float(nx.max(M)) + 1
+    if reg_m < -nx.max(M):
+        return nx.zeros((len(a), len(b)), type_as=M)
+
+    a0, b0, M0 = a, b, M
+    # convert to humpy
+    a, b, M = nx.to_numpy(a, b, M)
 
     eps = 1e-20
     M = np.asarray(M, dtype=np.float64)
@@ -149,10 +159,16 @@ def partial_wasserstein_lagrange(a, b, M, reg_m=None, nb_dummies=1, log=False,
     gamma = np.zeros((len(a), len(b)))
     gamma[np.ix_(idx_x, idx_y)] = gamma_extended[:-nb_dummies, :-nb_dummies]
 
+    # convert back to backend
+    gamma = nx.from_numpy(gamma, type_as=M0)
+
     if log_emd['warning'] is not None:
         raise ValueError("Error in the EMD resolution: try to increase the"
                          " number of dummy points")
-    log_emd['cost'] = np.sum(gamma * M)
+    log_emd['cost'] = nx.sum(gamma * M0)
+    log_emd['u'] = nx.from_numpy(log_emd['u'], type_as=a0)
+    log_emd['v'] = nx.from_numpy(log_emd['v'], type_as=b0)
+
     if log:
         return gamma, log_emd
     else:
@@ -250,15 +266,23 @@ def partial_wasserstein(a, b, M, m=None, nb_dummies=1, log=False, **kwargs):
     entropic regularization parameter
     """
 
+    a, b, M = list_to_array(a, b, M)
+
+    nx = get_backend(a, b, M)
+
     if m is None:
         return partial_wasserstein_lagrange(a, b, M, log=log, **kwargs)
     elif m < 0:
         raise ValueError("Problem infeasible. Parameter m should be greater"
                          " than 0.")
-    elif m > np.min((np.sum(a), np.sum(b))):
+    elif m > nx.min((nx.sum(a), nx.sum(b))):
         raise ValueError("Problem infeasible. Parameter m should lower or"
                          " equal than min(|a|_1, |b|_1).")
 
+    a0, b0, M0 = a, b, M
+    # convert to humpy
+    a, b, M = nx.to_numpy(a, b, M)
+
     b_extended = np.append(b, [(np.sum(a) - m) / nb_dummies] * nb_dummies)
     a_extended = np.append(a, [(np.sum(b) - m) / nb_dummies] * nb_dummies)
     M_extended = np.zeros((len(a_extended), len(b_extended)))
@@ -267,15 +291,20 @@ def partial_wasserstein(a, b, M, m=None, nb_dummies=1, log=False, **kwargs):
 
     gamma, log_emd = emd(a_extended, b_extended, M_extended, log=True,
                          **kwargs)
+
+    gamma = nx.from_numpy(gamma[:len(a), :len(b)], type_as=M)
+
     if log_emd['warning'] is not None:
         raise ValueError("Error in the EMD resolution: try to increase the"
                          " number of dummy points")
-    log_emd['partial_w_dist'] = np.sum(M * gamma[:len(a), :len(b)])
+    log_emd['partial_w_dist'] = nx.sum(M0 * gamma)
+    log_emd['u'] = nx.from_numpy(log_emd['u'][:len(a)], type_as=a0)
+    log_emd['v'] = nx.from_numpy(log_emd['v'][:len(b)], type_as=b0)
 
     if log:
-        return gamma[:len(a), :len(b)], log_emd
+        return gamma, log_emd
     else:
-        return gamma[:len(a), :len(b)]
+        return gamma
 
 
 def partial_wasserstein2(a, b, M, m=None, nb_dummies=1, log=False, **kwargs):

ot/smooth.py

@@ -44,6 +44,7 @@
 
 import numpy as np
 from scipy.optimize import minimize
+from .backend import get_backend
 
 
 def projection_simplex(V, z=1, axis=None):
@@ -511,22 +512,28 @@ def smooth_ot_dual(a, b, M, reg, reg_type='l2', method="L-BFGS-B", stopThr=1e-9,
 
     """
 
+    nx = get_backend(a, b, M)
+
     if reg_type.lower() in ['l2', 'squaredl2']:
         regul = SquaredL2(gamma=reg)
     elif reg_type.lower() in ['entropic', 'negentropy', 'kl']:
         regul = NegEntropy(gamma=reg)
     else:
         raise NotImplementedError('Unknown regularization')
 
+    a0, b0, M0 = a, b, M
+    # convert to humpy
+    a, b, M = nx.to_numpy(a, b, M)
+
     # solve dual
     alpha, beta, res = solve_dual(a, b, M, regul, max_iter=numItermax,
                                   tol=stopThr, verbose=verbose)
 
     # reconstruct transport matrix
-    G = get_plan_from_dual(alpha, beta, M, regul)
+    G = nx.from_numpy(get_plan_from_dual(alpha, beta, M, regul), type_as=M0)
 
     if log:
-        log = {'alpha': alpha, 'beta': beta, 'res': res}
+        log = {'alpha': nx.from_numpy(alpha, type_as=a0), 'beta': nx.from_numpy(beta, type_as=b0), 'res': res}
         return G, log
     else:
         return G