[MRG] Handle ratio=1 in sample distribution splitting

RELEASES.md

@@ -15,6 +15,7 @@
 
 #### Closed issues
 
+- Fix `ot.utils.split_sample_ratio` at `ratio=1`, returning the complete first distribution and an empty second distribution while preserving random permutations and selectors (PR #878, Issue #877).
 - Remove a leftover debug `print` from `ot.utils.projection_sparse_simplex` with `axis=1`, and make the `ot.datasets.make_gauss_hd` docstring a raw string so importing `ot` no longer emits a `SyntaxWarning` (PR #860)
 - Fix `ot.dist` ignoring the weights `w` for `metric="cityblock"`, which returned the unweighted distance although the weights are documented for this metric (PR #859)
 - Fix swapped arguments to `div_to_product` in `ot.gromov.fused_unbalanced_across_spaces_cost`: with `reg_type="independent"` (UCOOT) the entropic terms used the plan marginals as the reference measures and vice versa (PR #855, Issue #854)

ot/utils.py

@@ -2158,6 +2158,11 @@ def split_sample_ratio(
         X_a = X_a[perm]
         a = a[perm]
 
+    if ratio == 1:
+        sel_a1 = perm if random_split else slice(0, n_a)
+        sel_a2 = perm[n_a:] if random_split else slice(n_a, n_a)
+        return X_a, X_a[:0], a, a[:0], sel_a1, sel_a2
+
     # find the split indices
     acs = nx.cumsum(a)
     thr_a = ratio * nx.sum(a)

test/test_utils.py

@@ -1046,3 +1046,26 @@ def test_split_sample_ratio(nx, ratio, n, random):
         ot.utils.split_sample_ratio(X, ratio=1.5, a=a, random_state=seed, nx=nx)
 
     #
+
+
+@pytest.mark.parametrize("ratio", [0, 1])
+@pytest.mark.parametrize("random", [False, True])
+@pytest.mark.parametrize("weighted", [False, True])
+def test_split_sample_ratio_endpoints(nx, ratio, random, weighted):
+    values = np.arange(6, dtype=float).reshape(3, 2)
+    weights = np.array([0.2, 0.0, 0.8]) if weighted else np.full(3, 1 / 3)
+    X = nx.from_numpy(values)
+    a = nx.from_numpy(weights) if weighted else None
+    X1, X2, a1, a2, id1, id2 = ot.utils.split_sample_ratio(
+        X, a=a, ratio=ratio, random_split=random, random_state=42
+    )
+    full_X, full_a, full_id = (X1, a1, id1) if ratio == 1 else (X2, a2, id2)
+    empty_X, empty_a = (X2, a2) if ratio == 1 else (X1, a1)
+    np.testing.assert_array_equal(nx.to_numpy(full_X), nx.to_numpy(X[full_id]))
+    order = np.argsort(nx.to_numpy(full_X)[:, 0])
+    np.testing.assert_array_equal(nx.to_numpy(full_X)[order], values)
+    np.testing.assert_allclose(nx.to_numpy(full_a)[order], weights)
+    assert empty_X.shape == (0, 2)
+    assert empty_a.shape == (0,)
+    empty_id = id2 if ratio == 1 else id1
+    assert X[empty_id].shape == empty_X.shape