[MRG] Fix batched 1D Wasserstein costs and transport plans

RELEASES.md

@@ -18,6 +18,7 @@
 
 #### Closed issues
 
+- Fix multi-axis batches in `ot.wasserstein_1d` and `quantile_function`, and return transport plans for unbatched and pre-sorted inputs. Shared weights are repeated over omitted batch axes; plans retain the original sample indices (PR #899, related to Issue #518).
 - Fix overflow to NaN in `ot.stochastic.coordinate_grad_semi_dual` by shifting before exponentiating, as `ot.stochastic.c_transform_entropic` already does. The factor cancels in the normalisation, so results are unchanged where the old code did not overflow (PR #866, Issue #264)
 - 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)

ot/lp/solver_1d.py

@@ -25,22 +25,29 @@ def quantile_function(qs, cws, xs, return_index=False):
 
     Parameters
     ----------
-    qs: array-like, shape (n,)
-        Quantiles at which the quantile function is evaluated
+    qs: array-like, shape (k, ...)
+        Quantiles at which the quantile function is evaluated. The trailing
+        batch dimensions must match those of cws and xs.
     cws: array-like, shape (m, ...)
         cumulative weights of the 1D empirical distribution, if batched, must be similar to xs
-    xs: array-like, shape (n, ...)
-        locations of the 1D empirical distribution, batched against the `xs.ndim - 1` first dimensions
+    xs: array-like, shape (m, ...)
+        Sorted locations of the 1D empirical distributions. The first axis
+        contains samples and the remaining axes contain independent batches.
     return_index: bool, optional
         if True, returns also the associated indices. If False, do not return them
 
     Returns
     -------
-    q: array-like, shape (..., n)
+    q: array-like, shape (k, ...)
         The quantiles of the distribution
     """
     nx = get_backend(qs, cws)
     n = xs.shape[0]
+    qs_shape = qs.shape
+    # Some backends support only one batch axis in searchsorted.
+    if cws.ndim > 2:
+        cws = nx.reshape(cws, (cws.shape[0], -1))
+        qs = nx.reshape(qs, (qs.shape[0], -1))
     if nx.__name__ == "torch":
         # this is to ensure the best performance for torch searchsorted
         # and avoid a warning related to non-contiguous arrays
@@ -50,7 +57,10 @@ def quantile_function(qs, cws, xs, return_index=False):
         cws = cws.T
         qs = qs.T
 
-    idx = nx.clip(nx.searchsorted(cws, qs).T, 0, n - 1)
+    idx = nx.searchsorted(cws, qs)
+    if idx.ndim > 1:
+        idx = nx.transpose(idx)
+    idx = nx.reshape(nx.clip(idx, 0, n - 1), qs_shape)
 
     if return_index:
         return nx.take_along_axis(xs, idx, axis=0), idx
@@ -88,9 +98,11 @@ def wasserstein_1d(
     v_values: array-like, shape (m, ...)
         locations of the second empirical distributions
     u_weights: array-like, shape (n, ...), optional
-        weights of the first empirical distributions, if None then uniform weights are used
+        weights of the first empirical distributions, if None then uniform weights are used.
+        Omitted trailing batch dimensions are repeated over the corresponding values axes.
     v_weights: array-like, shape (m, ...), optional
-        weights of the second empirical distributions, if None then uniform weights are used
+        weights of the second empirical distributions, if None then uniform weights are used.
+        Omitted trailing batch dimensions are repeated over the corresponding values axes.
     p: int, optional
         order of the ground metric used, should be at least 1 (see [2, Chap. 2], default is 1
     require_sort: bool, optional
@@ -113,6 +125,9 @@ def wasserstein_1d(
         if return_plans is "coo_tuple", returns the plans as a list of namedTuple
         containing the data, rows and cols of the non-zero elements of the
         transportation matrix.
+        There is one plan per batch, in row-major order over the trailing
+        dimensions of the inputs, or a single-element list for unbatched inputs.
+        Row and column indices refer to the original input sample order.
 
     References
     ----------
@@ -129,12 +144,14 @@ def wasserstein_1d(
 
     if u_weights is None:
         u_weights = nx.full(u_values.shape, 1.0 / n, type_as=u_values)
-    elif u_weights.ndim != u_values.ndim:
-        u_weights = nx.repeat(u_weights[..., None], u_values.shape[-1], -1)
+    else:
+        for axis in range(u_weights.ndim, u_values.ndim):
+            u_weights = nx.repeat(u_weights[..., None], u_values.shape[axis], -1)
     if v_weights is None:
         v_weights = nx.full(v_values.shape, 1.0 / m, type_as=v_values)
-    elif v_weights.ndim != v_values.ndim:
-        v_weights = nx.repeat(v_weights[..., None], v_values.shape[-1], -1)
+    else:
+        for axis in range(v_weights.ndim, v_values.ndim):
+            v_weights = nx.repeat(v_weights[..., None], v_values.shape[axis], -1)
 
     if require_sort:
         u_sorter = nx.argsort(u_values, 0)
@@ -145,9 +162,6 @@ def wasserstein_1d(
 
         u_weights = nx.take_along_axis(u_weights, u_sorter, 0)
         v_weights = nx.take_along_axis(v_weights, v_sorter, 0)
-    else:
-        u_sorter = nx.arange(n)
-        v_sorter = nx.arange(m)
 
     u_cumweights = nx.cumsum(u_weights, 0)
     v_cumweights = nx.cumsum(v_weights, 0)
@@ -165,28 +179,32 @@ def wasserstein_1d(
     diff_quantiles = nx.abs(u_quantiles - v_quantiles)
 
     if return_plans is not False:
-        u_quantiles_idx = nx.take_along_axis(u_sorter, idx_u, axis=0)
-        v_quantiles_idx = nx.take_along_axis(v_sorter, idx_v, axis=0)
+        if require_sort:
+            idx_u = nx.take_along_axis(u_sorter, idx_u, axis=0)
+            idx_v = nx.take_along_axis(v_sorter, idx_v, axis=0)
+        u_quantiles_idx = nx.reshape(idx_u, (n + m, -1))
+        v_quantiles_idx = nx.reshape(idx_v, (n + m, -1))
+        plan_weights = nx.reshape(delta, (n + m, -1))
         if return_plans == "coo_tuple":
             PlanTuple = namedtuple("PlanTuple", ["data", "rows", "cols"])
             plans = [
                 PlanTuple(
-                    data=delta[:, k],
+                    data=plan_weights[:, k],
                     rows=u_quantiles_idx[:, k],
                     cols=v_quantiles_idx[:, k],
                 )
-                for k in range(delta.shape[1])
+                for k in range(plan_weights.shape[1])
             ]
         else:
             plans = [
                 nx.coo_matrix(
-                    delta[:, k],
+                    plan_weights[:, k],
                     u_quantiles_idx[:, k],
                     v_quantiles_idx[:, k],
                     shape=(n, m),
                     type_as=u_values,
                 )
-                for k in range(delta.shape[1])
+                for k in range(plan_weights.shape[1])
             ]
 
     if p == 1:

test/test_1d_solver.py

@@ -347,3 +347,148 @@ def test_wasserstein_1d_plan():
         shape=(x.shape[0], y.shape[0]),
     ).toarray()
     np.testing.assert_allclose(plan_1d_tuple, plan_emd, atol=1e-05)
+
+
+@pytest.mark.parametrize("batch_shape", [(), (2,), (2, 3)])
+@pytest.mark.parametrize("weights", ["uniform", "shared", "batched"])
+@pytest.mark.parametrize("require_sort", [False, True])
+@pytest.mark.parametrize("return_plans", [True, "coo_tuple"])
+def test_wasserstein_1d_batched_plans(
+    nx, batch_shape, weights, require_sort, return_plans
+):
+    rng = np.random.RandomState(42)
+    n, m = 7, 5
+    u = rng.randn(n, *batch_shape)
+    v = rng.randn(m, *batch_shape)
+    if not require_sort:
+        u = np.sort(u, axis=0)
+        v = np.sort(v, axis=0)
+
+    if weights == "uniform":
+        a, b = np.full(u.shape, 1 / n), np.full(v.shape, 1 / m)
+        ab = bb = None
+    else:
+        weight_shape = batch_shape if weights == "batched" else ()
+        a = rng.uniform(0.1, 1, (n, *weight_shape))
+        b = rng.uniform(0.1, 1, (m, *weight_shape))
+        a[0] = b[-1] = 0
+        a /= a.sum(axis=0)
+        b /= b.sum(axis=0)
+        ab, bb = nx.from_numpy(a, b)
+        if weights == "shared":
+            a = np.broadcast_to(a.reshape((n,) + (1,) * len(batch_shape)), u.shape)
+            b = np.broadcast_to(b.reshape((m,) + (1,) * len(batch_shape)), v.shape)
+
+    ub, vb = nx.from_numpy(u, v)
+    cost, plans = wasserstein_1d(
+        ub, vb, ab, bb, p=2, require_sort=require_sort, return_plans=return_plans
+    )
+    cost_only = wasserstein_1d(ub, vb, ab, bb, p=2, require_sort=require_sort)
+    assert cost.shape == batch_shape
+    nx.assert_same_dtype_device(ub, cost)
+    np.testing.assert_allclose(nx.to_numpy(cost), nx.to_numpy(cost_only))
+    batches = list(np.ndindex(batch_shape))
+    assert len(plans) == len(batches)
+    for plan, batch in zip(plans, batches):
+        if return_plans == "coo_tuple":
+            data, rows, cols = nx.to_numpy(plan.data, plan.rows, plan.cols)
+            plan = coo_array((data, (rows, cols)), shape=(n, m)).toarray()
+        else:
+            plan = nx.to_numpy(nx.todense(plan))
+        index = (slice(None),) + batch
+        ai, bi = np.ascontiguousarray(a[index]), np.ascontiguousarray(b[index])
+        matrix = (u[index][:, None] - v[index][None, :]) ** 2
+        reference = ot.emd(ai, bi, matrix)
+        np.testing.assert_allclose(plan, reference, atol=1e-12)
+        np.testing.assert_allclose(plan.sum(axis=1), ai, atol=1e-12)
+        np.testing.assert_allclose(plan.sum(axis=0), bi, atol=1e-12)
+        np.testing.assert_allclose(nx.to_numpy(cost)[batch], (plan * matrix).sum())
+
+
+@pytest.mark.parametrize("p", [1, 2])
+def test_wasserstein_1d_shared_batch_weights(nx, p):
+    rng = np.random.RandomState(4)
+    u, v = [rng.randn(n, 2, 3) for n in (7, 5)]
+    a, b = [rng.uniform(0.1, 1, (n, 2)) for n in (7, 5)]
+    a /= a.sum(axis=0)
+    b /= b.sum(axis=0)
+    cost = nx.to_numpy(wasserstein_1d(*nx.from_numpy(u, v, a, b), p=p))
+    assert cost.shape == (2, 3)
+    for i, j in np.ndindex(2, 3):
+        matrix = np.abs(u[:, i, j, None] - v[None, :, i, j]) ** p
+        reference = ot.emd2(
+            np.ascontiguousarray(a[:, i]), np.ascontiguousarray(b[:, i]), matrix
+        )
+        np.testing.assert_allclose(cost[i, j], reference)
+
+
+@pytest.mark.parametrize("batch_shape", [(), (2,), (2, 3), (2, 1, 3)])
+def test_quantile_function_batch_axes(nx, batch_shape):
+    rng = np.random.RandomState(0)
+    xs = np.sort(rng.randn(7, *batch_shape), axis=0)
+    weights = rng.uniform(0.1, 1, xs.shape)
+    cws = np.cumsum(weights / weights.sum(axis=0), axis=0)
+    qs = np.sort(rng.uniform(0, 1, (5, *batch_shape)), axis=0)
+    qs[0], qs[-1] = 0, 1
+    expected_idx = np.empty(qs.shape, dtype=int)
+    for batch in np.ndindex(batch_shape):
+        index = (slice(None),) + batch
+        expected_idx[index] = np.clip(np.searchsorted(cws[index], qs[index]), 0, 6)
+    q, idx = ot.lp.solver_1d.quantile_function(
+        *nx.from_numpy(qs, cws, xs), return_index=True
+    )
+    np.testing.assert_array_equal(nx.to_numpy(idx), expected_idx)
+    np.testing.assert_array_equal(
+        nx.to_numpy(q), np.take_along_axis(xs, expected_idx, axis=0)
+    )
+
+
+def test_wasserstein_1d_multibatch_gradients_torch():
+    torch = pytest.importorskip("torch")
+    rng = np.random.RandomState(5)
+    u, v = [
+        torch.tensor(rng.randn(n, 2, 3), requires_grad=True).transpose(1, 2)
+        for n in (4, 5)
+    ]
+    a, b = [torch.tensor(rng.randn(n), requires_grad=True) for n in (4, 5)]
+
+    def loss(u, v, a, b):
+        return wasserstein_1d(u, v, a.softmax(0), b.softmax(0), p=2).sum()
+
+    result = loss(u, v, a, b)
+    reference = sum(
+        wasserstein_1d(u[:, i, j], v[:, i, j], a.softmax(0), b.softmax(0), p=2)
+        for i in range(3)
+        for j in range(2)
+    )
+    np.testing.assert_allclose(result.detach(), reference.detach())
+    actual_grads = torch.autograd.grad(result, (u, v, a, b))
+    reference_grads = torch.autograd.grad(reference, (u, v, a, b))
+    for actual, expected in zip(actual_grads, reference_grads):
+        np.testing.assert_allclose(actual, expected, atol=1e-12)
+    assert torch.autograd.gradcheck(loss, (u, v, a, b))
+
+
+def test_wasserstein_1d_multibatch_gradients_jax():
+    jax = pytest.importorskip("jax")
+    jnp = jax.numpy
+    rng = np.random.RandomState(5)
+    u, v, a, b = [jnp.array(rng.randn(n, 2, 3)) for n in (4, 5, 4, 5)]
+
+    def loss(u, v, a, b):
+        return wasserstein_1d(
+            u, v, jax.nn.softmax(a, axis=0), jax.nn.softmax(b, axis=0), p=2
+        ).sum()
+
+    def reference(u, v, a, b):
+        return sum(
+            loss(u[:, i, j], v[:, i, j], a[:, i, j], b[:, i, j])
+            for i in range(2)
+            for j in range(3)
+        )
+
+    np.testing.assert_allclose(jax.jit(loss)(u, v, a, b), reference(u, v, a, b))
+    actual_grads = jax.grad(loss, argnums=(0, 1, 2, 3))(u, v, a, b)
+    reference_grads = jax.grad(reference, argnums=(0, 1, 2, 3))(u, v, a, b)
+    for actual, expected in zip(actual_grads, reference_grads):
+        np.testing.assert_allclose(actual, expected, atol=1e-12)