integrate comments of jmassich

Author: Hicham Janati

examples/plot_UOT_1D.py

@@ -66,9 +66,9 @@
 
 #%% Sinkhorn
 
-lambd = 0.1
-alpha = 1.
-Gs = ot.unbalanced.sinkhorn_unbalanced(a, b, M, lambd, alpha, verbose=True)
+epsilon = 0.1  # entropy parameter
+alpha = 1.  # Unbalanced KL relaxation parameter
+Gs = ot.unbalanced.sinkhorn_unbalanced(a, b, M, epsilon, alpha, verbose=True)
 
 pl.figure(4, figsize=(5, 5))
 ot.plot.plot1D_mat(a, b, Gs, 'UOT matrix Sinkhorn')

ot/unbalanced.py

@@ -6,6 +6,7 @@
 # Author: Hicham Janati <hicham.janati@inria.fr>
 # License: MIT License
 
+import warnings
 import numpy as np
 # from .utils import unif, dist
 
@@ -29,7 +30,7 @@ def sinkhorn_unbalanced(a, b, M, reg, alpha, method='sinkhorn', numItermax=1000,
     - a and b are source and target weights
     - KL is the Kullback-Leibler divergence
 
-    The algorithm used for solving the problem is the generalized Sinkhorn-Knopp matrix scaling algorithm as proposed in [10]_
+    The algorithm used for solving the problem is the generalized Sinkhorn-Knopp matrix scaling algorithm as proposed in [10, 23]_
 
 
     Parameters
@@ -85,33 +86,23 @@ def sinkhorn_unbalanced(a, b, M, reg, alpha, method='sinkhorn', numItermax=1000,
 
     .. [10] Chizat, L., Peyré, G., Schmitzer, B., & Vialard, F. X. (2016). Scaling algorithms for unbalanced transport problems. arXiv preprint arXiv:1607.05816.
 
+    .. [23] Frogner C., Zhang C., Mobahi H., Araya-Polo M., Poggio T. : Learning with a Wasserstein Loss,  Advances in Neural Information Processing Systems (NIPS) 2015
 
 
     See Also
     --------
-    ot.lp.emd : Unregularized OT
-    ot.optim.cg : General regularized OT
-    ot.bregman.sinkhorn_knopp : Classic Sinkhorn [2]
-    ot.bregman.sinkhorn_stabilized: Stabilized sinkhorn [9][10]
-    ot.bregman.sinkhorn_epsilon_scaling: Sinkhorn with epslilon scaling [9][10]
+    ot.unbalanced.sinkhorn_knopp : Unbalanced Classic Sinkhorn [10]
+    ot.unbalanced.sinkhorn_stabilized: Unbalanced Stabilized sinkhorn [9][10]
+    ot.unbalanced.sinkhorn_epsilon_scaling: Unbalanced Sinkhorn with epslilon scaling [9][10]
 
     """
 
     if method.lower() == 'sinkhorn':
         def sink():
             return sinkhorn_knopp(a, b, M, reg, alpha, numItermax=numItermax,
                                   stopThr=stopThr, verbose=verbose, log=log, **kwargs)
-    # elif method.lower() == 'sinkhorn_stabilized':
-    #     def sink():
-    #         return sinkhorn_stabilized(a, b, M, reg, numItermax=numItermax,
-    #                                    stopThr=stopThr, verbose=verbose, log=log, **kwargs)
-    # elif method.lower() == 'sinkhorn_epsilon_scaling':
-    #     def sink():
-    #         return sinkhorn_epsilon_scaling(
-    #             a, b, M, reg, numItermax=numItermax,
-    #             stopThr=stopThr, verbose=verbose, log=log, **kwargs)
     else:
-        print('Warning : unknown method. Falling back to classic Sinkhorn Knopp')
+        warnings.warn('Unknown method. Falling back to classic Sinkhorn Knopp')
 
         def sink():
             return sinkhorn_knopp(a, b, M, reg, alpha, numItermax=numItermax,
@@ -139,7 +130,7 @@ def sinkhorn2(a, b, M, reg, alpha, method='sinkhorn', numItermax=1000,
     - a and b are source and target weights
     - KL is the Kullback-Leibler divergence
 
-    The algorithm used for solving the problem is the generalized Sinkhorn-Knopp matrix scaling algorithm as proposed in [10]_
+    The algorithm used for solving the problem is the generalized Sinkhorn-Knopp matrix scaling algorithm as proposed in [10, 23]_
 
 
     Parameters
@@ -196,36 +187,22 @@ def sinkhorn2(a, b, M, reg, alpha, method='sinkhorn', numItermax=1000,
 
     .. [10] Chizat, L., Peyré, G., Schmitzer, B., & Vialard, F. X. (2016). Scaling algorithms for unbalanced transport problems. arXiv preprint arXiv:1607.05816.
 
-       [21] Altschuler J., Weed J., Rigollet P. : Near-linear time approximation algorithms for optimal transport via Sinkhorn iteration, Advances in Neural Information Processing Systems (NIPS) 31, 2017
-
-
+    .. [23] Frogner C., Zhang C., Mobahi H., Araya-Polo M., Poggio T. : Learning with a Wasserstein Loss,  Advances in Neural Information Processing Systems (NIPS) 2015
 
     See Also
     --------
-    ot.lp.emd : Unregularized OT
-    ot.optim.cg : General regularized OT
-    ot.bregman.sinkhorn_knopp : Classic Sinkhorn [2]
-    ot.bregman.greenkhorn : Greenkhorn [21]
-    ot.bregman.sinkhorn_stabilized: Stabilized sinkhorn [9][10]
-    ot.bregman.sinkhorn_epsilon_scaling: Sinkhorn with epslilon scaling [9][10]
+    ot.unbalanced.sinkhorn_knopp : Unbalanced Classic Sinkhorn [10]
+    ot.unbalanced.sinkhorn_stabilized: Unbalanced Stabilized sinkhorn [9][10]
+    ot.unbalanced.sinkhorn_epsilon_scaling: Unbalanced Sinkhorn with epslilon scaling [9][10]
 
     """
 
     if method.lower() == 'sinkhorn':
         def sink():
             return sinkhorn_knopp(a, b, M, reg, alpha, numItermax=numItermax,
                                   stopThr=stopThr, verbose=verbose, log=log, **kwargs)
-    # elif method.lower() == 'sinkhorn_stabilized':
-    #     def sink():
-    #         return sinkhorn_stabilized(a, b, M, reg, numItermax=numItermax,
-    #                                    stopThr=stopThr, verbose=verbose, log=log, **kwargs)
-    # elif method.lower() == 'sinkhorn_epsilon_scaling':
-    #     def sink():
-    #         return sinkhorn_epsilon_scaling(
-    #             a, b, M, reg, numItermax=numItermax,
-    #             stopThr=stopThr, verbose=verbose, log=log, **kwargs)
     else:
-        print('Warning : unknown method using classic Sinkhorn Knopp')
+        warnings.warn('Unknown method using classic Sinkhorn Knopp')
 
         def sink():
             return sinkhorn_knopp(a, b, M, reg, alpha, **kwargs)
@@ -256,7 +233,7 @@ def sinkhorn_knopp(a, b, M, reg, alpha, numItermax=1000,
     - a and b are source and target weights
     - KL is the Kullback-Leibler divergence
 
-    The algorithm used for solving the problem is the generalized Sinkhorn-Knopp matrix scaling algorithm as proposed in [10]_
+    The algorithm used for solving the problem is the generalized Sinkhorn-Knopp matrix scaling algorithm as proposed in [10, 23]_
 
 
     Parameters
@@ -306,6 +283,7 @@ def sinkhorn_knopp(a, b, M, reg, alpha, numItermax=1000,
 
     .. [10] Chizat, L., Peyré, G., Schmitzer, B., & Vialard, F. X. (2016). Scaling algorithms for unbalanced transport problems. arXiv preprint arXiv:1607.05816.
 
+    .. [23] Frogner C., Zhang C., Mobahi H., Araya-Polo M., Poggio T. : Learning with a Wasserstein Loss,  Advances in Neural Information Processing Systems (NIPS) 2015
 
     See Also
     --------
@@ -368,7 +346,7 @@ def sinkhorn_knopp(a, b, M, reg, alpha, numItermax=1000,
                 or np.any(np.isinf(u)) or np.any(np.isinf(v))):
             # we have reached the machine precision
             # come back to previous solution and quit loop
-            print('Warning: numerical errors at iteration', cpt)
+            warnings.warn('Numerical errors at iteration', cpt)
             u = uprev
             v = vprev
             break

test/test_unbalanced.py

@@ -6,15 +6,19 @@
 
 import numpy as np
 import ot
+import pytest
 
 
-def test_unbalanced():
+@pytest.mark.parametrize("metric", ["sinkhorn"])
+def test_unbalanced_convergence(method):
     # test generalized sinkhorn for unbalanced OT
     n = 100
     rng = np.random.RandomState(42)
 
     x = rng.randn(n, 2)
     a = ot.utils.unif(n)
+
+    # make dists unbalanced
     b = ot.utils.unif(n) * 1.5
 
     M = ot.dist(x, x)
@@ -23,7 +27,8 @@ def test_unbalanced():
     K = np.exp(- M / epsilon)
 
     G, log = ot.unbalanced.sinkhorn_unbalanced(a, b, M, reg=epsilon, alpha=alpha,
-                                               stopThr=1e-10, log=True)
+                                               stopThr=1e-10, method=method,
+                                               log=True)
 
     # check fixed point equations
     fi = alpha / (alpha + epsilon)