diff --git a/RELEASES.md b/RELEASES.md index 5d58d2128..7786f6ddf 100644 --- a/RELEASES.md +++ b/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) diff --git a/ot/lp/solver_1d.py b/ot/lp/solver_1d.py index 802621c14..ab02d0488 100644 --- a/ot/lp/solver_1d.py +++ b/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: diff --git a/test/test_1d_solver.py b/test/test_1d_solver.py index d8907d3b9..2491adcca 100644 --- a/test/test_1d_solver.py +++ b/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)