diff --git a/RELEASES.md b/RELEASES.md index 8b5aff51d..e3bba2f1e 100644 --- a/RELEASES.md +++ b/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 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) - 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) diff --git a/ot/utils.py b/ot/utils.py index 69a4854b5..4056fd782 100644 --- a/ot/utils.py +++ b/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) diff --git a/test/test_utils.py b/test/test_utils.py index 9c93d520a..56ece83b2 100644 --- a/test/test_utils.py +++ b/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