From 1fc45baf76c8f0889098b18d455ecc52fdec08da Mon Sep 17 00:00:00 2001 From: Tom Vercauteren Date: Fri, 18 Sep 2026 12:51:32 +0100 Subject: [PATCH 1/5] [WIP] Add exact EMD-L1 solver for Cartesian grid histograms Introduce ot.lp.emd_grid_l1(A, B), an exact solver for the Earth Mover's Distance with the cityblock ground metric between histograms sharing a d-dimensional Cartesian grid support. Rather than solving the min-cost flow on the full bipartite graph (as ot.emd would), this reduces the problem to a min-cost flow on the much sparser grid adjacency graph following Ling & Okada (2007), then hands it to POT's existing off-the-shelf network simplex LP solver instead of their bespoke tree-based one. - New ot/lp/sparse_digraph.h: a general (non-bipartite) sparse digraph, needed because every grid cell can carry both supply and demand, unlike the existing bipartite sparse_bipartitegraph.h. - New EMD_wrap_grid_l1 in EMD_wrapper.cpp/EMD.h: builds the grid graph from a shape array, runs NetworkSimplexSimple, and decomposes the resulting flow into a transportation plan (including same-bin "self-mass" that already overlaps between A and B, which the flow decomposition alone would miss). - New ot/lp/_grid.py: the Python-facing emd_grid_l1, with a dedicated 1D fast path (POT's own emd_1d_sorted for the plan; a closed-form, fully backend-native O(n) reduction for the cost-only case, so a 1D GPU array never leaves the device). Backend-compatible throughout; the sparse plan is returned as G via the backend's coo_matrix, same convention as ot.emd2_lazy's return_matrix. - test/test_grid.py: correctness against the dense solver, plan/coupling marginal checks, backend round-trips, and the 1D dispatch. Deliberately scoped to L1 only, with no gradient support yet -- see the PR description for benchmarks and open questions on L2 and batching. Co-Authored-By: Claude Sonnet 5 --- MANIFEST.in | 1 + RELEASES.md | 1 + ot/lp/EMD.h | 16 ++ ot/lp/EMD_wrapper.cpp | 254 ++++++++++++++++++++++++++++- ot/lp/__init__.py | 2 + ot/lp/_grid.py | 287 +++++++++++++++++++++++++++++++++ ot/lp/emd_wrap.pyx | 87 ++++++++++ ot/lp/sparse_digraph.h | 209 ++++++++++++++++++++++++ test/test_grid.py | 352 +++++++++++++++++++++++++++++++++++++++++ 9 files changed, 1208 insertions(+), 1 deletion(-) create mode 100644 ot/lp/_grid.py create mode 100644 ot/lp/sparse_digraph.h create mode 100644 test/test_grid.py diff --git a/MANIFEST.in b/MANIFEST.in index d48924e46..a6e4b5ca9 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -11,6 +11,7 @@ include ot/lp/full_bipartitegraph_omp.h include ot/lp/network_simplex_simple.h include ot/lp/network_simplex_simple_omp.h include ot/lp/sparse_bipartitegraph.h +include ot/lp/sparse_digraph.h include ot/partial/partial_cython.pyx include ot/bsp/BSP-OT_header_only.h include ot/bsp/bsp_wrapper.cpp diff --git a/RELEASES.md b/RELEASES.md index a83c8ba2d..10da4137d 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -4,6 +4,7 @@ #### New features +- Add `ot.lp.emd_grid_l1`, an exact EMD-L1 solver for histograms sharing a multi-dimensional Cartesian grid support. It solves a min-cost flow on the grid's adjacency graph rather than the full bipartite graph (Ling & Okada, 2007), which is one to two orders of magnitude faster for this case (PR #TODO) - Use `ot.utils.check_marginal` (and shape-tuple support in `ot.utils.unif`) to fill and validate default marginals consistently across solvers (Gromov, low-rank, stochastic, barycenter, factored) (PR #856) - Add stereographic spherical sliced Wasserstein distance in `ot.sliced.stereographic_sliced_wasserstein_sphere`, with its rotationally invariant extension (PR #836) - Add Quasi-Monte Carlo sliced Wasserstein sampling (QSW/RQSW) via generalized diff --git a/ot/lp/EMD.h b/ot/lp/EMD.h index 70ee01844..c8e2cae19 100644 --- a/ot/lp/EMD.h +++ b/ot/lp/EMD.h @@ -52,6 +52,22 @@ int EMD_wrap_sparse( double *beta_init // Initial dual variables for targets (warmstart) ); +int EMD_wrap_grid_l1( + int ndim, // Number of grid dimensions + int64_t *shape, // Grid shape (ndim entries) + double *X, // Source histogram, flattened C-order (prod(shape)) + double *Y, // Target histogram, flattened C-order (prod(shape)) + bool return_plan, // If false, skip decomposing the flow into a + // transportation plan and only compute cost + uint64_t *plan_sources_out, // Output: source bin index of each plan entry + uint64_t *plan_targets_out, // Output: target bin index of each plan entry + double *plan_values_out, // Output: mass moved by each plan entry + uint64_t *n_plan_entries_out, + uint64_t max_plan_entries, + double *cost, // Output: total transportation cost + uint64_t maxIter // Maximum iterations for solver +); + int EMD_wrap_lazy( int n1, // Number of source points int n2, // Number of target points diff --git a/ot/lp/EMD_wrapper.cpp b/ot/lp/EMD_wrapper.cpp index b6a572015..3e8ff89ea 100644 --- a/ot/lp/EMD_wrapper.cpp +++ b/ot/lp/EMD_wrapper.cpp @@ -15,7 +15,10 @@ #include "network_simplex_simple.h" #include "sparse_bipartitegraph.h" +#include "sparse_digraph.h" #include "EMD.h" +#include +#include #include #include #include @@ -202,6 +205,98 @@ inline bool extract_sparse_solution( return true; } +// An arc carrying positive flow, keyed by its head node. Used to decompose +// grid min-cost-flow arc flows (which move mass between *adjacent* grid +// cells) into a direct (source_bin, target_bin, mass) transport plan. +struct GridFlowEdge { + int head; + double flow; +}; + +// Decomposes multi-hop arc flows into direct (source, target, mass) entries. +// +// A min-cost flow on the grid adjacency graph reports how much mass crosses +// each arc between neighbouring cells, but a transport plan has to say which +// bin each unit of mass came from and where it ended up. This walks from +// each node with leftover supply along arcs that still carry flow until it +// reaches one with a deficit, records that path's bottleneck as one plan +// entry, and subtracts it from every arc on the path, repeating until no +// supply is left. `flow_adj` is consumed in place and `rem_supply` is taken +// by value, both as scratch. +// +// Returns false if the plan would exceed `max_plan_entries` entries. +inline bool decompose_grid_flows( + std::vector>& flow_adj, + std::vector rem_supply, + uint64_t* plan_sources_out, + uint64_t* plan_targets_out, + double* plan_values_out, + uint64_t* n_plan_entries_out, + uint64_t max_plan_entries +) { + const double eps = 1e-10; + const std::size_t n_nodes = flow_adj.size(); + std::vector ptr(n_nodes, 0); + + for (std::size_t src = 0; src < n_nodes; ++src) { + while (rem_supply[src] > eps) { + std::vector> path_edges; + std::size_t cur = src; + + while (true) { + if (rem_supply[cur] < -eps && cur != src) { + break; + } + auto& list = flow_adj[cur]; + std::size_t p = ptr[cur]; + while (p < list.size() && list[p].flow <= eps) { + ++p; + } + ptr[cur] = p; + if (p >= list.size()) { + break; + } + path_edges.emplace_back(cur, p); + cur = static_cast(list[p].head); + } + + if (path_edges.empty()) { + break; + } + + const std::size_t target = cur; + if (rem_supply[target] >= -eps) { + break; + } + + double bottleneck = rem_supply[src]; + bottleneck = std::min(bottleneck, -rem_supply[target]); + for (const auto& edge : path_edges) { + bottleneck = std::min(bottleneck, flow_adj[edge.first][edge.second].flow); + } + + if (bottleneck <= eps) { + break; + } + + for (const auto& edge : path_edges) { + flow_adj[edge.first][edge.second].flow -= bottleneck; + } + rem_supply[src] -= bottleneck; + rem_supply[target] += bottleneck; + + if (*n_plan_entries_out >= max_plan_entries) { + return false; + } + plan_sources_out[*n_plan_entries_out] = static_cast(src); + plan_targets_out[*n_plan_entries_out] = static_cast(target); + plan_values_out[*n_plan_entries_out] = bottleneck; + ++(*n_plan_entries_out); + } + } + return true; +} + } // namespace @@ -453,7 +548,164 @@ int EMD_wrap_sparse( return ret; } -int EMD_wrap_lazy(int n1, int n2, double *X, double *Y, double *coords_a, double *coords_b, +int EMD_wrap_grid_l1( + int ndim, + int64_t *shape, + double *X, + double *Y, + bool return_plan, + uint64_t *plan_sources_out, + uint64_t *plan_targets_out, + double *plan_values_out, + uint64_t *n_plan_entries_out, + uint64_t max_plan_entries, + double *cost, + uint64_t maxIter +) { + using namespace lemon; + + int64_t n_nodes = 1; + for (int d = 0; d < ndim; ++d) { + if (shape[d] <= 0) { + return INFEASIBLE; + } + n_nodes *= shape[d]; + } + + double total_x = 0.0; + double total_y = 0.0; + bool any_diff = false; + for (int64_t i = 0; i < n_nodes; ++i) { + if (X[i] < 0 || Y[i] < 0) { + return INFEASIBLE; + } + total_x += X[i]; + total_y += Y[i]; + any_diff = any_diff || (X[i] != Y[i]); + } + if (std::abs(total_x - total_y) > 1e-8 * std::max(1.0, total_x)) { + return INFEASIBLE; + } + + *cost = 0.0; + *n_plan_entries_out = 0; + + if (!any_diff) { + // Histograms are identical: nothing to transport, but if a plan is + // requested, the identity coupling is still the (trivially optimal) + // transportation plan. + if (return_plan) { + for (int64_t i = 0; i < n_nodes; ++i) { + if (X[i] > 1e-10) { + if (*n_plan_entries_out >= max_plan_entries) { + return (int)MAX_ITER_REACHED; + } + plan_sources_out[*n_plan_entries_out] = static_cast(i); + plan_targets_out[*n_plan_entries_out] = static_cast(i); + plan_values_out[*n_plan_entries_out] = X[i]; + ++(*n_plan_entries_out); + } + } + } + return OPTIMAL; + } + + // Grid-adjacent arcs: one forward and one backward arc per adjacent cell + // pair, unit cost each. On a unit-spaced Cartesian grid this reduces the + // cityblock-EMD problem to a min-cost flow on the grid graph, which is + // far sparser than the full bipartite graph (Ling & Okada, 2007). Unlike + // that paper's bespoke tree-based solver, the reduced graph below is + // handed to the off-the-shelf NetworkSimplexSimple LP solver. + std::vector stride(ndim); + stride[ndim - 1] = 1; + for (int d = ndim - 2; d >= 0; --d) { + stride[d] = stride[d + 1] * shape[d + 1]; + } + + std::vector> edges; + for (int d = 0; d < ndim; ++d) { + const int64_t extent = shape[d]; + if (extent < 2) { + continue; + } + const int64_t st = stride[d]; + for (int64_t u = 0; u < n_nodes; ++u) { + if ((u / st) % extent < extent - 1) { + edges.emplace_back(static_cast(u), static_cast(u + st)); + edges.emplace_back(static_cast(u + st), static_cast(u)); + } + } + } + + typedef SparseDigraph Digraph; + Digraph di(static_cast(n_nodes)); + di.buildFromEdges(edges); + const int64_t total_arcs = static_cast(edges.size()); + + std::vector supply(n_nodes); + for (int64_t i = 0; i < n_nodes; ++i) { + supply[i] = X[i] - Y[i]; + } + + typedef NetworkSimplexSimple Simplex; + Simplex::SimplexOptions simplex_options(true); + Simplex net(di, simplex_options, static_cast(n_nodes), total_arcs, maxIter); + net.supplyMap(supply); + for (int64_t k = 0; k < total_arcs; ++k) { + net.setCost(Digraph::arcFromId(k), 1.0); + } + + int ret = net.run(); + if (ret != (int)net.OPTIMAL && ret != (int)net.MAX_ITER_REACHED) { + return ret; + } + + *cost = net.totalCost(); + + if (!return_plan) { + // The caller only wants the cost: skip decomposing the Beckmann-style + // arc flow into a transportation plan (coupling) entirely. + return ret; + } + + // A bin's mass that already overlaps between X and Y needs no transport, + // so the min-cost flow above never routes it and the arc-flow + // decomposition below never reports it. Emit it directly as a same-bin + // plan entry so the plan is a genuine coupling (row sums X, column sums + // Y), not just the net residual. + for (int64_t i = 0; i < n_nodes; ++i) { + const double self_mass = std::min(X[i], Y[i]); + if (self_mass > 1e-10) { + if (*n_plan_entries_out >= max_plan_entries) { + return (int)net.MAX_ITER_REACHED; + } + plan_sources_out[*n_plan_entries_out] = static_cast(i); + plan_targets_out[*n_plan_entries_out] = static_cast(i); + plan_values_out[*n_plan_entries_out] = self_mass; + ++(*n_plan_entries_out); + } + } + + // Decompose the arc flow into a direct (source_bin, target_bin, mass) + // transportation plan. + std::vector> flow_adj(n_nodes); + for (int64_t k = 0; k < total_arcs; ++k) { + const Digraph::Arc a = Digraph::arcFromId(k); + const double f = net.flow(a); + if (f > 1e-10) { + flow_adj[di.source(a)].push_back({di.target(a), f}); + } + } + + if (!decompose_grid_flows(flow_adj, supply, plan_sources_out, plan_targets_out, + plan_values_out, n_plan_entries_out, max_plan_entries)) { + return (int)net.MAX_ITER_REACHED; + } + + return ret; +} + +int EMD_wrap_lazy(int n1, int n2, double *X, double *Y, double *coords_a, double *coords_b, int dim, int metric, uint64_t *flow_sources_out, uint64_t *flow_targets_out, double *flow_values_out, uint64_t *n_flows_out, uint64_t max_flows_out, diff --git a/ot/lp/__init__.py b/ot/lp/__init__.py index c9fa676c4..e7f1356dd 100644 --- a/ot/lp/__init__.py +++ b/ot/lp/__init__.py @@ -10,6 +10,7 @@ from .dmmot import dmmot_monge_1dgrid_loss, dmmot_monge_1dgrid_optimize from ._network_simplex import emd, emd2, emd2_lazy +from ._grid import emd_grid_l1 from ._barycenter_solvers import ( barycenter, free_support_barycenter, @@ -40,6 +41,7 @@ "emd", "emd2", "emd2_lazy", + "emd_grid_l1", "barycenter", "free_support_barycenter", "cvx", diff --git a/ot/lp/_grid.py b/ot/lp/_grid.py new file mode 100644 index 000000000..717a96b94 --- /dev/null +++ b/ot/lp/_grid.py @@ -0,0 +1,287 @@ +""" +Exact EMD-L1 solver for histograms sharing a Cartesian grid support. +""" + +# Author: Tom Vercauteren +# +# License: MIT License + +import numpy as np + +from ..backend import get_backend +from ..utils import list_to_array +from .emd_wrap import check_result, emd_1d_sorted, emd_c_grid_l1 + +# Mirrors ot::ProblemType in ot/lp/EMD.h (INFEASIBLE=0, OPTIMAL=1, ...). Not +# otherwise exposed to Python, since it is a Cython ``cdef enum``. +_RESULT_INFEASIBLE = 0 +_RESULT_OPTIMAL = 1 + +_EMPTY_U64 = np.empty(0, dtype=np.uint64) +_EMPTY_F64 = np.empty(0, dtype=np.float64) + + +def _finalize_native_cost(cost, nx): + """For the numpy backend, ``nx.sum`` etc. return a numpy scalar (e.g. + ``np.float64``), whereas every other POT solver returns a plain Python + float in that case (e.g. numpy scalars repr as ``np.float64(3.0)``, + breaking naive equality/display expectations). Normalize to a plain + float there; other backends keep their native scalar tensor as is. + """ + if isinstance(cost, np.generic): + return float(cost) + return nx.detach(cost) + + +def _emd_grid_l1_1d_cost(A, B, nx): + r"""Closed-form, fully backend-native cost for a 1D grid. + + A 1D grid with a shared, sorted, unit-spaced support has a classic + :math:`\mathcal{O}(n)` closed form for the L1 (cityblock) Wasserstein + cost: the L1 norm of the difference of cumulative sums. Every step is a + generic backend reduction (``nx.cumsum``, ``nx.abs``, ``nx.sum``), so + this never leaves the device `A`/`B` are already on -- unlike every + other path in :any:`emd_grid_l1`, which needs a CPU round-trip. + + Returns ``(cost, result_code)``; `cost` is already a backend-native + scalar with `A`'s dtype and device, detached from any autodiff graph + (see :any:`_finalize_native_cost`). + """ + if A.shape[0] == 0 or nx.any(A < 0) or nx.any(B < 0): + return _finalize_native_cost(0.0 * nx.sum(A), nx), _RESULT_INFEASIBLE + + total_a = nx.sum(A) + total_b = nx.sum(B) + if abs(float(nx.to_numpy(total_a)) - float(nx.to_numpy(total_b))) > 1e-8 * max( + 1.0, float(nx.to_numpy(total_a)) + ): + return _finalize_native_cost(0.0 * total_a, nx), _RESULT_INFEASIBLE + + cum_diff = nx.cumsum(A, 0) - nx.cumsum(B, 0) + cost = nx.sum(nx.abs(cum_diff[:-1])) + return _finalize_native_cost(cost, nx), _RESULT_OPTIMAL + + +def _emd_grid_l1_1d_plan(A, B, nx): + r"""Cost and sparse transportation plan for a 1D grid. + + A 1D grid is just a sorted, shared support, for which POT already has an + exact :math:`\mathcal{O}(n)` solver (:any:`ot.lp.emd_1d_sorted`, the same + routine backing :any:`ot.lp.emd_1d`). Calling it directly here, with the + grid's integer positions passed in as already sorted, skips both the + network simplex setup of the general grid solver and the argsort/pairwise + distance overhead that the generic, arbitrary-support :any:`ot.lp.emd_1d` + would otherwise pay. + + Unlike :any:`_emd_grid_l1_1d_cost`, recovering the actual coupling needs + the compiled O(n) merge behind :any:`ot.lp.emd_1d_sorted`, which requires + plain CPU numpy arrays -- so, backend-compatible via `nx`, this converts + `A`/`B` right here, immediately before that one call. + """ + a = np.ascontiguousarray(nx.to_numpy(A), dtype=np.float64) + b = np.ascontiguousarray(nx.to_numpy(B), dtype=np.float64) + n = a.shape[0] + + if n == 0 or np.any(a < 0) or np.any(b < 0): + return _EMPTY_U64, _EMPTY_U64, _EMPTY_F64, 0.0, _RESULT_INFEASIBLE + + total_a = a.sum() + total_b = b.sum() + if abs(total_a - total_b) > 1e-8 * max(1.0, total_a): + return _EMPTY_U64, _EMPTY_U64, _EMPTY_F64, 0.0, _RESULT_INFEASIBLE + + x = np.arange(n, dtype=np.float64) + plan_values, indices, cost = emd_1d_sorted(a, b, x, x, metric="cityblock") + + # The merge algorithm behind emd_1d_sorted can emit exact-zero-mass + # entries at ties (e.g. two identical histograms); drop them so the plan + # only lists actual mass movements, matching the general grid solver. + nonzero = plan_values > 0 + plan_sources = np.ascontiguousarray(indices[nonzero, 0], dtype=np.uint64) + plan_targets = np.ascontiguousarray(indices[nonzero, 1], dtype=np.uint64) + plan_values = np.ascontiguousarray(plan_values[nonzero], dtype=np.float64) + return plan_sources, plan_targets, plan_values, cost, _RESULT_OPTIMAL + + +def emd_grid_l1( + A, B, numItermax=100000, return_plan=False, log=False, check_marginals=True +): + r"""Solves the Earth Mover's Distance with the cityblock ground metric + between two histograms sharing the same :math:`d`-dimensional Cartesian + grid support. + + .. math:: + \gamma = \mathop{\arg \min}_\gamma \quad \langle \gamma, \mathbf{M} \rangle_F + + s.t. \ \gamma \mathbf{1} = \mathbf{a} + + \gamma^T \mathbf{1} = \mathbf{b} + + \gamma \geq 0 + + where :math:`M_{i,j}` is the cityblock (:math:`\ell_1`) distance between + grid nodes :math:`i` and :math:`j` on the integer grid of shape + `A.shape`, and :math:`\mathbf{a}`, :math:`\mathbf{b}` are `A`, `B` + flattened in C order. + + Rather than solving the min-cost flow on the full bipartite graph between + the :math:`n=\prod(\text{A.shape})` source and target bins (as + :any:`ot.emd` would), this reduces the problem to a min-cost flow on the + much sparser grid adjacency graph (:math:`\mathcal{O}(d \cdot n)` arcs + instead of :math:`\mathcal{O}(n^2)`), using the graph formulation of Ling + & Okada [1]_. Unlike [1]_, which introduces a bespoke tree-based solver + for that reduced graph, this reuses POT's existing (off-the-shelf) + network simplex LP solver on it. This is exact and typically one to two + orders of magnitude faster than :any:`ot.emd` for grid histograms. + + .. note:: The min-cost flow solved internally is a Beckmann-style flow on + the grid's adjacency graph: it moves mass between *neighbouring* + cells and does not by itself say which source bin any given unit of + mass originally came from. Recovering the actual transportation plan + :math:`\gamma` (a coupling between source and target bins) requires + decomposing that flow into paths, which has a non-negligible cost of + its own. Set `return_plan` to request it; leave it False (the + default) when only the transport cost is needed. + + .. note:: For a 1D grid (``A.ndim == 1``), this instead delegates to + POT's :math:`\mathcal{O}(n)` sorted-support solver + (:any:`ot.lp.emd_1d_sorted`), which is exact here since the grid's + integer positions are already a shared, sorted support. This avoids + the network simplex setup entirely for that case. Better still, when + `return_plan` is False, the cost itself has a closed form that runs + as generic backend reductions with no CPU round-trip at all, so a 1D + grid on a GPU array is solved entirely on-device. + + .. note:: This function is backend-compatible and will work on arrays + from all compatible backends. Beyond the 1D, cost-only case above, + the algorithm uses a C++/Cython CPU solver, so GPU arrays are copied + to CPU before solving (and the transportation plan's bin indices, if + requested, are returned as CPU arrays). Gradients are not currently + supported: `cost` is always detached from any computation graph the + inputs were part of. + + Parameters + ---------- + A : array-like, float64, shape (n_1, ..., n_d) + Source histogram on a :math:`d`-dimensional Cartesian grid + B : array-like, float64, shape (n_1, ..., n_d) + Target histogram, on the same grid as `A` + numItermax : int, optional (default=100000) + The maximum number of iterations before stopping the optimization + algorithm if it has not converged. + return_plan : bool, optional (default=False) + If True, additionally recovers an explicit transportation plan (a + sparse coupling), returned in `log`. Computing it has a cost of its + own (a CPU round-trip, and either a network simplex solve or an O(n) + merge), so it is skipped by default when only the transport cost + `cost` is needed. + log : bool, optional (default=False) + If True, also returns a dictionary with the solver status and, + if `return_plan` is True, the sparse transportation plan. + check_marginals : bool, optional (default=True) + If True, checks that `A` and `B` have the same total mass. + + Returns + ------- + cost : float + Optimal transportation cost. + log : dict, optional + If input `log` is True, a dictionary containing the solver status + (`warning`, `result_code`) and, if `return_plan` is True, the sparse + transportation plan `G` (built via the backend's `coo_matrix`, same + as :any:`ot.emd2_lazy`'s `return_matrix`; a real sparse matrix for + NumPy/PyTorch/TensorFlow/CuPy, silently densified for JAX, which has + no sparse array type) of shape :math:`(n, n)` with + :math:`n=\prod(\text{A.shape})`, indexing into `A.reshape(-1)` and + `B.reshape(-1)`. + + Examples + -------- + >>> import numpy as np + >>> A = np.array([1.0, 0.0, 0.0, 0.0]) + >>> B = np.array([0.0, 0.0, 0.0, 1.0]) + >>> emd_grid_l1(A, B) + 3.0 + + References + ---------- + .. [1] Ling, H., & Okada, K. (2007). An efficient earth mover's distance + algorithm for robust histogram comparison. IEEE Transactions on + Pattern Analysis and Machine Intelligence, 29(5), 840-853. + + See Also + -------- + ot.emd : Exact OT solver with a general, precomputed cost matrix + """ + A, B = list_to_array(A, B) + nx = get_backend(A, B) + + if A.shape != B.shape: + raise ValueError( + f"A and B must have the same shape, got {A.shape} and {B.shape}" + ) + + if check_marginals: + np.testing.assert_allclose( + nx.to_numpy(nx.sum(A)), + nx.to_numpy(nx.sum(B)), + rtol=1e-7, + atol=1e-7, + err_msg="A and B must have the same total mass", + ) + + if A.ndim == 1: + # A 1D grid is backend-native either way (works on any backend's + # arrays via `nx`, GPU included). Only the cost-only case below + # avoids a CPU round-trip entirely though: recovering the plan needs + # the compiled O(n) merge in _emd_grid_l1_1d_plan, which is CPU-only. + if not return_plan: + cost, result_code = _emd_grid_l1_1d_cost(A, B, nx) + if log: + return cost, { + "warning": check_result(result_code), + "result_code": result_code, + } + check_result(result_code) + return cost + + plan_sources, plan_targets, plan_values, cost, result_code = ( + _emd_grid_l1_1d_plan(A, B, nx) + ) + else: + shape = np.array(A.shape, dtype=np.int64) + # The C++ solver only understands flattened (CPU) numpy arrays: + # `to_numpy` also does the GPU -> CPU copy for backends such as + # torch or jax. + a_np = np.ascontiguousarray(nx.to_numpy(A), dtype=np.float64).reshape(-1) + b_np = np.ascontiguousarray(nx.to_numpy(B), dtype=np.float64).reshape(-1) + plan_sources, plan_targets, plan_values, cost, result_code = emd_c_grid_l1( + a_np, b_np, shape, numItermax, return_plan + ) + + cost = nx.from_numpy(cost, type_as=A) + + if log: + log_dict = { + "warning": check_result(result_code), + "result_code": result_code, + } + if return_plan: + # A.size is a numpy property but a torch method: go through + # A.shape (uniformly a tuple of ints/plain integers across + # backends) instead. + n = int(np.prod(A.shape)) + plan_values_b = nx.from_numpy(plan_values, type_as=A) + plan_sources_b = nx.from_numpy(plan_sources.astype(np.int64), type_as=A) + plan_targets_b = nx.from_numpy(plan_targets.astype(np.int64), type_as=A) + log_dict["G"] = nx.coo_matrix( + plan_values_b, + plan_sources_b, + plan_targets_b, + shape=(n, n), + type_as=A, + ) + return cost, log_dict + + check_result(result_code) + return cost diff --git a/ot/lp/emd_wrap.pyx b/ot/lp/emd_wrap.pyx index 64905efe9..5a9ddfc7f 100644 --- a/ot/lp/emd_wrap.pyx +++ b/ot/lp/emd_wrap.pyx @@ -22,6 +22,7 @@ import warnings cdef extern from "EMD.h": int EMD_wrap(int n1,int n2, double *X, double *Y,double *D, double *G, double* alpha, double* beta, double *cost, uint64_t maxIter, double* alpha_init, double* beta_init) nogil int EMD_wrap_sparse(int n1, int n2, double *X, double *Y, uint64_t n_edges, uint64_t *edge_sources, uint64_t *edge_targets, double *edge_costs, uint64_t *flow_sources_out, uint64_t *flow_targets_out, double *flow_values_out, uint64_t *n_flows_out, uint64_t max_flows_out, double *alpha, double *beta, double *cost, uint64_t maxIter, double* alpha_init, double* beta_init) nogil + int EMD_wrap_grid_l1(int ndim, int64_t *shape, double *X, double *Y, bint return_plan, uint64_t *plan_sources_out, uint64_t *plan_targets_out, double *plan_values_out, uint64_t *n_plan_entries_out, uint64_t max_plan_entries, double *cost, uint64_t maxIter) nogil int EMD_wrap_lazy(int n1, int n2, double *X, double *Y, double *coords_a, double *coords_b, int dim, int metric, uint64_t *flow_sources_out, uint64_t *flow_targets_out, double *flow_values_out, uint64_t *n_flows_out, uint64_t max_flows_out, double* alpha, double* beta, double *cost, uint64_t maxIter, double* alpha_init, double* beta_init) nogil cdef enum ProblemType: INFEASIBLE, OPTIMAL, UNBOUNDED, MAX_ITER_REACHED @@ -367,3 +368,89 @@ def emd_c_lazy(np.ndarray[double, ndim=1, mode="c"] a, np.ndarray[double, ndim=1 flow_values = flow_values[:n_flows_out] return flow_sources, flow_targets, flow_values, cost, alpha, beta, result_code + + +@cython.boundscheck(False) +@cython.wraparound(False) +def emd_c_grid_l1(np.ndarray[double, ndim=1, mode="c"] a, + np.ndarray[double, ndim=1, mode="c"] b, + np.ndarray[int64_t, ndim=1, mode="c"] shape, + uint64_t max_iter, + bint return_plan=False): + """ + Grid EMD-L1 solver. + + Solves the Earth Mover's Distance with the cityblock ground metric between + two histograms sharing the same Cartesian grid support, by running a + min-cost flow on the grid's adjacency graph instead of the full bipartite + graph, using the graph formulation of Ling & Okada (2007). Unlike that + paper's bespoke tree-based solver, this reuses the off-the-shelf + NetworkSimplexSimple LP solver on the reduced graph. This is exact and + typically one to two orders of magnitude faster than the dense or lazy + solvers for this case. + + The min-cost flow itself is a Beckmann-style flow on the grid's adjacency + graph: it does not directly say which source bin each unit of mass came + from. When `return_plan` is True, that arc flow is decomposed (at some + extra cost) into an explicit transportation plan, i.e. a coupling given as + sparse `(source_bin, target_bin, mass)` entries. When only the transport + cost is needed, leave `return_plan` False to skip this decomposition. + + Parameters + ---------- + a : (n,) array, float64 + Source histogram, flattened in C order, with n = prod(shape) + b : (n,) array, float64 + Target histogram, flattened in C order, with n = prod(shape) + shape : (ndim,) array, int64 + Grid shape + max_iter : uint64_t + Maximum number of iterations + return_plan : bool, optional (default=False) + If True, also decompose the solver's arc flow into a sparse + transportation plan. If False, only the cost is computed. + + Returns + ------- + plan_sources : (n_plan_entries,) array, uint64 + Flattened source bin index of each transportation plan entry (empty + if `return_plan` is False) + plan_targets : (n_plan_entries,) array, uint64 + Flattened target bin index of each transportation plan entry (empty + if `return_plan` is False) + plan_values : (n_plan_entries,) array, float64 + Mass moved by each transportation plan entry (empty if `return_plan` + is False) + cost : float + Total transportation cost + result_code : int + Result status + """ + cdef int ndim = shape.shape[0] + cdef uint64_t n_plan_entries_out = 0 + cdef int result_code = 0 + cdef double cost = 0 + # Up to one same-bin plan entry per node for the overlapping ("self") + # mass, plus a path decomposition of the residual flow, which uses at + # most one plan entry per node plus one per (bidirectional) grid arc + # that ends up carrying flow. + cdef uint64_t max_plan_entries = a.shape[0] * (2 * ndim + 2) if return_plan else 0 + + cdef np.ndarray[uint64_t, ndim=1, mode="c"] plan_sources = np.zeros(max_plan_entries, dtype=np.uint64) + cdef np.ndarray[uint64_t, ndim=1, mode="c"] plan_targets = np.zeros(max_plan_entries, dtype=np.uint64) + cdef np.ndarray[double, ndim=1, mode="c"] plan_values = np.zeros(max_plan_entries, dtype=np.float64) + + with nogil: + result_code = EMD_wrap_grid_l1( + ndim, shape.data, + a.data, b.data, + return_plan, + plan_sources.data, plan_targets.data, plan_values.data, + &n_plan_entries_out, max_plan_entries, &cost, max_iter + ) + + plan_sources = plan_sources[:n_plan_entries_out] + plan_targets = plan_targets[:n_plan_entries_out] + plan_values = plan_values[:n_plan_entries_out] + + return plan_sources, plan_targets, plan_values, cost, result_code diff --git a/ot/lp/sparse_digraph.h b/ot/lp/sparse_digraph.h new file mode 100644 index 000000000..b4e1872e4 --- /dev/null +++ b/ot/lp/sparse_digraph.h @@ -0,0 +1,209 @@ +/* -*- mode: C++; indent-tabs-mode: nil; -*- + * + * General (non-bipartite) sparse directed graph for optimal transport. + * + * Unlike SparseBipartiteDigraph (see sparse_bipartitegraph.h), nodes are not + * split into a source half and a target half: any node may carry supply or + * demand. This is required for min-cost-flow formulations on grid graphs + * (e.g. EMD-L1 on Cartesian grids), where every cell is both a potential + * source and a potential sink. + * + * Uses CSR (Compressed Sparse Row) format for cache-friendly arc iteration. + * Requires edges to be provided in sorted order during construction (sorting + * is done internally by buildFromEdges). + */ + +#pragma once + +#include "core.h" +#include +#include +#include +#include +#include + +namespace lemon { + + class SparseDigraph { + public: + + typedef SparseDigraph Digraph; + typedef int Node; + typedef int64_t Arc; + + private: + + int _node_num; + int64_t _arc_num; + + std::vector _arc_sources; // _arc_sources[arc_id] = source node + std::vector _arc_targets; // _arc_targets[arc_id] = target node + + std::vector _row_ptr; // CSR row pointers (size: node_num+1) + std::vector _arc_ids; // arc IDs in source order + + mutable std::vector> _in_arcs; // _in_arcs[node] = incoming arc IDs + mutable bool _in_arcs_built; + + mutable std::vector _arc_to_out_pos; // _arc_to_out_pos[arc_id] = position in _arc_ids + mutable std::vector _arc_to_in_pos; // _arc_to_in_pos[arc_id] = position in _in_arcs[target] + mutable bool _position_maps_built; + + void build_in_arcs() const { + if (_in_arcs_built) return; + + _in_arcs.resize(_node_num); + for (Arc a = 0; a < _arc_num; ++a) { + _in_arcs[_arc_targets[a]].push_back(a); + } + _in_arcs_built = true; + } + + void build_position_maps() const { + if (_position_maps_built) return; + + _arc_to_out_pos.resize(_arc_num); + _arc_to_in_pos.resize(_arc_num); + + for (int64_t pos = 0; pos < _arc_num; ++pos) { + _arc_to_out_pos[_arc_ids[pos]] = pos; + } + + build_in_arcs(); + for (int node = 0; node < _node_num; ++node) { + const std::vector& in = _in_arcs[node]; + for (size_t pos = 0; pos < in.size(); ++pos) { + _arc_to_in_pos[in[pos]] = pos; + } + } + + _position_maps_built = true; + } + + public: + + explicit SparseDigraph(int n) + : _node_num(n), _arc_num(0), + _in_arcs_built(false), _position_maps_built(false) {} + + void buildFromEdges(const std::vector>& edges) { + _arc_num = edges.size(); + _arc_sources.resize(_arc_num); + _arc_targets.resize(_arc_num); + _arc_ids.resize(_arc_num); + _in_arcs_built = false; + _position_maps_built = false; + _in_arcs.clear(); + _arc_to_out_pos.clear(); + _arc_to_in_pos.clear(); + + // Create indexed edges: (source, target, original_arc_id) + std::vector> indexed_edges; + indexed_edges.reserve(_arc_num); + for (Arc i = 0; i < _arc_num; ++i) { + indexed_edges.emplace_back(edges[i].first, edges[i].second, i); + } + + // Sort by source node, then by target node (CSR requirement) + std::sort(indexed_edges.begin(), indexed_edges.end(), + [](const auto& a, const auto& b) { + if (std::get<0>(a) != std::get<0>(b)) + return std::get<0>(a) < std::get<0>(b); + return std::get<1>(a) < std::get<1>(b); + }); + + _row_ptr.assign(_node_num + 1, 0); + int current_row = 0; + + for (int64_t i = 0; i < _arc_num; ++i) { + Node src = std::get<0>(indexed_edges[i]); + Node tgt = std::get<1>(indexed_edges[i]); + Arc orig_arc_id = std::get<2>(indexed_edges[i]); + + // Fill out row_ptr for rows with no outgoing edges + while (current_row < src) { + _row_ptr[++current_row] = i; + } + + _arc_sources[orig_arc_id] = src; + _arc_targets[orig_arc_id] = tgt; + _arc_ids[i] = orig_arc_id; + } + + // Fill remaining row_ptr entries + while (current_row < _node_num) { + _row_ptr[++current_row] = _arc_num; + } + } + + int nodeNum() const { return _node_num; } + int64_t arcNum() const { return _arc_num; } + + int maxNodeId() const { return _node_num - 1; } + int64_t maxArcId() const { return _arc_num - 1; } + + Node source(Arc arc) const { return _arc_sources[arc]; } + Node target(Arc arc) const { return _arc_targets[arc]; } + + static int id(Node node) { return node; } + static int64_t id(Arc arc) { return arc; } + + static Node nodeFromId(int id) { return Node(id); } + static Arc arcFromId(int64_t id) { return Arc(id); } + + void first(Node& node) const { node = _node_num - 1; } + static void next(Node& node) { --node; } + + void first(Arc& arc) const { arc = _arc_num - 1; } + static void next(Arc& arc) { --arc; } + + void firstOut(Arc& arc, const Node& node) const { + if (node < 0 || node >= _node_num) { + arc = -1; + return; + } + + int64_t start = _row_ptr[node]; + int64_t end = _row_ptr[node + 1]; + + arc = (start < end) ? _arc_ids[start] : Arc(-1); + } + + void nextOut(Arc& arc) const { + if (arc < 0) return; + + build_position_maps(); + + int64_t pos = _arc_to_out_pos[arc]; + Node src = _arc_sources[arc]; + int64_t end = _row_ptr[src + 1]; + + arc = (pos + 1 < end) ? _arc_ids[pos + 1] : Arc(-1); + } + + void firstIn(Arc& arc, const Node& node) const { + build_in_arcs(); + + if (node < 0 || node >= _node_num) { + arc = -1; + return; + } + + const std::vector& in = _in_arcs[node]; + arc = in.empty() ? Arc(-1) : in[0]; + } + + void nextIn(Arc& arc) const { + if (arc < 0) return; + + build_position_maps(); + + int64_t pos = _arc_to_in_pos[arc]; + Node tgt = _arc_targets[arc]; + const std::vector& in = _in_arcs[tgt]; + + arc = (pos + 1 < static_cast(in.size())) ? in[pos + 1] : Arc(-1); + } + }; + +} //namespace lemon diff --git a/test/test_grid.py b/test/test_grid.py new file mode 100644 index 000000000..af46e3365 --- /dev/null +++ b/test/test_grid.py @@ -0,0 +1,352 @@ +"""Tests for the EMD-L1 grid solver (ot.lp.emd_grid_l1)""" + +# Author: Tom Vercauteren +# +# License: MIT License + +import itertools +from unittest import mock + +import numpy as np +import pytest + +import ot +from ot.lp import emd_grid_l1 +from ot.lp._grid import _emd_grid_l1_1d_cost, _emd_grid_l1_1d_plan + + +def _grid_coords(shape): + axes = [np.arange(s) for s in shape] + return np.array(np.meshgrid(*axes, indexing="ij")).reshape(len(shape), -1).T + + +def _dense_cost_and_check(a, b, shape): + coords = _grid_coords(shape) + M = ot.dist(coords, coords, metric="cityblock") + return ot.emd2(a, b, M) + + +def _check_plan(G, a, b, shape, cost, nx): + """G must be a genuine coupling: nonnegative, row sums a, column sums b, + and total mass-times-distance equal to the reported cost.""" + dense = nx.to_numpy(nx.todense(G)) + assert np.all(dense >= -1e-12) + np.testing.assert_allclose(dense.sum(axis=1), a, atol=1e-8) + np.testing.assert_allclose(dense.sum(axis=0), b, atol=1e-8) + + coords = _grid_coords(shape) + M = ot.dist(coords, coords, metric="cityblock") + np.testing.assert_allclose((dense * M).sum(), cost, rtol=1e-6, atol=1e-8) + + +@pytest.mark.parametrize( + "shape", + [ + (4,), + (5,), + (2, 3), + (3, 4), + (2, 2, 2), + (2, 3, 2), + ], +) +def test_emd_grid_l1_vs_dense(shape): + """The grid solver must match the dense bipartite solver exactly.""" + rng = np.random.RandomState(42) + n = int(np.prod(shape)) + a = rng.rand(n) + a /= a.sum() + b = rng.rand(n) + b /= b.sum() + + cost_dense = _dense_cost_and_check(a, b, shape) + cost_grid = emd_grid_l1(a.reshape(shape), b.reshape(shape)) + + np.testing.assert_allclose(cost_grid, cost_dense, rtol=1e-6, atol=1e-8) + + +@pytest.mark.parametrize("shape", [(4,), (3, 3), (2, 3, 2)]) +def test_emd_grid_l1_identical_histograms(shape): + """Identical histograms cost 0; the plan, if requested, is the (trivially + optimal) identity coupling, not an empty one.""" + n = int(np.prod(shape)) + rng = np.random.RandomState(0) + a = rng.rand(n) + a /= a.sum() + A = a.reshape(shape) + nx = ot.backend.NumpyBackend() + + cost, log = emd_grid_l1(A, A, return_plan=True, log=True) + + assert cost == 0.0 + _check_plan(log["G"], a, a, shape, cost, nx) + np.testing.assert_allclose(nx.to_numpy(nx.todense(log["G"])).sum(), 1.0) + + +def test_emd_grid_l1_return_plan_default_false(): + """By default, no plan is computed or returned (cost-only fast path).""" + shape = (3, 3) + n = int(np.prod(shape)) + rng = np.random.RandomState(1) + a = rng.rand(n) + a /= a.sum() + b = rng.rand(n) + b /= b.sum() + A, B = a.reshape(shape), b.reshape(shape) + + cost_no_log = emd_grid_l1(A, B) + cost, log = emd_grid_l1(A, B, log=True) + + assert cost == cost_no_log + assert "G" not in log + + # Explicitly requesting the plan must add it to the log, with a matching + # cost. + cost_with_plan, log_with_plan = emd_grid_l1(A, B, return_plan=True, log=True) + assert cost_with_plan == cost + assert "G" in log_with_plan + + +def test_emd_grid_l1_one_hot(): + """Point mass moving between two corners costs their Manhattan distance.""" + shape = (2, 3, 4) + A = np.zeros(shape) + B = np.zeros(shape) + A[0, 0, 0] = 1.0 + B[1, 2, 3] = 1.0 + + cost = emd_grid_l1(A, B) + np.testing.assert_allclose(cost, 1 + 2 + 3) + + +def test_emd_grid_l1_plan_marginals(): + """The transportation plan must be a genuine coupling: row sums A, + column sums B, including the mass A and B already share at the same + bin (which the underlying flow decomposition alone would miss).""" + shape = (3, 4) + n = int(np.prod(shape)) + rng = np.random.RandomState(7) + a = rng.rand(n) + a /= a.sum() + b = rng.rand(n) + b /= b.sum() + nx = ot.backend.NumpyBackend() + + cost, log = emd_grid_l1( + a.reshape(shape), b.reshape(shape), return_plan=True, log=True + ) + + _check_plan(log["G"], a, b, shape, cost, nx) + + +def test_emd_grid_l1_mass_mismatch(): + shape = (2, 2) + A = np.array([1.0, 0.0, 0.0, 0.0]).reshape(shape) + B = np.array([0.0, 0.0, 0.0, 0.5]).reshape(shape) + + with pytest.raises(AssertionError): + emd_grid_l1(A, B) + + # Explicitly skipping the check should not raise from the Python side. + _cost, log = emd_grid_l1(A, B, check_marginals=False, log=True) + assert log["result_code"] != 1 # not OPTIMAL: infeasible + + +def test_emd_grid_l1_shape_mismatch(): + a = np.ones(4) / 4 + b = np.ones(5) / 5 + with pytest.raises(ValueError): + emd_grid_l1(a, b) # different total number of elements + + # Same total number of elements, but a different grid shape/structure + # must still raise: the grid geometry, not just the element count, has + # to match. + A = np.ones((2, 3)) / 6 + B = np.ones((3, 2)) / 6 + with pytest.raises(ValueError): + emd_grid_l1(A, B) + + +@pytest.mark.parametrize("ndim", [1, 2, 3]) +def test_emd_grid_l1_random_grids_batch(ndim): + """Broader randomized cross-check across many small grid shapes.""" + rng = np.random.RandomState(123) + for extents in itertools.product(range(2, 4), repeat=ndim): + n = int(np.prod(extents)) + a = rng.rand(n) + a /= a.sum() + b = rng.rand(n) + b /= b.sum() + + cost_dense = _dense_cost_and_check(a, b, extents) + cost_grid = emd_grid_l1(a.reshape(extents), b.reshape(extents)) + np.testing.assert_allclose(cost_grid, cost_dense, rtol=1e-6, atol=1e-8) + + +def test_emd_grid_l1_1d_uses_native_cost_path(): + """A 1D grid with return_plan=False must skip both the general C++ + solver and the O(n) merge that recovers the plan.""" + a = np.array([1.0, 0.0, 0.0, 0.0]) + b = np.array([0.0, 0.0, 0.0, 1.0]) + + with ( + mock.patch("ot.lp._grid.emd_c_grid_l1") as mocked_cpp, + mock.patch( + "ot.lp._grid._emd_grid_l1_1d_plan", + wraps=_emd_grid_l1_1d_plan, + ) as mocked_plan, + ): + cost = emd_grid_l1(a, b) + mocked_cpp.assert_not_called() + mocked_plan.assert_not_called() + np.testing.assert_allclose(cost, 3.0) + + # Requesting the plan on a 1D grid must still avoid the C++ solver, but + # does need the O(n) merge. + with ( + mock.patch("ot.lp._grid.emd_c_grid_l1") as mocked_cpp, + mock.patch( + "ot.lp._grid._emd_grid_l1_1d_plan", + wraps=_emd_grid_l1_1d_plan, + ) as mocked_plan, + ): + emd_grid_l1(a, b, return_plan=True) + mocked_cpp.assert_not_called() + mocked_plan.assert_called_once() + + # A 2D grid, by contrast, must go through the general C++ solver. + with mock.patch( + "ot.lp._grid.emd_c_grid_l1", wraps=ot.lp._grid.emd_c_grid_l1 + ) as mocked_cpp: + emd_grid_l1(a.reshape(2, 2), b.reshape(2, 2)) + mocked_cpp.assert_called_once() + + +@pytest.mark.parametrize("n", [1, 2, 5, 8]) +def test_emd_grid_l1_1d_vs_dense(n): + """Cross-check the dedicated 1D path against the dense solver directly, + beyond what test_emd_grid_l1_vs_dense already covers.""" + rng = np.random.RandomState(11) + a = rng.rand(n) + a /= a.sum() + b = rng.rand(n) + b /= b.sum() + + cost_dense = _dense_cost_and_check(a, b, (n,)) + cost_grid = emd_grid_l1(a, b) + np.testing.assert_allclose(cost_grid, cost_dense, rtol=1e-6, atol=1e-8) + + +def test_emd_grid_l1_1d_plan_marginals(): + n = 6 + rng = np.random.RandomState(5) + a = rng.rand(n) + a /= a.sum() + b = rng.rand(n) + b /= b.sum() + nx = ot.backend.NumpyBackend() + + cost, log = emd_grid_l1(a, b, return_plan=True, log=True) + + _check_plan(log["G"], a, b, (n,), cost, nx) + + +def test_emd_grid_l1_1d_identical_histograms(): + """Identical histograms cost 0 and the plan is the identity coupling, + even though the underlying merge could in principle emit zero-mass + entries at ties.""" + a = np.array([0.25, 0.25, 0.25, 0.25]) + nx = ot.backend.NumpyBackend() + + cost, log = emd_grid_l1(a, a, return_plan=True, log=True) + + assert cost == 0.0 + dense = nx.to_numpy(nx.todense(log["G"])) + assert np.all(np.diag(dense) > 0) + np.testing.assert_allclose(dense, np.diag(a)) + + +def test_emd_grid_l1_1d_mass_mismatch(): + a = np.array([1.0, 0.0, 0.0, 0.0]) + b = np.array([0.0, 0.0, 0.0, 0.5]) + + with pytest.raises(AssertionError): + emd_grid_l1(a, b) + + _cost, log = emd_grid_l1(a, b, check_marginals=False, log=True) + assert log["result_code"] != 1 # not OPTIMAL: infeasible + + +def test_emd_grid_l1_1d_direct_cost_helper(): + nx = ot.backend.NumpyBackend() + + a = np.array([1.0, 0.0, 0.0, 0.0]) + b = np.array([0.0, 0.0, 0.0, 1.0]) + cost, result_code = _emd_grid_l1_1d_cost(a, b, nx) + assert result_code == 1 # OPTIMAL + np.testing.assert_allclose(cost, 3.0) + + # Negative values and mass mismatches must be reported as infeasible, + # like the general (C++) path. + a_neg = np.array([1.0, -0.1, 0.0, 0.0]) + _cost, result_code = _emd_grid_l1_1d_cost(a_neg, b, nx) + assert result_code != 1 + + a_mismatch = np.array([1.0, 0.0, 0.0, 0.0]) + b_mismatch = np.array([0.0, 0.0, 0.0, 0.5]) + _cost, result_code = _emd_grid_l1_1d_cost(a_mismatch, b_mismatch, nx) + assert result_code != 1 + + +def test_emd_grid_l1_1d_direct_plan_helper_negative_values(): + a = np.array([1.0, -0.1, 0.0, 0.0]) + b = np.array([0.0, 0.0, 0.0, 0.9]) + nx = ot.backend.NumpyBackend() + _sources, _targets, _values, cost, result_code = _emd_grid_l1_1d_plan(a, b, nx) + assert result_code != 1 # not OPTIMAL: infeasible + assert cost == 0.0 + + +def test_emd_grid_l1_backends(nx): + """Non-numpy inputs (e.g. torch, jax) must be accepted and the outputs + returned in the same backend/dtype/device as the inputs, including a + round trip through GPU arrays where the backend supports them.""" + shape = (2, 2, 2) + n = int(np.prod(shape)) + rng = np.random.RandomState(0) + a = rng.rand(n) + a /= a.sum() + b = rng.rand(n) + b /= b.sum() + A, B = a.reshape(shape), b.reshape(shape) + + cost_np = emd_grid_l1(A, B) + + for tp in nx.__type_list__: + Ab, Bb = nx.from_numpy(A, B, type_as=tp) + + cost_b = emd_grid_l1(Ab, Bb) + nx.assert_same_dtype_device(tp, cost_b) + np.testing.assert_allclose(nx.to_numpy(cost_b), cost_np, rtol=1e-6, atol=1e-8) + + cost_plan_b, log_b = emd_grid_l1(Ab, Bb, return_plan=True, log=True) + nx.assert_same_dtype_device(tp, cost_plan_b) + np.testing.assert_allclose( + nx.to_numpy(cost_plan_b), cost_np, rtol=1e-6, atol=1e-8 + ) + _check_plan(log_b["G"], a, b, shape, nx.to_numpy(cost_plan_b), nx) + + +def test_emd_grid_l1_1d_backends_native_cost(nx): + """The 1D, cost-only path must also work across backends, without a plan + (and, for backends that support it, without ever leaving the device).""" + a = np.array([0.3, 0.1, 0.2, 0.4]) + b = np.array([0.1, 0.2, 0.3, 0.4]) + + cost_np = emd_grid_l1(a, b) + + for tp in nx.__type_list__: + ab, bb = nx.from_numpy(a, b, type_as=tp) + cost_b = emd_grid_l1(ab, bb) + nx.assert_same_dtype_device(tp, cost_b) + np.testing.assert_allclose(nx.to_numpy(cost_b), cost_np, rtol=1e-6, atol=1e-8) From b9785762e7714a591eea36b488a92f4c46d080c0 Mon Sep 17 00:00:00 2001 From: Tom Vercauteren Date: Fri, 18 Sep 2026 12:54:30 +0100 Subject: [PATCH 2/5] Fill in PR number in RELEASES.md Co-Authored-By: Claude Sonnet 5 --- RELEASES.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/RELEASES.md b/RELEASES.md index 10da4137d..7844fab59 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -4,7 +4,7 @@ #### New features -- Add `ot.lp.emd_grid_l1`, an exact EMD-L1 solver for histograms sharing a multi-dimensional Cartesian grid support. It solves a min-cost flow on the grid's adjacency graph rather than the full bipartite graph (Ling & Okada, 2007), which is one to two orders of magnitude faster for this case (PR #TODO) +- Add `ot.lp.emd_grid_l1`, an exact EMD-L1 solver for histograms sharing a multi-dimensional Cartesian grid support. It solves a min-cost flow on the grid's adjacency graph rather than the full bipartite graph (Ling & Okada, 2007), which is one to two orders of magnitude faster for this case (PR #863) - Use `ot.utils.check_marginal` (and shape-tuple support in `ot.utils.unif`) to fill and validate default marginals consistently across solvers (Gromov, low-rank, stochastic, barycenter, factored) (PR #856) - Add stereographic spherical sliced Wasserstein distance in `ot.sliced.stereographic_sliced_wasserstein_sphere`, with its rotationally invariant extension (PR #836) - Add Quasi-Monte Carlo sliced Wasserstein sampling (QSW/RQSW) via generalized From bab45cc34954ffdff5e1180e1d205b1a084911db Mon Sep 17 00:00:00 2001 From: Tom Vercauteren Date: Fri, 18 Sep 2026 14:46:19 +0100 Subject: [PATCH 3/5] Add test coverage for the 1D plan helper's mass-mismatch branch test_emd_grid_l1_1d_direct_plan_helper_mass_mismatch exercises the mass-mismatch early return in _emd_grid_l1_1d_plan, which codecov flagged as uncovered. Co-Authored-By: Claude Sonnet 5 --- test/test_grid.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/test/test_grid.py b/test/test_grid.py index af46e3365..ad4db9b9a 100644 --- a/test/test_grid.py +++ b/test/test_grid.py @@ -307,6 +307,15 @@ def test_emd_grid_l1_1d_direct_plan_helper_negative_values(): assert cost == 0.0 +def test_emd_grid_l1_1d_direct_plan_helper_mass_mismatch(): + a = np.array([1.0, 0.0, 0.0, 0.0]) + b = np.array([0.0, 0.0, 0.0, 0.5]) + nx = ot.backend.NumpyBackend() + _sources, _targets, _values, cost, result_code = _emd_grid_l1_1d_plan(a, b, nx) + assert result_code != 1 # not OPTIMAL: infeasible + assert cost == 0.0 + + def test_emd_grid_l1_backends(nx): """Non-numpy inputs (e.g. torch, jax) must be accepted and the outputs returned in the same backend/dtype/device as the inputs, including a From d83455ef46c1440f66576ed40e1a07b3c1e1dc7c Mon Sep 17 00:00:00 2001 From: Tom Vercauteren Date: Sat, 19 Sep 2026 16:48:31 +0100 Subject: [PATCH 4/5] Fix arange() silently ignoring type_as for dtype/device across all backends Every other array-creation helper (zeros/ones/full) respects type_as for both dtype and device, but arange() ignored it entirely for dtype (and for device on Numpy/Jax/Tensorflow/Cupy). Fixes #864. Co-Authored-By: Claude Sonnet 5 --- ot/backend.py | 27 ++++++++++++++++++++++----- test/test_backend.py | 13 +++++++++++++ 2 files changed, 35 insertions(+), 5 deletions(-) diff --git a/ot/backend.py b/ot/backend.py index fc087495c..beab6c0a4 100644 --- a/ot/backend.py +++ b/ot/backend.py @@ -1292,7 +1292,10 @@ def ones(self, shape, type_as=None): return np.ones(shape, dtype=type_as.dtype) def arange(self, stop, start=0, step=1, type_as=None): - return np.arange(start, stop, step) + if type_as is None: + return np.arange(start, stop, step) + else: + return np.arange(start, stop, step, dtype=type_as.dtype) def full(self, shape, fill_value, type_as=None): if type_as is None: @@ -1730,7 +1733,12 @@ def ones(self, shape, type_as=None): return self._change_device(jnp.ones(shape, dtype=type_as.dtype), type_as) def arange(self, stop, start=0, step=1, type_as=None): - return jnp.arange(start, stop, step) + if type_as is None: + return jnp.arange(start, stop, step) + else: + return self._change_device( + jnp.arange(start, stop, step, dtype=type_as.dtype), type_as + ) def full(self, shape, fill_value, type_as=None): if type_as is None: @@ -2237,7 +2245,9 @@ def arange(self, stop, start=0, step=1, type_as=None): if type_as is None: return torch.arange(start, stop, step) else: - return torch.arange(start, stop, step, device=type_as.device) + return torch.arange( + start, stop, step, dtype=type_as.dtype, device=type_as.device + ) def full(self, shape, fill_value, type_as=None): if isinstance(shape, int): @@ -2787,7 +2797,11 @@ def ones(self, shape, type_as=None): return cp.ones(shape, dtype=type_as.dtype) def arange(self, stop, start=0, step=1, type_as=None): - return cp.arange(start, stop, step) + if type_as is None: + return cp.arange(start, stop, step) + else: + with cp.cuda.Device(type_as.device): + return cp.arange(start, stop, step, dtype=type_as.dtype) def full(self, shape, fill_value, type_as=None): if isinstance(shape, (list, tuple)): @@ -3240,7 +3254,10 @@ def ones(self, shape, type_as=None): return tnp.ones(shape, dtype=type_as.dtype) def arange(self, stop, start=0, step=1, type_as=None): - return tnp.arange(start, stop, step) + if type_as is None: + return tnp.arange(start, stop, step) + else: + return tnp.arange(start, stop, step, dtype=type_as.dtype) def full(self, shape, fill_value, type_as=None): if type_as is None: diff --git a/test/test_backend.py b/test/test_backend.py index c88ee5052..4c3c43045 100644 --- a/test/test_backend.py +++ b/test/test_backend.py @@ -813,6 +813,19 @@ def test_func_backends(nx): ) +def test_arange_type_as(nx): + # `arange` used to silently ignore `type_as` for dtype (and device, in + # some backends), unlike `zeros`/`ones`/`full`, which do respect it. + for tp in nx.__type_list__: + a = nx.arange(5, type_as=tp) + nx.assert_same_dtype_device(tp, a) + np.testing.assert_allclose(nx.to_numpy(a), np.arange(5)) + + b = nx.arange(8, 2, 2, type_as=tp) + nx.assert_same_dtype_device(tp, b) + np.testing.assert_allclose(nx.to_numpy(b), np.arange(2, 8, 2)) + + def test_random_backends(nx): tmp_u = nx.rand() From c35d037eca470ea89150f933a3e6dbf1b6977bde Mon Sep 17 00:00:00 2001 From: Tom Vercauteren Date: Sat, 19 Sep 2026 17:21:48 +0100 Subject: [PATCH 5/5] Add exact gradient support to emd_grid_l1 Expose the dual potentials alpha/beta (the gradient of cost w.r.t. A, B) in log, via LEMON's network-simplex node potentials for d >= 2 (a free byproduct of the solve) and a closed form for 1D grids. Add a grad argument ('envelope' by default, or None) to control whether the non-free 1D gradient pass runs. Co-Authored-By: Claude Sonnet 5 --- ot/lp/EMD.h | 5 + ot/lp/EMD_wrapper.cpp | 20 ++- ot/lp/_grid.py | 368 +++++++++++++++++++++++++++++------------- ot/lp/emd_wrap.pyx | 12 +- test/test_grid.py | 232 ++++++++++++++++++++++++-- 5 files changed, 505 insertions(+), 132 deletions(-) diff --git a/ot/lp/EMD.h b/ot/lp/EMD.h index c8e2cae19..26b1edbb8 100644 --- a/ot/lp/EMD.h +++ b/ot/lp/EMD.h @@ -64,6 +64,11 @@ int EMD_wrap_grid_l1( double *plan_values_out, // Output: mass moved by each plan entry uint64_t *n_plan_entries_out, uint64_t max_plan_entries, + double *alpha, // Output: node potentials / dual variables + // (n_nodes = prod(shape)); the Beckmann + // formulation uses a single graph (not + // bipartite), so there is one potential + // array, not alpha/beta; beta = -alpha double *cost, // Output: total transportation cost uint64_t maxIter // Maximum iterations for solver ); diff --git a/ot/lp/EMD_wrapper.cpp b/ot/lp/EMD_wrapper.cpp index 3e8ff89ea..6a7c0dabf 100644 --- a/ot/lp/EMD_wrapper.cpp +++ b/ot/lp/EMD_wrapper.cpp @@ -559,6 +559,7 @@ int EMD_wrap_grid_l1( double *plan_values_out, uint64_t *n_plan_entries_out, uint64_t max_plan_entries, + double *alpha, double *cost, uint64_t maxIter ) { @@ -591,9 +592,12 @@ int EMD_wrap_grid_l1( *n_plan_entries_out = 0; if (!any_diff) { - // Histograms are identical: nothing to transport, but if a plan is - // requested, the identity coupling is still the (trivially optimal) - // transportation plan. + // Histograms are identical: the cost is 0 and constant in a + // neighbourhood of X == Y, so the zero potential is a valid + // (sub)gradient here. + std::fill(alpha, alpha + n_nodes, 0.0); + // Nothing to transport, but if a plan is requested, the identity + // coupling is still the (trivially optimal) transportation plan. if (return_plan) { for (int64_t i = 0; i < n_nodes; ++i) { if (X[i] > 1e-10) { @@ -662,6 +666,16 @@ int EMD_wrap_grid_l1( *cost = net.totalCost(); + // Node potentials (dual variables) are a byproduct of the solve, cheap + // to extract regardless of whether a plan was requested: dW/dX[i] = + // alpha[i], dW/dY[i] = -alpha[i] (beta = -alpha, since supply[i] = + // X[i] - Y[i] uses a single graph, not a bipartite source/target split). + // Negated to match LEMON's sign convention, same as the bipartite + // extract_compressed_support above (alpha = -potential). + for (int64_t i = 0; i < n_nodes; ++i) { + alpha[i] = -net.potential(Digraph::nodeFromId(static_cast(i))); + } + if (!return_plan) { // The caller only wants the cost: skip decomposing the Beckmann-style // arc flow into a transportation plan (coupling) entirely. diff --git a/ot/lp/_grid.py b/ot/lp/_grid.py index 717a96b94..e3e3281dd 100644 --- a/ot/lp/_grid.py +++ b/ot/lp/_grid.py @@ -10,16 +10,13 @@ from ..backend import get_backend from ..utils import list_to_array -from .emd_wrap import check_result, emd_1d_sorted, emd_c_grid_l1 +from .emd_wrap import check_result, emd_c_grid_l1 # Mirrors ot::ProblemType in ot/lp/EMD.h (INFEASIBLE=0, OPTIMAL=1, ...). Not # otherwise exposed to Python, since it is a Cython ``cdef enum``. _RESULT_INFEASIBLE = 0 _RESULT_OPTIMAL = 1 -_EMPTY_U64 = np.empty(0, dtype=np.uint64) -_EMPTY_F64 = np.empty(0, dtype=np.float64) - def _finalize_native_cost(cost, nx): """For the numpy backend, ``nx.sum`` etc. return a numpy scalar (e.g. @@ -33,78 +30,155 @@ def _finalize_native_cost(cost, nx): return nx.detach(cost) -def _emd_grid_l1_1d_cost(A, B, nx): - r"""Closed-form, fully backend-native cost for a 1D grid. +def _emd_grid_l1_1d(A, B, return_plan, return_alpha, nx): + r"""Fully backend-native cost, gradient, and (optionally) transportation + plan for a 1D grid. A 1D grid with a shared, sorted, unit-spaced support has a classic - :math:`\mathcal{O}(n)` closed form for the L1 (cityblock) Wasserstein - cost: the L1 norm of the difference of cumulative sums. Every step is a - generic backend reduction (``nx.cumsum``, ``nx.abs``, ``nx.sum``), so - this never leaves the device `A`/`B` are already on -- unlike every - other path in :any:`emd_grid_l1`, which needs a CPU round-trip. - - Returns ``(cost, result_code)``; `cost` is already a backend-native - scalar with `A`'s dtype and device, detached from any autodiff graph - (see :any:`_finalize_native_cost`). + closed form for both the L1 (cityblock) Wasserstein cost and its + gradient: with :math:`\text{CumA}(k) = \sum_{i \leq k} A_i` (and + likewise for `B`), + + .. math:: + \text{cost} = \sum_{k=0}^{n-2} |\text{CumA}(k) - \text{CumB}(k)|, + \quad + \alpha_i = \frac{\partial \text{cost}}{\partial A_i} = + \sum_{k=i}^{n-2} \text{sign}(\text{CumA}(k) - \text{CumB}(k)) + + (the gradient is a reverse/suffix cumulative sum of the sign of the CDF + difference, an application of the envelope theorem to this LP). Both + share the same ``nx.cumsum(A) - nx.cumsum(B)``, so are computed together + here in :math:`\mathcal{O}(n)`, with no CPU round-trip: every step is a + generic backend reduction (``nx.cumsum``, ``nx.sign``, ``nx.flip``, ...). + Unlike the cost, `alpha` is not otherwise needed, so it is only computed + when `return_alpha` is True. + + When `return_plan` is True, :any:`_emd_grid_l1_1d_monotone_plan` recovers + the actual transportation plan too, via a fully vectorized merge of the + two CDFs (also backend-native, :math:`\mathcal{O}(n \log n)` for the + sort it needs); left `None` otherwise, since it is not needed for the + cost or the gradient. + + Returns ``(cost, alpha, plan_sources, plan_targets, plan_values, + result_code)``, all backend-native, matching `A`'s dtype and device + (`alpha`, `plan_*` are `None` when not requested/not needed, see below). + `alpha` (raw, uncentred; the caller centres it) is `None` if + `return_alpha` is False. `plan_*` are `None` if `return_plan` is False. """ - if A.shape[0] == 0 or nx.any(A < 0) or nx.any(B < 0): - return _finalize_native_cost(0.0 * nx.sum(A), nx), _RESULT_INFEASIBLE + n = A.shape[0] + + if n == 0 or nx.any(A < 0) or nx.any(B < 0): + zero_cost = _finalize_native_cost(0.0 * nx.sum(A), nx) + zero_alpha = nx.zeros(A.shape, type_as=A) if return_alpha else None + return zero_cost, zero_alpha, None, None, None, _RESULT_INFEASIBLE total_a = nx.sum(A) total_b = nx.sum(B) if abs(float(nx.to_numpy(total_a)) - float(nx.to_numpy(total_b))) > 1e-8 * max( 1.0, float(nx.to_numpy(total_a)) ): - return _finalize_native_cost(0.0 * total_a, nx), _RESULT_INFEASIBLE - - cum_diff = nx.cumsum(A, 0) - nx.cumsum(B, 0) - cost = nx.sum(nx.abs(cum_diff[:-1])) - return _finalize_native_cost(cost, nx), _RESULT_OPTIMAL - - -def _emd_grid_l1_1d_plan(A, B, nx): - r"""Cost and sparse transportation plan for a 1D grid. - - A 1D grid is just a sorted, shared support, for which POT already has an - exact :math:`\mathcal{O}(n)` solver (:any:`ot.lp.emd_1d_sorted`, the same - routine backing :any:`ot.lp.emd_1d`). Calling it directly here, with the - grid's integer positions passed in as already sorted, skips both the - network simplex setup of the general grid solver and the argsort/pairwise - distance overhead that the generic, arbitrary-support :any:`ot.lp.emd_1d` - would otherwise pay. - - Unlike :any:`_emd_grid_l1_1d_cost`, recovering the actual coupling needs - the compiled O(n) merge behind :any:`ot.lp.emd_1d_sorted`, which requires - plain CPU numpy arrays -- so, backend-compatible via `nx`, this converts - `A`/`B` right here, immediately before that one call. + zero_cost = _finalize_native_cost(0.0 * total_a, nx) + zero_alpha = nx.zeros(A.shape, type_as=A) if return_alpha else None + return zero_cost, zero_alpha, None, None, None, _RESULT_INFEASIBLE + + cum_a = nx.cumsum(A, 0) + cum_b = nx.cumsum(B, 0) + cum_diff = cum_a - cum_b + cost = _finalize_native_cost(nx.sum(nx.abs(cum_diff[:-1])), nx) + + if return_alpha: + sign = nx.sign(cum_diff[:-1]) + suffix = nx.flip(nx.cumsum(nx.flip(sign, 0), 0), 0) + alpha = nx.concatenate([suffix, nx.zeros((1,), type_as=A)], axis=0) + else: + alpha = None + + if not return_plan: + return cost, alpha, None, None, None, _RESULT_OPTIMAL + + # cum_a, cum_b are reused as is (the plan needs the raw, unpinned CDFs + # too), avoiding a second nx.cumsum(A) / nx.cumsum(B). + plan_sources, plan_targets, plan_values = _emd_grid_l1_1d_monotone_plan( + cum_a, cum_b, nx + ) + return cost, alpha, plan_sources, plan_targets, plan_values, _RESULT_OPTIMAL + + +def _emd_grid_l1_1d_monotone_plan(cum_a, cum_b, nx): + r"""Exact monotone 1D transportation plan on a shared grid support, in + sparse COO form. + + Works by merging the two CDFs: sorting their :math:`2n` combined + breakpoints together, the cumulative count of A-breakpoints and + B-breakpoints seen so far at each step of the merge gives that step's + (source bin, target bin) pair, and the gap to the previous breakpoint + gives the mass moved. A source bin can touch more than two target bins + in general (and vice versa), which is why this has to be indexed by + merge step rather than by bin -- the same reason the classic sequential + merge algorithm behind :any:`ot.lp.emd_1d_sorted` needs to walk source + and target pointers independently. + + Takes the raw (unpinned) CDFs `cum_a`, `cum_b` (``nx.cumsum(A, 0)``, + ``nx.cumsum(B, 0)``) rather than `A`, `B` themselves, since the caller + (:any:`_emd_grid_l1_1d`) already needs them for the cost and can pass + them along here, avoiding computing them twice. Assumes `A`, `B` are + feasible (nonnegative, matching total mass); the caller is responsible + for checking this first. + + Returns backend-native ``(plan_sources, plan_targets, plan_values)``, + matching `cum_a`'s dtype and device, with exactly :math:`2n-1` entries. + Some entries may carry zero flow (e.g. at ties between the two CDFs, + such as when `A` and `B` are identical): these are harmless, explicit + zeros in the sparse coupling built from them, and filtering them out + would need a backend-native "keep the nonzero entries" op that does not + exist as a generic, shape-static reduction, so it is not worth it for + what is already a tiny, fixed-size result. """ - a = np.ascontiguousarray(nx.to_numpy(A), dtype=np.float64) - b = np.ascontiguousarray(nx.to_numpy(B), dtype=np.float64) - n = a.shape[0] - - if n == 0 or np.any(a < 0) or np.any(b < 0): - return _EMPTY_U64, _EMPTY_U64, _EMPTY_F64, 0.0, _RESULT_INFEASIBLE - - total_a = a.sum() - total_b = b.sum() - if abs(total_a - total_b) > 1e-8 * max(1.0, total_a): - return _EMPTY_U64, _EMPTY_U64, _EMPTY_F64, 0.0, _RESULT_INFEASIBLE - - x = np.arange(n, dtype=np.float64) - plan_values, indices, cost = emd_1d_sorted(a, b, x, x, metric="cityblock") - - # The merge algorithm behind emd_1d_sorted can emit exact-zero-mass - # entries at ties (e.g. two identical histograms); drop them so the plan - # only lists actual mass movements, matching the general grid solver. - nonzero = plan_values > 0 - plan_sources = np.ascontiguousarray(indices[nonzero, 0], dtype=np.uint64) - plan_targets = np.ascontiguousarray(indices[nonzero, 1], dtype=np.uint64) - plan_values = np.ascontiguousarray(plan_values[nonzero], dtype=np.float64) - return plan_sources, plan_targets, plan_values, cost, _RESULT_OPTIMAL + n = cum_a.shape[0] + + # Pin the last entry of each CDF to exactly 1: avoids floating-point + # drift making the tie at the very end (both CDFs must reach the same + # total mass) not an exact tie, which would otherwise show up as a + # spurious tiny "flow" entry below. + one = nx.reshape(0.0 * cum_a[-1] + 1.0, (1,)) + cum_a = nx.concatenate([cum_a[:-1], one], axis=0) + cum_b = nx.concatenate([cum_b[:-1], one], axis=0) + all_cum = nx.concatenate([cum_a, cum_b], axis=0) # length 2n + + # One-hot markers of where each of the 2n merged breakpoints came from, + # carried through the sort below to recover, at each step, how many of + # each type preceded it. + ones_n = nx.ones((n,), type_as=cum_a) + zeros_n = nx.zeros((n,), type_as=cum_a) + origin_source = nx.concatenate([ones_n, zeros_n], axis=0) + origin_target = nx.concatenate([zeros_n, ones_n], axis=0) + + perm = nx.argsort(all_cum, axis=-1) + sort_val = nx.take_along_axis(all_cum, perm, axis=-1) + sort_source = nx.take_along_axis(origin_source, perm, axis=-1) + sort_target = nx.take_along_axis(origin_target, perm, axis=-1) + + origin = nx.stack([sort_source, sort_target], axis=0) # (2, 2n) + excl_cumsum = nx.cumsum(origin, axis=-1) - origin + idx = nx.clip(excl_cumsum, None, n - 1) # (2, 2n): [0] source, [1] target + + left = nx.zero_pad(sort_val[:-1], [(1, 0)], value=0.0) + flow = nx.clip(sort_val - left, 0.0, None) + + # The very last merge step is the simultaneous end of both CDFs (both + # pinned to 1 above), carrying no flow; drop it, leaving exactly 2n-1 + # entries, as expected for a monotone coupling of two n-atom measures. + return idx[0, :-1], idx[1, :-1], flow[:-1] def emd_grid_l1( - A, B, numItermax=100000, return_plan=False, log=False, check_marginals=True + A, + B, + numItermax=100000, + return_plan=False, + log=False, + check_marginals=True, + grad="envelope", ): r"""Solves the Earth Mover's Distance with the cityblock ground metric between two histograms sharing the same :math:`d`-dimensional Cartesian @@ -143,22 +217,40 @@ def emd_grid_l1( its own. Set `return_plan` to request it; leave it False (the default) when only the transport cost is needed. - .. note:: For a 1D grid (``A.ndim == 1``), this instead delegates to - POT's :math:`\mathcal{O}(n)` sorted-support solver - (:any:`ot.lp.emd_1d_sorted`), which is exact here since the grid's - integer positions are already a shared, sorted support. This avoids - the network simplex setup entirely for that case. Better still, when - `return_plan` is False, the cost itself has a closed form that runs - as generic backend reductions with no CPU round-trip at all, so a 1D - grid on a GPU array is solved entirely on-device. + .. note:: For a 1D grid (``A.ndim == B.ndim == 1``), this instead uses a + closed form (see :any:`_emd_grid_l1_1d`), exact here since the + grid's integer positions are already a shared, sorted support. This + avoids the network simplex setup entirely for that case, and runs + as generic backend reductions with no CPU round-trip at all -- for + the cost and gradient, but also for the plan, when requested -- so + a 1D grid on a GPU array is solved entirely on-device. .. note:: This function is backend-compatible and will work on arrays - from all compatible backends. Beyond the 1D, cost-only case above, - the algorithm uses a C++/Cython CPU solver, so GPU arrays are copied + from all compatible backends. Beyond the 1D case above, the + algorithm uses a C++/Cython CPU solver, so GPU arrays are copied to CPU before solving (and the transportation plan's bin indices, if - requested, are returned as CPU arrays). Gradients are not currently - supported: `cost` is always detached from any computation graph the - inputs were part of. + requested, are returned as CPU arrays). `cost` is always detached + from any computation graph the inputs were part of: this does not + support automatic differentiation. The exact gradient of `cost` + with respect to `A` and `B` is instead exposed explicitly as + `alpha`/`beta` in `log` (see below). + + .. note:: `log["alpha"]` and `log["beta"]`, when present, are the + (centred) dual potentials, which give the gradient of `cost` with + respect to `A` and `B` directly: :math:`\partial \text{cost}/\partial + A_i = \text{alpha}_i` and :math:`\partial \text{cost}/\partial B_i = + \text{beta}_i = -\text{alpha}_i` (a single graph is used, not a + bipartite source/target split, so there is one potential array, not + two). For :math:`d \geq 2` these are LEMON's network-simplex node + potentials, an unavoidable byproduct of the solve itself (a free + application of the envelope theorem to this LP), so they are always + computed and returned in `log` regardless of `grad`. For a 1D grid + they are instead a closed form (see :any:`_emd_grid_l1_1d`) that + needs its own, separate :math:`\mathcal{O}(n)` pass -- not otherwise + free -- so `grad` controls whether it runs there. Either way they + are defined up to an additive constant; centring (subtracting their + mean) picks the canonical representative, the Riemannian gradient of + `cost` on the probability simplex. Parameters ---------- @@ -172,14 +264,25 @@ def emd_grid_l1( return_plan : bool, optional (default=False) If True, additionally recovers an explicit transportation plan (a sparse coupling), returned in `log`. Computing it has a cost of its - own (a CPU round-trip, and either a network simplex solve or an O(n) - merge), so it is skipped by default when only the transport cost - `cost` is needed. + own (for :math:`d \geq 2`, a network simplex solve, plus a CPU + round-trip since that solver is CPU-only; for a 1D grid, an + :math:`\mathcal{O}(n \log n)` merge, entirely on-device), so it is + skipped by default when only the transport cost `cost` is needed. log : bool, optional (default=False) If True, also returns a dictionary with the solver status and, if `return_plan` is True, the sparse transportation plan. check_marginals : bool, optional (default=True) If True, checks that `A` and `B` have the same total mass. + grad : {'envelope', None}, optional (default='envelope') + Controls whether the dual potentials `alpha`/`beta` (the gradient of + `cost`) are computed for a 1D grid, where doing so has a cost of its + own (see the note above). For :math:`d \geq 2` they come for free as + a byproduct of the network-simplex solve, so this has no effect + there: they are always computed and returned in `log`. If + `'envelope'` (the default), also compute them for a 1D grid, via the + envelope theorem applied to that grid's closed form. If None, skip + that computation for a 1D grid, and `log` will not contain `alpha` + or `beta` in that case. Has no effect unless `log` is True. Returns ------- @@ -188,12 +291,14 @@ def emd_grid_l1( log : dict, optional If input `log` is True, a dictionary containing the solver status (`warning`, `result_code`) and, if `return_plan` is True, the sparse - transportation plan `G` (built via the backend's `coo_matrix`, same - as :any:`ot.emd2_lazy`'s `return_matrix`; a real sparse matrix for - NumPy/PyTorch/TensorFlow/CuPy, silently densified for JAX, which has - no sparse array type) of shape :math:`(n, n)` with - :math:`n=\prod(\text{A.shape})`, indexing into `A.reshape(-1)` and - `B.reshape(-1)`. + transportation plan `G` (built via the backend's + `coo_matrix`, same as :any:`ot.emd2_lazy`'s `return_matrix`; a real + sparse matrix for NumPy/PyTorch/TensorFlow/CuPy, silently densified + for JAX, which has no sparse array type) of shape :math:`(n, n)` + with :math:`n=\prod(\text{A.shape})`, indexing into `A.reshape(-1)` + and `B.reshape(-1)`. Unless `A.ndim == 1` and `grad` is None, it also + contains the (centred) dual potentials `alpha`, `beta` (the gradient + of `cost` with respect to `A`, `B`; see the note above). Examples -------- @@ -213,6 +318,9 @@ def emd_grid_l1( -------- ot.emd : Exact OT solver with a general, precomputed cost matrix """ + if grad not in (None, "envelope"): + raise ValueError(f"grad must be None or 'envelope', got {grad!r}") + A, B = list_to_array(A, B) nx = get_backend(A, B) @@ -231,40 +339,78 @@ def emd_grid_l1( ) if A.ndim == 1: - # A 1D grid is backend-native either way (works on any backend's - # arrays via `nx`, GPU included). Only the cost-only case below - # avoids a CPU round-trip entirely though: recovering the plan needs - # the compiled O(n) merge in _emd_grid_l1_1d_plan, which is CPU-only. - if not return_plan: - cost, result_code = _emd_grid_l1_1d_cost(A, B, nx) - if log: - return cost, { - "warning": check_result(result_code), - "result_code": result_code, - } - check_result(result_code) - return cost - - plan_sources, plan_targets, plan_values, cost, result_code = ( - _emd_grid_l1_1d_plan(A, B, nx) - ) - else: - shape = np.array(A.shape, dtype=np.int64) - # The C++ solver only understands flattened (CPU) numpy arrays: - # `to_numpy` also does the GPU -> CPU copy for backends such as - # torch or jax. - a_np = np.ascontiguousarray(nx.to_numpy(A), dtype=np.float64).reshape(-1) - b_np = np.ascontiguousarray(nx.to_numpy(B), dtype=np.float64).reshape(-1) - plan_sources, plan_targets, plan_values, cost, result_code = emd_c_grid_l1( - a_np, b_np, shape, numItermax, return_plan + # A 1D grid is backend-native throughout -- cost, gradient, and + # (when requested) the transportation plan -- so it gets its own, + # fully self-contained branch. See _emd_grid_l1_1d. Unlike the + # general (d >= 2) path below, the gradient is not a free byproduct + # here, so it is only computed when actually requested. + return_alpha = log and grad == "envelope" + cost, alpha, plan_sources, plan_targets, plan_values, result_code = ( + _emd_grid_l1_1d(A, B, return_plan, return_alpha, nx) ) + if log: + log_dict = { + "warning": check_result(result_code), + "result_code": result_code, + } + if alpha is not None: + # `alpha` is only defined up to an additive constant; + # centring it picks the canonical representative, the + # Riemannian gradient of the cost on the probability + # simplex. + alpha = alpha - nx.mean(alpha) + log_dict["alpha"] = alpha + log_dict["beta"] = -alpha + if return_plan and result_code == _RESULT_OPTIMAL: + # plan_sources/targets/values are already backend-native + # (matching A), so this is a plain packaging step, with no + # conversion needed. + n = A.shape[0] + log_dict["G"] = nx.coo_matrix( + plan_values, + plan_sources, + plan_targets, + shape=(n, n), + type_as=A, + ) + return cost, log_dict + + check_result(result_code) + return cost + + # ndim >= 2: the general grid solver, backed by network simplex in C++. + shape = np.array(A.shape, dtype=np.int64) + # The C++ solver only understands flattened (CPU) numpy arrays: + # `to_numpy` also does the GPU -> CPU copy for backends such as torch or + # jax. + a_np = np.ascontiguousarray(nx.to_numpy(A), dtype=np.float64).reshape(-1) + b_np = np.ascontiguousarray(nx.to_numpy(B), dtype=np.float64).reshape(-1) + ( + plan_sources, + plan_targets, + plan_values, + alpha, + cost, + result_code, + ) = emd_c_grid_l1(a_np, b_np, shape, numItermax, return_plan) + + # `alpha` (the node potentials) is only defined up to an additive + # constant; centring it picks the canonical representative, the + # Riemannian gradient of the cost on the probability simplex. + # beta = -alpha since this is a single graph, not a bipartite + # source/target split (supply[i] = A[i] - B[i] for every node). + alpha = alpha - alpha.mean() + beta = -alpha + cost = nx.from_numpy(cost, type_as=A) if log: log_dict = { "warning": check_result(result_code), "result_code": result_code, + "alpha": nx.from_numpy(alpha, type_as=A), + "beta": nx.from_numpy(beta, type_as=A), } if return_plan: # A.size is a numpy property but a torch method: go through diff --git a/ot/lp/emd_wrap.pyx b/ot/lp/emd_wrap.pyx index 5a9ddfc7f..03c0b578f 100644 --- a/ot/lp/emd_wrap.pyx +++ b/ot/lp/emd_wrap.pyx @@ -22,7 +22,7 @@ import warnings cdef extern from "EMD.h": int EMD_wrap(int n1,int n2, double *X, double *Y,double *D, double *G, double* alpha, double* beta, double *cost, uint64_t maxIter, double* alpha_init, double* beta_init) nogil int EMD_wrap_sparse(int n1, int n2, double *X, double *Y, uint64_t n_edges, uint64_t *edge_sources, uint64_t *edge_targets, double *edge_costs, uint64_t *flow_sources_out, uint64_t *flow_targets_out, double *flow_values_out, uint64_t *n_flows_out, uint64_t max_flows_out, double *alpha, double *beta, double *cost, uint64_t maxIter, double* alpha_init, double* beta_init) nogil - int EMD_wrap_grid_l1(int ndim, int64_t *shape, double *X, double *Y, bint return_plan, uint64_t *plan_sources_out, uint64_t *plan_targets_out, double *plan_values_out, uint64_t *n_plan_entries_out, uint64_t max_plan_entries, double *cost, uint64_t maxIter) nogil + int EMD_wrap_grid_l1(int ndim, int64_t *shape, double *X, double *Y, bint return_plan, uint64_t *plan_sources_out, uint64_t *plan_targets_out, double *plan_values_out, uint64_t *n_plan_entries_out, uint64_t max_plan_entries, double *alpha, double *cost, uint64_t maxIter) nogil int EMD_wrap_lazy(int n1, int n2, double *X, double *Y, double *coords_a, double *coords_b, int dim, int metric, uint64_t *flow_sources_out, uint64_t *flow_targets_out, double *flow_values_out, uint64_t *n_flows_out, uint64_t max_flows_out, double* alpha, double* beta, double *cost, uint64_t maxIter, double* alpha_init, double* beta_init) nogil cdef enum ProblemType: INFEASIBLE, OPTIMAL, UNBOUNDED, MAX_ITER_REACHED @@ -421,6 +421,11 @@ def emd_c_grid_l1(np.ndarray[double, ndim=1, mode="c"] a, plan_values : (n_plan_entries,) array, float64 Mass moved by each transportation plan entry (empty if `return_plan` is False) + alpha : (n,) array, float64 + Raw (uncentred) node potentials, i.e. d(cost)/d(a) up to the additive + constant LEMON's network simplex happens to settle on; d(cost)/d(b) + is -alpha, since this is a single graph, not a bipartite source/ + target split. Centre before use: `alpha -= alpha.mean()`. cost : float Total transportation cost result_code : int @@ -439,6 +444,7 @@ def emd_c_grid_l1(np.ndarray[double, ndim=1, mode="c"] a, cdef np.ndarray[uint64_t, ndim=1, mode="c"] plan_sources = np.zeros(max_plan_entries, dtype=np.uint64) cdef np.ndarray[uint64_t, ndim=1, mode="c"] plan_targets = np.zeros(max_plan_entries, dtype=np.uint64) cdef np.ndarray[double, ndim=1, mode="c"] plan_values = np.zeros(max_plan_entries, dtype=np.float64) + cdef np.ndarray[double, ndim=1, mode="c"] alpha = np.zeros(a.shape[0], dtype=np.float64) with nogil: result_code = EMD_wrap_grid_l1( @@ -446,11 +452,11 @@ def emd_c_grid_l1(np.ndarray[double, ndim=1, mode="c"] a, a.data, b.data, return_plan, plan_sources.data, plan_targets.data, plan_values.data, - &n_plan_entries_out, max_plan_entries, &cost, max_iter + &n_plan_entries_out, max_plan_entries, alpha.data, &cost, max_iter ) plan_sources = plan_sources[:n_plan_entries_out] plan_targets = plan_targets[:n_plan_entries_out] plan_values = plan_values[:n_plan_entries_out] - return plan_sources, plan_targets, plan_values, cost, result_code + return plan_sources, plan_targets, plan_values, alpha, cost, result_code diff --git a/test/test_grid.py b/test/test_grid.py index ad4db9b9a..175cf573c 100644 --- a/test/test_grid.py +++ b/test/test_grid.py @@ -12,7 +12,7 @@ import ot from ot.lp import emd_grid_l1 -from ot.lp._grid import _emd_grid_l1_1d_cost, _emd_grid_l1_1d_plan +from ot.lp._grid import _emd_grid_l1_1d, _emd_grid_l1_1d_monotone_plan def _grid_coords(shape): @@ -185,15 +185,15 @@ def test_emd_grid_l1_random_grids_batch(ndim): def test_emd_grid_l1_1d_uses_native_cost_path(): """A 1D grid with return_plan=False must skip both the general C++ - solver and the O(n) merge that recovers the plan.""" + solver and the merge that recovers the plan.""" a = np.array([1.0, 0.0, 0.0, 0.0]) b = np.array([0.0, 0.0, 0.0, 1.0]) with ( mock.patch("ot.lp._grid.emd_c_grid_l1") as mocked_cpp, mock.patch( - "ot.lp._grid._emd_grid_l1_1d_plan", - wraps=_emd_grid_l1_1d_plan, + "ot.lp._grid._emd_grid_l1_1d_monotone_plan", + wraps=_emd_grid_l1_1d_monotone_plan, ) as mocked_plan, ): cost = emd_grid_l1(a, b) @@ -202,12 +202,12 @@ def test_emd_grid_l1_1d_uses_native_cost_path(): np.testing.assert_allclose(cost, 3.0) # Requesting the plan on a 1D grid must still avoid the C++ solver, but - # does need the O(n) merge. + # does need the merge. with ( mock.patch("ot.lp._grid.emd_c_grid_l1") as mocked_cpp, mock.patch( - "ot.lp._grid._emd_grid_l1_1d_plan", - wraps=_emd_grid_l1_1d_plan, + "ot.lp._grid._emd_grid_l1_1d_monotone_plan", + wraps=_emd_grid_l1_1d_monotone_plan, ) as mocked_plan, ): emd_grid_l1(a, b, return_plan=True) @@ -277,43 +277,65 @@ def test_emd_grid_l1_1d_mass_mismatch(): assert log["result_code"] != 1 # not OPTIMAL: infeasible -def test_emd_grid_l1_1d_direct_cost_helper(): +def test_emd_grid_l1_1d_direct_helper_cost_only(): nx = ot.backend.NumpyBackend() a = np.array([1.0, 0.0, 0.0, 0.0]) b = np.array([0.0, 0.0, 0.0, 1.0]) - cost, result_code = _emd_grid_l1_1d_cost(a, b, nx) + cost, alpha, sources, targets, values, result_code = _emd_grid_l1_1d( + a, b, False, False, nx + ) assert result_code == 1 # OPTIMAL np.testing.assert_allclose(cost, 3.0) + assert alpha is None + assert sources is None and targets is None and values is None + + # return_alpha=True must compute it even without a plan. + _cost, alpha, _sources, _targets, _values, result_code = _emd_grid_l1_1d( + a, b, False, True, nx + ) + assert alpha is not None # Negative values and mass mismatches must be reported as infeasible, # like the general (C++) path. a_neg = np.array([1.0, -0.1, 0.0, 0.0]) - _cost, result_code = _emd_grid_l1_1d_cost(a_neg, b, nx) + _cost, _alpha, _sources, _targets, _values, result_code = _emd_grid_l1_1d( + a_neg, b, False, False, nx + ) assert result_code != 1 a_mismatch = np.array([1.0, 0.0, 0.0, 0.0]) b_mismatch = np.array([0.0, 0.0, 0.0, 0.5]) - _cost, result_code = _emd_grid_l1_1d_cost(a_mismatch, b_mismatch, nx) + _cost, _alpha, _sources, _targets, _values, result_code = _emd_grid_l1_1d( + a_mismatch, b_mismatch, False, False, nx + ) assert result_code != 1 -def test_emd_grid_l1_1d_direct_plan_helper_negative_values(): +def test_emd_grid_l1_1d_direct_helper_plan_negative_values(): a = np.array([1.0, -0.1, 0.0, 0.0]) b = np.array([0.0, 0.0, 0.0, 0.9]) nx = ot.backend.NumpyBackend() - _sources, _targets, _values, cost, result_code = _emd_grid_l1_1d_plan(a, b, nx) + cost, alpha, sources, targets, values, result_code = _emd_grid_l1_1d( + a, b, True, True, nx + ) assert result_code != 1 # not OPTIMAL: infeasible assert cost == 0.0 + assert alpha is not None + assert sources is None and targets is None and values is None -def test_emd_grid_l1_1d_direct_plan_helper_mass_mismatch(): +def test_emd_grid_l1_1d_direct_helper_plan_mass_mismatch(): a = np.array([1.0, 0.0, 0.0, 0.0]) b = np.array([0.0, 0.0, 0.0, 0.5]) nx = ot.backend.NumpyBackend() - _sources, _targets, _values, cost, result_code = _emd_grid_l1_1d_plan(a, b, nx) + cost, alpha, sources, targets, values, result_code = _emd_grid_l1_1d( + a, b, True, True, nx + ) assert result_code != 1 # not OPTIMAL: infeasible assert cost == 0.0 + assert alpha is not None + assert sources is None and targets is None and values is None def test_emd_grid_l1_backends(nx): @@ -359,3 +381,183 @@ def test_emd_grid_l1_1d_backends_native_cost(nx): cost_b = emd_grid_l1(ab, bb) nx.assert_same_dtype_device(tp, cost_b) np.testing.assert_allclose(nx.to_numpy(cost_b), cost_np, rtol=1e-6, atol=1e-8) + + +@pytest.mark.parametrize("shape", [(8,), (3, 4), (2, 3, 2)]) +def test_emd_grid_l1_gradient_strong_duality(shape): + """The (centred) node potentials alpha, beta=-alpha are the gradient of + cost w.r.t. A, B, and must satisfy strong duality: for the optimal + coupling, dot(alpha, A) + dot(beta, B) == cost.""" + n = int(np.prod(shape)) + rng = np.random.RandomState(42) + for _trial in range(5): + a = rng.rand(n) + a /= a.sum() + b = rng.rand(n) + b /= b.sum() + A, B = a.reshape(shape), b.reshape(shape) + + cost, log = emd_grid_l1(A, B, log=True) + alpha = log["alpha"] + beta = log["beta"] + + np.testing.assert_allclose(beta, -alpha, atol=1e-10) + np.testing.assert_allclose( + np.dot(alpha, a) + np.dot(beta, b), cost, rtol=1e-6, atol=1e-8 + ) + # Equivalent, since beta = -alpha. + np.testing.assert_allclose(np.dot(alpha, a - b), cost, rtol=1e-6, atol=1e-8) + + +@pytest.mark.parametrize("shape", [(8,), (3, 4), (2, 3, 2)]) +def test_emd_grid_l1_gradient_finite_differences(shape): + """alpha[i] - alpha[j] must match the centered finite difference of cost + with respect to moving mass eps from bin j to bin i.""" + n = int(np.prod(shape)) + rng = np.random.RandomState(0) + a = rng.rand(n) + a /= a.sum() + b = rng.rand(n) + b /= b.sum() + A, B = a.reshape(shape), b.reshape(shape) + + _cost, log = emd_grid_l1(A, B, log=True) + alpha = log["alpha"] + + eps = 1e-5 + pairs = { + (int(i), int(j)) + for i, j in zip(rng.randint(0, n, size=8), rng.randint(0, n, size=8)) + if i != j + } + for i, j in pairs: + a_plus = a.copy() + a_plus[i] += eps + a_plus[j] -= eps + a_minus = a.copy() + a_minus[i] -= eps + a_minus[j] += eps + + cost_plus = emd_grid_l1(a_plus.reshape(shape), B) + cost_minus = emd_grid_l1(a_minus.reshape(shape), B) + fd = (cost_plus - cost_minus) / (2 * eps) + + np.testing.assert_allclose(fd, alpha[i] - alpha[j], atol=1e-4) + + +def test_emd_grid_l1_1d_gradient_matches_return_plan(): + """The gradient must not depend on whether return_plan is also + requested: both the cost-only closed form and the with-plan branch + dispatch to the same closed-form alpha.""" + n = 7 + rng = np.random.RandomState(9) + a = rng.rand(n) + a /= a.sum() + b = rng.rand(n) + b /= b.sum() + + cost, log = emd_grid_l1(a, b, log=True) + cost_p, log_p = emd_grid_l1(a, b, return_plan=True, log=True) + + np.testing.assert_allclose(cost, cost_p, rtol=1e-6, atol=1e-8) + np.testing.assert_allclose(log["alpha"], log_p["alpha"], atol=1e-10) + np.testing.assert_allclose(log["beta"], log_p["beta"], atol=1e-10) + np.testing.assert_allclose( + np.dot(log_p["alpha"], a - b), cost_p, rtol=1e-6, atol=1e-8 + ) + + +def test_emd_grid_l1_1d_gradient_infeasible_is_zero(): + """Matches the general (C++) path: no meaningful gradient when + infeasible, so alpha/beta are exactly zero rather than nonsense.""" + a = np.array([1.0, 0.0, 0.0, 0.0]) + b = np.array([0.0, 0.0, 0.0, 0.5]) + + _cost, log = emd_grid_l1(a, b, check_marginals=False, log=True) + assert log["result_code"] != 1 # not OPTIMAL: infeasible + np.testing.assert_allclose(log["alpha"], 0.0) + np.testing.assert_allclose(log["beta"], 0.0) + + +def test_emd_grid_l1_1d_gradient_backends(nx): + """The 1D gradient must also be backend-native: dtype/device matching + the inputs, computed without a CPU round-trip.""" + n = 6 + rng = np.random.RandomState(4) + a = rng.rand(n) + a /= a.sum() + b = rng.rand(n) + b /= b.sum() + + cost_np, log_np = emd_grid_l1(a, b, log=True) + + for tp in nx.__type_list__: + ab, bb = nx.from_numpy(a, b, type_as=tp) + cost_b, log_b = emd_grid_l1(ab, bb, log=True) + nx.assert_same_dtype_device(tp, log_b["alpha"]) + nx.assert_same_dtype_device(tp, log_b["beta"]) + np.testing.assert_allclose( + nx.to_numpy(log_b["alpha"]), log_np["alpha"], rtol=1e-6, atol=1e-8 + ) + np.testing.assert_allclose(nx.to_numpy(cost_b), cost_np, rtol=1e-6, atol=1e-8) + + +def test_emd_grid_l1_grad_invalid_raises(): + a = np.array([1.0, 0.0, 0.0, 0.0]) + b = np.array([0.0, 0.0, 0.0, 1.0]) + with pytest.raises(ValueError): + emd_grid_l1(a, b, grad="bogus") + + +def test_emd_grid_l1_1d_grad_none_omits_alpha_and_skips_computation(): + """grad=None on a 1D grid must skip the (non-free) O(n) gradient pass + entirely, and log must not contain alpha/beta.""" + a = np.array([1.0, 0.0, 0.0, 0.0]) + b = np.array([0.0, 0.0, 0.0, 1.0]) + + with mock.patch("ot.lp._grid._emd_grid_l1_1d", wraps=_emd_grid_l1_1d) as mocked_1d: + cost, log = emd_grid_l1(a, b, grad=None, log=True) + mocked_1d.assert_called_once_with(a, b, False, False, mock.ANY) + np.testing.assert_allclose(cost, 3.0) + assert "alpha" not in log + assert "beta" not in log + + # The default ('envelope') must still compute and return it. + with mock.patch("ot.lp._grid._emd_grid_l1_1d", wraps=_emd_grid_l1_1d) as mocked_1d: + _cost, log = emd_grid_l1(a, b, log=True) + mocked_1d.assert_called_once_with(a, b, False, True, mock.ANY) + assert "alpha" in log + assert "beta" in log + + +def test_emd_grid_l1_1d_grad_none_without_log_unaffected(): + """grad only matters when log is requested; without log, behavior and + return value must be unchanged.""" + a = np.array([1.0, 0.0, 0.0, 0.0]) + b = np.array([0.0, 0.0, 0.0, 1.0]) + + cost_default = emd_grid_l1(a, b) + cost_grad_none = emd_grid_l1(a, b, grad=None) + np.testing.assert_allclose(cost_default, cost_grad_none) + + +@pytest.mark.parametrize("shape", [(3, 4), (2, 2, 2)]) +def test_emd_grid_l1_grad_none_multid_still_has_alpha(shape): + """For d >= 2, alpha/beta are a free byproduct of the network-simplex + solve, so grad=None must not remove them from log.""" + n = int(np.prod(shape)) + rng = np.random.RandomState(3) + a = rng.rand(n) + a /= a.sum() + b = rng.rand(n) + b /= b.sum() + A, B = a.reshape(shape), b.reshape(shape) + + cost, log = emd_grid_l1(A, B, grad=None, log=True) + cost_default, log_default = emd_grid_l1(A, B, log=True) + + assert "alpha" in log + assert "beta" in log + np.testing.assert_allclose(cost, cost_default) + np.testing.assert_allclose(log["alpha"], log_default["alpha"]) + np.testing.assert_allclose(log["beta"], log_default["beta"])