[MRG] Fix sparse simplex projection with oversized budgets

RELEASES.md

@@ -15,6 +15,7 @@
 
 #### Closed issues
 
+- Fix oversized `max_nz` budgets in `ot.utils.projection_sparse_simplex` so they return the documented unconstrained simplex projection for all axis modes (PR #876, Issue #875).
 - 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

@@ -193,6 +193,7 @@ def projection_sparse_simplex(V, max_nz, z=1, axis=None, nx=None):
         raise ValueError("V.ndim must be <= 2")
 
     if axis == 1:
+        max_nz = min(max_nz, V.shape[1])
         # For each row of V, find top max_nz values; arrange the
         # corresponding column indices such that their values are
         # in a descending order.

test/test_utils.py

@@ -142,6 +142,26 @@ def double_sort_projection_sparse_simplex(X, max_nz, z=1, axis=None):
         np.testing.assert_allclose(slow_sparse_proj, fast_sparse_proj)
 
 
+@pytest.mark.parametrize("axis", [None, 0, 1])
+@pytest.mark.parametrize("extra_budget", [0, 2])
+def test_projection_sparse_simplex_unconstrained_budget(nx, axis, extra_budget):
+    values = np.array([[0.8, 0.1, -1.0], [0.4, 0.3, 0.2]])
+    z = 0.5
+    if axis is None:
+        dimension = values.size
+        expected = np.array([0.45, 0.0, 0.0, 0.05, 0.0, 0.0])
+    elif axis == 0:
+        dimension = values.shape[0]
+        expected = np.array([[0.45, 0.15, 0.0], [0.05, 0.35, 0.5]])
+    else:
+        dimension = values.shape[1]
+        expected = np.array([[0.5, 0.0, 0.0], [4 / 15, 1 / 6, 1 / 15]])
+    result = ot.utils.projection_sparse_simplex(
+        nx.from_numpy(values), dimension + extra_budget, z=z, axis=axis
+    )
+    np.testing.assert_allclose(nx.to_numpy(result), expected, atol=1e-12)
+
+
 def test_parmap():
     n = 10