add test and rearrange arguments

Author: 6Ulm

ot/unbalanced.py

@@ -19,8 +19,8 @@
 from .utils import list_to_array, get_parameter_pair
 
 
-def sinkhorn_unbalanced(a, b, M, reg, reg_m, reg_type="entropy", warmstart=None,
-                        method='sinkhorn', numItermax=1000,
+def sinkhorn_unbalanced(a, b, M, reg, reg_m, method='sinkhorn',
+                        reg_type="entropy", warmstart=None, numItermax=1000,
                         stopThr=1e-6, verbose=False, log=False, **kwargs):
     r"""
     Solve the unbalanced entropic regularization optimal transport problem
@@ -67,6 +67,9 @@ def sinkhorn_unbalanced(a, b, M, reg, reg_m, reg_type="entropy", warmstart=None,
         For semi-relaxed case, use either
         `reg_m=(float("inf"), scalar)` or `reg_m=(scalar, float("inf"))`.
         If reg_m is an array, it must have the same backend as input arrays (a, b, M).
+    method : str
+        method used for the solver either 'sinkhorn', 'sinkhorn_stabilized' or
+        'sinkhorn_reg_scaling', see those function for specific parameters
     reg_type : string, optional
         Regularizer term. Can take two values:
         'entropy' (negative entropy)
@@ -75,10 +78,7 @@ def sinkhorn_unbalanced(a, b, M, reg, reg_m, reg_type="entropy", warmstart=None,
         :math:`\Omega(\gamma) = \text{KL}(\gamma, \mathbf{a} \mathbf{b}^T)`.
     warmstart: tuple of arrays, shape (dim_a, dim_b), optional
         Initialization of dual potentials. If provided, the dual potentials should be given
-        (that is the logarithm of the u,v sinkhorn scaling vectors).s
-    method : str
-        method used for the solver either 'sinkhorn', 'sinkhorn_stabilized' or
-        'sinkhorn_reg_scaling', see those function for specific parameters
+        (that is the logarithm of the u,v sinkhorn scaling vectors).
     numItermax : int, optional
         Max number of iterations
     stopThr : float, optional
@@ -165,8 +165,8 @@ def sinkhorn_unbalanced(a, b, M, reg, reg_m, reg_type="entropy", warmstart=None,
         raise ValueError("Unknown method '%s'." % method)
 
 
-def sinkhorn_unbalanced2(a, b, M, reg, reg_m, reg_type="entropy", warmstart=None,
-                         method='sinkhorn', numItermax=1000,
+def sinkhorn_unbalanced2(a, b, M, reg, reg_m, method='sinkhorn',
+                         reg_type="entropy", warmstart=None, numItermax=1000,
                          stopThr=1e-6, verbose=False, log=False, **kwargs):
     r"""
     Solve the entropic regularization unbalanced optimal transport problem and
@@ -212,6 +212,9 @@ def sinkhorn_unbalanced2(a, b, M, reg, reg_m, reg_type="entropy", warmstart=None
         For semi-relaxed case, use either
         `reg_m=(float("inf"), scalar)` or `reg_m=(scalar, float("inf"))`.
         If reg_m is an array, it must have the same backend as input arrays (a, b, M).
+    method : str
+        method used for the solver either 'sinkhorn', 'sinkhorn_stabilized' or
+        'sinkhorn_reg_scaling', see those function for specific parameterss
     reg_type : string, optional
         Regularizer term. Can take two values:
         'entropy' (negative entropy)
@@ -221,9 +224,6 @@ def sinkhorn_unbalanced2(a, b, M, reg, reg_m, reg_type="entropy", warmstart=None
     warmstart: tuple of arrays, shape (dim_a, dim_b), optional
         Initialization of dual potentials. If provided, the dual potentials should be given
         (that is the logarithm of the u,v sinkhorn scaling vectors).
-    method : str
-        method used for the solver either 'sinkhorn', 'sinkhorn_stabilized' or
-        'sinkhorn_reg_scaling', see those function for specific parameters
     numItermax : int, optional
         Max number of iterations
     stopThr : float, optional
@@ -435,12 +435,12 @@ def sinkhorn_knopp_unbalanced(a, b, M, reg, reg_m, reg_type="entropy",
     # distances
     if warmstart is None:
         if n_hists:
-            u = nx.ones((dim_a, 1), type_as=M) / dim_a
-            v = nx.ones((dim_b, n_hists), type_as=M) / dim_b
+            u = nx.ones((dim_a, 1), type_as=M)
+            v = nx.ones((dim_b, n_hists), type_as=M)
             a = a.reshape(dim_a, 1)
         else:
-            u = nx.ones(dim_a, type_as=M) / dim_a
-            v = nx.ones(dim_b, type_as=M) / dim_b
+            u = nx.ones(dim_a, type_as=M)
+            v = nx.ones(dim_b, type_as=M)
     else:
         u, v = nx.exp(warmstart[0]), nx.exp(warmstart[1])
 
@@ -644,12 +644,12 @@ def sinkhorn_stabilized_unbalanced(a, b, M, reg, reg_m, reg_type="entropy",
     # distances
     if warmstart is None:
         if n_hists:
-            u = nx.ones((dim_a, n_hists), type_as=M) / dim_a
-            v = nx.ones((dim_b, n_hists), type_as=M) / dim_b
+            u = nx.ones((dim_a, n_hists), type_as=M)
+            v = nx.ones((dim_b, n_hists), type_as=M)
             a = a.reshape(dim_a, 1)
         else:
-            u = nx.ones(dim_a, type_as=M) / dim_a
-            v = nx.ones(dim_b, type_as=M) / dim_b
+            u = nx.ones(dim_a, type_as=M)
+            v = nx.ones(dim_b, type_as=M)
     else:
         u, v = nx.exp(warmstart[0]), nx.exp(warmstart[1])
 

test/test_unbalanced.py

@@ -33,12 +33,12 @@ def test_unbalanced_convergence(nx, method, reg_type):
     reg_m = 1.
 
     G, log = ot.unbalanced.sinkhorn_unbalanced(
-        a, b, M, reg=epsilon, reg_m=reg_m, reg_type=reg_type,
-        method=method, log=True, verbose=True
+        a, b, M, reg=epsilon, reg_m=reg_m, method=method,
+        reg_type=reg_type, log=True, verbose=True
     )
     loss = nx.to_numpy(ot.unbalanced.sinkhorn_unbalanced2(
-        a, b, M, reg=epsilon, reg_m=reg_m, reg_type=reg_type,
-        method=method, verbose=True
+        a, b, M, reg=epsilon, reg_m=reg_m, method=method,
+        reg_type=reg_type, verbose=True
     ))
     # check fixed point equations
     # in log-domain
@@ -70,15 +70,61 @@ def test_unbalanced_convergence(nx, method, reg_type):
 
     G = ot.unbalanced.sinkhorn_unbalanced(
         a, b, M, reg=epsilon, reg_m=reg_m,
-        reg_type=reg_type, method=method, verbose=True
+        method=method, reg_type=reg_type, verbose=True
     )
     G_np = ot.unbalanced.sinkhorn_unbalanced(
         a_np, b_np, M_np, reg=epsilon, reg_m=reg_m,
-        reg_type=reg_type, method=method, verbose=True
+        method=method, reg_type=reg_type, verbose=True
     )
     np.testing.assert_allclose(G_np, nx.to_numpy(G))
 
 
+@pytest.mark.parametrize("method,reg_type", itertools.product(["sinkhorn", "sinkhorn_stabilized"], ["kl", "entropy"]))
+# @pytest.mark.parametrize("method", ["sinkhorn", "sinkhorn_stabilized"])
+def test_unbalanced_warmstart(nx, method, reg_type):
+    # 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)
+    a, b, M = nx.from_numpy(a, b, M)
+
+    epsilon = 1.
+    reg_m = 1.
+
+    dim_a, dim_b = M.shape
+    warmstart = (nx.zeros(dim_a, type_as=M), nx.zeros(dim_b, type_as=M))
+    G, log = ot.unbalanced.sinkhorn_unbalanced(
+        a, b, M, reg=epsilon, reg_m=reg_m, method=method,
+        reg_type=reg_type, warmstart=warmstart, log=True, verbose=True
+    )
+    loss = nx.to_numpy(ot.unbalanced.sinkhorn_unbalanced2(
+        a, b, M, reg=epsilon, reg_m=reg_m, method=method,
+        reg_type=reg_type, warmstart=warmstart, verbose=True
+    ))
+
+    G0, log0 = ot.unbalanced.sinkhorn_unbalanced(
+        a, b, M, reg=epsilon, reg_m=reg_m, method=method,
+        reg_type=reg_type, warmstart=None, log=True, verbose=True
+    )
+    loss0 = nx.to_numpy(ot.unbalanced.sinkhorn_unbalanced2(
+        a, b, M, reg=epsilon, reg_m=reg_m, method=method,
+        reg_type=reg_type, warmstart=None, verbose=True
+    ))
+
+    np.testing.assert_allclose(loss, loss0, atol=1e-6)
+    np.testing.assert_allclose(
+        nx.to_numpy(log["logu"]), nx.to_numpy(log0["logu"]), atol=1e-05)
+    np.testing.assert_allclose(
+        nx.to_numpy(log["logv"]), nx.to_numpy(log0["logv"]), atol=1e-05)
+    np.testing.assert_allclose(nx.to_numpy(G), nx.to_numpy(G0), atol=1e-05)
+
+
 @pytest.mark.parametrize("method,reg_m", itertools.product(["sinkhorn", "sinkhorn_stabilized"], [1, float("inf")]))
 def test_unbalanced_relaxation_parameters(nx, method, reg_m):
     # test generalized sinkhorn for unbalanced OT