[MRG] Stabilize simplex projection against large common offsets

RELEASES.md

@@ -18,6 +18,7 @@
 
 #### Closed issues
 
+- Center floating-point inputs in `ot.utils.proj_simplex` to prevent large common offsets from erasing the simplex mass (Issue #879, PR #880).
 - 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

@@ -111,6 +111,10 @@ def proj_simplex(v, z=1):
     .. note:: This function is backend-compatible and will work on arrays
         from all compatible backends.
 
+    Floating-point inputs are centered by their column maximum before projection.
+    A common offset does not change the simplex projection; removing it prevents
+    a large offset from erasing the simplex size in the cumulative sums.
+
     Parameters
     ----------
     v : {array-like}, shape (n, d)
@@ -123,6 +127,9 @@ def proj_simplex(v, z=1):
         Array of projections on the simplex
     """
     nx = get_backend(v)
+    # The simplex projection is invariant to a common offset per column.
+    if nx.is_floating_point(v):
+        v = v - nx.max(v, axis=0)
     n = v.shape[0]
     if v.ndim == 1:
         d1 = 1

test/test_utils.py

@@ -99,6 +99,24 @@ def test_proj_simplex(nx):
     np.testing.assert_allclose(l1, l2, atol=1e-5)
 
 
+@pytest.mark.parametrize(
+    "dtype, offset",
+    [
+        (np.float32, 0),
+        (np.float32, 1e8),
+        (np.float64, 0),
+        (np.float64, 1e16),
+        (np.int64, 0),
+    ],
+)
+@pytest.mark.parametrize("matrix", [False, True])
+def test_proj_simplex_large_translation(nx, dtype, offset, matrix):
+    values = np.full((3, 2) if matrix else (3,), offset, dtype=dtype)
+    result = nx.to_numpy(ot.utils.proj_simplex(nx.from_numpy(values)))
+    np.testing.assert_allclose(result, np.full(values.shape, 1 / 3), rtol=1e-6)
+    np.testing.assert_allclose(result.sum(axis=0), 1, rtol=1e-6)
+
+
 def test_projection_sparse_simplex():
     def double_sort_projection_sparse_simplex(X, max_nz, z=1, axis=None):
         r"""This is an equivalent but less efficient version