better coverage

Author: rflamary

ot/solvers/_linear.py

@@ -490,13 +490,13 @@ def solve(
                 else:
                     # stabilize kl of 0 mass
                     if nx.any(a == 0) or nx.any(b == 0):
-                        a = a + 1e-16 * nx.min(a)
-                        b = b + 1e-16 * nx.min(b)
+                        a = a + 1e-15 * nx.min(a)
+                        b = b + 1e-15 * nx.min(b)
                         print(
                             "Warning: a or b has 0 mass, adding small mass to avoid NaN in KL divergence."
                         )
                     value = value_linear + reg * nx.kl_div(
-                        plan, a[:, None] * b[None, :]
+                        plan, a[:, None] * b[None, :] + 1e-15
                     )
 
                 if grad == "envelope":  # set the gradient at convergence

test/test_solvers.py

@@ -33,11 +33,14 @@
     {"method": "gaussian"},
     {"method": "gaussian", "reg": 1},
     {"method": "gaussian_hd", "rank": 1},
+    {"method": "gaussian_hd", "rank": 3},
     {"method": "factored", "rank": 2},
     {"method": "lowrank", "rank": 2, "max_iter": 5},
     {"method": "nystroem", "rank": 2},
     {"method": "sliced", "n_projections": 10},
+    {"method": "sliced", "n_projections": 10, "metric": "euclidean"},
     {"method": "max_sliced", "n_projections": 10},
+    {"method": "max_sliced", "n_projections": 10, "metric": "euclidean"},
 ]
 
 lst_parameters_solve_sample_NotImplemented = [
@@ -808,6 +811,32 @@ def test_solve_sample_methods(nx, method_params):
         np.testing.assert_allclose(sol2.value, 0, atol=1e-10)
 
 
+def test_zero_mass_solvers():
+    # test that solvers handle zero mass distributions correctly
+    n_samples_s = 10
+    n_samples_t = 9
+    n_features = 2
+    rng = np.random.RandomState(42)
+
+    x = rng.randn(n_samples_s, n_features)
+    y = rng.randn(n_samples_t, n_features)
+    a = ot.utils.unif(n_samples_s)
+    b = ot.utils.unif(n_samples_t)
+
+    # Set one of the distributions to zero mass
+    a[0] = 0.0
+    b[0] = 0.0
+
+    a /= a.sum()
+    b /= b.sum()
+
+    res = ot.solve_sample(x, y, a, b)
+    res_reg = ot.solve_sample(x, y, a, b, reg=1.0)
+
+    assert np.isnan(res.value) == False
+    assert np.isnan(res_reg.value) == False
+
+
 @pytest.skip_backend("tf", reason="Not implemented for tf backend")
 @pytest.mark.parametrize("debias", [True, False, "split"])
 @pytest.mark.parametrize("reg", [None, 10])

test/test_utils.py

@@ -975,7 +975,7 @@ def test_split_sample_ratio(nx, ratio, n, random):
 
     # test uniform weights
     X1, X2, a1, a2, id1, id2 = ot.utils.split_sample_ratio(
-        X, ratio=ratio, random_state=seed, nx=nx
+        X, ratio=ratio, random_split=random, random_state=seed, nx=nx
     )
 
     np.testing.assert_allclose(nx.to_numpy(a1).sum(), ratio, atol=1e-5)
@@ -986,7 +986,7 @@ def test_split_sample_ratio(nx, ratio, n, random):
 
     # test non uniform weights
     X1, X2, a1, a2, id1, id2 = ot.utils.split_sample_ratio(
-        X, ratio=ratio, a=a, random_state=seed
+        X, ratio=ratio, a=a, random_split=random, random_state=seed
     )
 
     with pytest.raises(ValueError):