Handle both endpoints in sample-ratio splitting

RELEASES.md

@@ -18,6 +18,7 @@
 
 #### Closed issues
 
+- Fix `ot.utils.split_sample_ratio` at `ratio=1` and handle `ratio=0` without scalar backend search, returning full/empty distributions while preserving random permutations and selectors (PR #878, Issue #877).
 - Fix duplicate function and class descriptions in generated module API pages while preserving top-level aliases and example backreferences (Issue #527)
 - Fix internal arrays created on the torch default device instead of the input device in `ot.partial.partial_wasserstein`, `ot.wasserstein_circle` (and `ot.sliced_wasserstein_sphere`), the high-dimensional Bures-Wasserstein functions in `ot.gaussian`, `ot.gmm`, `ot.lowrank_sinkhorn` with `init="deterministic"`, the semi-relaxed (F)GW barycenters and `ot.utils.projection_sparse_simplex`, which made them fail on GPU inputs. Add a `meta_default_device` test fixture to catch such allocations on CPU (PR #883, Issue #852)
 - Allow `NumpyBackend.seed` to adopt an existing `np.random.RandomState` instance and remove NumPy-specific random sampling paths in sliced utilities (PR #849, Issue #848)

ot/utils.py

@@ -2157,6 +2157,16 @@ 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
+
+    if ratio == 0:
+        sel_a1 = perm[:0] if random_split else slice(0, 0)
+        sel_a2 = perm if random_split else slice(0, n_a)
+        return X_a[:0], X_a, a[:0], a, 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