From 2ca49e95d4d686a1b8fd11b39d34ff7080f704b8 Mon Sep 17 00:00:00 2001 From: Tom Vercauteren Date: Sun, 20 Sep 2026 14:01:34 +0100 Subject: [PATCH 1/2] Fix Backend.logsumexp API conformance across Cupy and Tensorflow Backend.logsumexp is declared as logsumexp(self, a, axis=None, keepdims=False) and documented to follow the scipy.special.logsumexp API. NumPy, Jax and Torch honour that signature, but CupyBackend and TensorflowBackend were declared without keepdims and raised a TypeError if it was passed. CupyBackend.logsumexp also reimplemented a 2021 SciPy snippet by hand; it now delegates to cupyx.scipy.special.logsumexp, which matches the hand-rolled all-(-inf) handling exactly (verified against the current cupy source) and adds keepdims support in one call. TensorflowBackend.logsumexp now forwards keepdims to tf.math.reduce_logsumexp, which already supports it. No current caller passes keepdims, so this changes no behaviour for existing callers. Extends test_func_backends in test/test_backend.py with axis=0, axis=1, keepdims=True and an all-(-inf)-row case, verified on numpy/jax/torch/tf; the Cupy path is unverified by CI (no Cupy/GPU available here). Fixes #867 Co-Authored-By: Claude Sonnet 5 --- RELEASES.md | 1 + ot/backend.py | 23 +++++------------------ test/test_backend.py | 20 ++++++++++++++++++++ 3 files changed, 26 insertions(+), 18 deletions(-) diff --git a/RELEASES.md b/RELEASES.md index 063701229..0169d261f 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -15,6 +15,7 @@ #### Closed issues +- Fix `Backend.logsumexp` API conformance: `CupyBackend.logsumexp` and `TensorflowBackend.logsumexp` did not accept `keepdims`, unlike the NumPy/Jax/Torch backends and the base class signature; `CupyBackend.logsumexp` now delegates to `cupyx.scipy.special.logsumexp` instead of a hand-rolled port (Issue #867) - Remove a leftover debug `print` from `ot.utils.projection_sparse_simplex` with `axis=1`, and make the `ot.datasets.make_gauss_hd` docstring a raw string so importing `ot` no longer emits a `SyntaxWarning` (PR #860) - Fix `ot.dist` ignoring the weights `w` for `metric="cityblock"`, which returned the unweighted distance although the weights are documented for this metric (PR #859) - Fix swapped arguments to `div_to_product` in `ot.gromov.fused_unbalanced_across_spaces_cost`: with `reg_type="independent"` (UCOOT) the entropic terms used the plan marginals as the reference measures and vice versa (PR #855, Issue #854) diff --git a/ot/backend.py b/ot/backend.py index fc087495c..9d17e14e2 100644 --- a/ot/backend.py +++ b/ot/backend.py @@ -174,6 +174,7 @@ def norm_1d_jax_jvp(primals, tangents): try: import cupy as cp import cupyx + import cupyx.scipy.special cp_type = cp.ndarray except ImportError: @@ -2946,22 +2947,8 @@ def diag(self, a, k=0): def unique(self, a, return_inverse=False): return cp.unique(a, return_inverse=return_inverse) - def logsumexp(self, a, axis=None): - # Taken from - # https://github.com/scipy/scipy/blob/v1.7.1/scipy/special/_logsumexp.py#L7-L127 - a_max = cp.amax(a, axis=axis, keepdims=True) - - if a_max.ndim > 0: - a_max[~cp.isfinite(a_max)] = 0 - elif not cp.isfinite(a_max): - a_max = 0 - - tmp = cp.exp(a - a_max) - s = cp.sum(tmp, axis=axis) - out = cp.log(s) - a_max = cp.squeeze(a_max, axis=axis) - out += a_max - return out + def logsumexp(self, a, axis=None, keepdims=False): + return cupyx.scipy.special.logsumexp(a, axis=axis, keepdims=keepdims) def stack(self, arrays, axis=0): return cp.stack(arrays, axis) @@ -3410,8 +3397,8 @@ def unique(self, a, return_inverse=False): else: return y_prime - def logsumexp(self, a, axis=None): - return tf.math.reduce_logsumexp(a, axis=axis) + def logsumexp(self, a, axis=None, keepdims=False): + return tf.math.reduce_logsumexp(a, axis=axis, keepdims=keepdims) def stack(self, arrays, axis=0): return tnp.stack(arrays, axis) diff --git a/test/test_backend.py b/test/test_backend.py index c88ee5052..2a3602232 100644 --- a/test/test_backend.py +++ b/test/test_backend.py @@ -350,6 +350,9 @@ def test_func_backends(nx): M_complex = M + 1j * rnd.randn(10, 3) v_acos = np.clip(v, -0.99, 0.99) + M_neginf = M.copy() + M_neginf[0, :] = -np.inf + lst_tot = [] for nx in [ot.backend.NumpyBackend(), nx]: @@ -616,6 +619,23 @@ def test_func_backends(nx): lst_b.append(nx.to_numpy(A)) lst_name.append("logsumexp") + A = nx.logsumexp(Mb, axis=0) + lst_b.append(nx.to_numpy(A)) + lst_name.append("logsumexp(axis=0)") + + A = nx.logsumexp(Mb, axis=1) + lst_b.append(nx.to_numpy(A)) + lst_name.append("logsumexp(axis=1)") + + A = nx.logsumexp(Mb, axis=1, keepdims=True) + lst_b.append(nx.to_numpy(A)) + lst_name.append("logsumexp(axis=1, keepdims=True)") + + M_neginf_b = nx.from_numpy(M_neginf) + A = nx.logsumexp(M_neginf_b, axis=1) + lst_b.append(nx.to_numpy(A)) + lst_name.append("logsumexp(all -inf row)") + A = nx.stack([Mb, Mb]) lst_b.append(nx.to_numpy(A)) lst_name.append("stack") From d7b5a5dcdbe83696092ae00cedaabaa00706f458 Mon Sep 17 00:00:00 2001 From: Tom Vercauteren Date: Sun, 20 Sep 2026 14:51:38 +0100 Subject: [PATCH 2/2] Generalize convolutional Wasserstein barycenters to N-D grids convolutional_barycenter2d and convolutional_barycenter2d_debiased were hardcoded to 2D images, using a separate width/height convolution operator (_get_convol_img_fn) and, in log-domain, a Python loop over histograms with in-place writes that jax and tf could not trace. Replace _get_convol_img_fn with a general _SeparableKernel: it applies one 1D Gaussian factor per grid axis via a dense matmul (exp-domain) or a stabilized log-matmul-exp (log-domain), treating any leading axes as batch. This works unchanged for 1D signals, 2D images, 3D volumes, or higher, and processes all histograms in a batch instead of looping, which removes the mutable state that blocked jax/tf under method="sinkhorn_log". Add the N-D public API: convolutional_grid_barycenter and convolutional_grid_barycenter_debiased, for A of shape (n_hists, *grid_shape). convolutional_barycenter2d and convolutional_barycenter2d_debiased become thin wrappers that check A.ndim == 3 and delegate; their signatures, defaults, docstrings and return contract are unchanged, and they are not deprecated. The log=True return contract keeps its existing asymmetry (only the exp-domain solvers populate log["U"]/log["V"]), now applied consistently across all four grid functions. Along the way, fix bugs found while touching this code: array-method calls (.std()/.sum()) instead of nx.std()/nx.sum(), an in-place log_bar += rebind that is invalid under jax/tf, and a dead `F` variable from an unused triple-zeros allocation. Verified: - _SeparableKernel matches the removed _get_convol_img_fn to ~1.8e-15 (exp-domain) and ~4.4e-16 (log-domain, fast) / exactly 0.0 against a module-private exact-logsumexp test helper, on random (4, 9, 7) input. - The fast, shifted log-matmul-exp agrees with the exact per-axis logsumexp reduction to ~1e-15 for reg >= 1e-3 on a moderate grid; for an exact delta input at reg <= 1e-4 it underflows to -inf where the exact reduction stays finite (regression-tested as a known, documented limitation of method="sinkhorn_log"). - jax and tf now genuinely pass sinkhorn_log (previously NotImplementedError), verified in a separate venv with tf installed. - On the 4-image, reg=0.004 example: sinkhorn 230->30ms (7.7x), sinkhorn_log 8653->133ms (65x), debiased sinkhorn 350->48ms (7.3x), debiased sinkhorn_log 12924->235ms (55x), with old-vs-new output agreement to ~1e-17 (see local_sandbox/bench_convolutional_barycenter.py, not part of this commit). Co-Authored-By: Claude Sonnet 5 --- README.md | 2 +- RELEASES.md | 2 + docs/source/user_guide.rst | 6 +- ot/bregman/__init__.py | 4 + ot/bregman/_convolutional.py | 522 +++++++++++++++++++++++++++-------- test/test_bregman.py | 389 +++++++++++++++++--------- 6 files changed, 664 insertions(+), 261 deletions(-) diff --git a/README.md b/README.md index 543e3869e..9b0ae7881 100644 --- a/README.md +++ b/README.md @@ -40,7 +40,7 @@ POT provides the following generic OT solvers: Algorithm](https://pythonot.github.io/auto_examples/plot_OT_1D.html) \[2] , stabilized version \[9] \[10] \[34], lazy CPU/GPU solver from geomloss \[60] \[61], greedy Sinkhorn \[22] and Screening Sinkhorn \[26]. -* Bregman projections for [Wasserstein barycenter](https://pythonot.github.io/auto_examples/barycenters/plot_barycenter_lp_vs_entropic.html) \[3], [convolutional barycenter](https://pythonot.github.io/auto_examples/barycenters/plot_convolutional_barycenter.html) \[21] and unmixing \[4]. +* Bregman projections for [Wasserstein barycenter](https://pythonot.github.io/auto_examples/barycenters/plot_barycenter_lp_vs_entropic.html) \[3], [convolutional barycenter](https://pythonot.github.io/auto_examples/barycenters/plot_convolutional_barycenter.html) \[21] (`ot.bregman.convolutional_grid_barycenter` generalizes it to grids of any dimension, e.g. 1D signals or 3D volumes) and unmixing \[4]. * Sinkhorn divergence \[23] and entropic regularization OT from empirical data. * Debiased Sinkhorn barycenters [Sinkhorn divergence barycenter](https://pythonot.github.io/auto_examples/barycenters/plot_debiased_barycenter.html) \[37] * Smooth optimal transport solvers (dual and semi-dual) for KL and squared L2 regularizations \[17]. diff --git a/RELEASES.md b/RELEASES.md index 0169d261f..1445a22ea 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -4,6 +4,7 @@ #### New features +- Generalize the separable-kernel convolutional Wasserstein barycenter to regular grids of any dimension (1D signals, 3D volumes, ...) via `ot.bregman.convolutional_grid_barycenter` and `ot.bregman.convolutional_grid_barycenter_debiased`; `ot.bregman.convolutional_barycenter2d` and `ot.bregman.convolutional_barycenter2d_debiased` are now thin `A.ndim == 3` wrappers around them, with unchanged signatures, defaults and docstrings. As before, `log["U"]`/`log["V"]` are only populated by the `method="sinkhorn"` solvers, not `"sinkhorn_log"`, and this asymmetry is now consistent across all four grid functions - 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 @@ -15,6 +16,7 @@ #### Closed issues +- Vectorize `ot.bregman.convolutional_barycenter2d`'s log-domain solvers (`method="sinkhorn_log"`) across histograms instead of a per-histogram Python loop with in-place writes, which unblocks Jax and Tensorflow for `sinkhorn_log` (previously a `NotImplementedError`) and speeds up the exp-domain solvers too via a separable, batched kernel application (~7-8x for `sinkhorn` and ~55-65x for `sinkhorn_log` measured on the 4-image, `reg=0.004` example) - Fix `Backend.logsumexp` API conformance: `CupyBackend.logsumexp` and `TensorflowBackend.logsumexp` did not accept `keepdims`, unlike the NumPy/Jax/Torch backends and the base class signature; `CupyBackend.logsumexp` now delegates to `cupyx.scipy.special.logsumexp` instead of a hand-rolled port (Issue #867) - Remove a leftover debug `print` from `ot.utils.projection_sparse_simplex` with `axis=1`, and make the `ot.datasets.make_gauss_hd` docstring a raw string so importing `ot` no longer emits a `SyntaxWarning` (PR #860) - Fix `ot.dist` ignoring the weights `w` for `metric="cityblock"`, which returned the unweighted distance although the weights are documented for this metric (PR #859) diff --git a/docs/source/user_guide.rst b/docs/source/user_guide.rst index ecb314a1a..0182da963 100644 --- a/docs/source/user_guide.rst +++ b/docs/source/user_guide.rst @@ -575,7 +575,11 @@ accelerate the estimation of Wasserstein barycenter when the support has a separable structure [21]_. In the case of 2D images for instance one can replace the matrix vector production in the Bregman projections by convolution operators. We provide an implementation of this algorithm in function -:any:`ot.bregman.convolutional_barycenter2d`. +:any:`ot.bregman.convolutional_barycenter2d`, and its debiased variant [37]_ in +:any:`ot.bregman.convolutional_barycenter2d_debiased`. The same separable-kernel +algorithm generalizes to a regular grid of any dimension (1D signals, 3D +volumes, ...) in :any:`ot.bregman.convolutional_grid_barycenter` and +:any:`ot.bregman.convolutional_grid_barycenter_debiased`. diff --git a/ot/bregman/__init__.py b/ot/bregman/__init__.py index 008c85219..4b87aa814 100644 --- a/ot/bregman/__init__.py +++ b/ot/bregman/__init__.py @@ -32,6 +32,8 @@ from ._convolutional import ( convolutional_barycenter2d, convolutional_barycenter2d_debiased, + convolutional_grid_barycenter, + convolutional_grid_barycenter_debiased, ) from ._empirical import ( @@ -69,6 +71,8 @@ "jcpot_barycenter", "convolutional_barycenter2d", "convolutional_barycenter2d_debiased", + "convolutional_grid_barycenter", + "convolutional_grid_barycenter_debiased", "empirical_sinkhorn", "empirical_sinkhorn2", "empirical_sinkhorn2_geomloss", diff --git a/ot/bregman/_convolutional.py b/ot/bregman/_convolutional.py index 9a8253240..53468fde3 100644 --- a/ot/bregman/_convolutional.py +++ b/ot/bregman/_convolutional.py @@ -20,36 +20,117 @@ ) -def _get_convol_img_fn(nx, width, height, reg, type_as, log_domain=False): - """Return the convolution operator for 2D images. +def _move_axis_perm(ndim, src): + """Return (perm, inv_perm) moving axis `src` to the last position of an + `ndim`-dimensional array, and the permutation that moves it back.""" + perm = tuple(i for i in range(ndim) if i != src) + (src,) + inv_perm = [0] * ndim + for i, p in enumerate(perm): + inv_perm[p] = i + return perm, tuple(inv_perm) - The function constructed is equivalent to blurring on horizontal then vertical directions.""" - t1 = nx.linspace(0, 1, width, type_as=type_as) - Y1, X1 = nx.meshgrid(t1, t1) - M1 = -((X1 - Y1) ** 2) / reg - t2 = nx.linspace(0, 1, height, type_as=type_as) - Y2, X2 = nx.meshgrid(t2, t2) - M2 = -((X2 - Y2) ** 2) / reg +def _log_matmul_exp(nx, y, K): + """Stabilized log(exp(y) @ K). - # If normal domain is selected, we can use M1 and M2 to compute the convolution - if not log_domain: - K1, K2 = nx.exp(M1), nx.exp(M2) + y : array-like, shape (..., n) + K : array-like, shape (n, m), non-negative + """ + c = nx.max(y, axis=-1, keepdims=True) + c = nx.where(nx.isfinite(c), c, nx.zeros(c.shape, type_as=c)) + return c + nx.log(nx.matmul(nx.exp(y - c), K)) - def convol_imgs(imgs): - kx = nx.einsum("...ij,kjl->kil", K1, imgs) - kxy = nx.einsum("...ij,klj->kli", K2, kx) - return kxy - # Else, we can use M1 and M2 to compute the convolution in log-domain - else: +class _SeparableKernel: + r"""Apply the separable kernel :math:`K = K_1 \otimes \dots \otimes K_d` to + the trailing `d` axes of an array. + + Leading axes are treated as batch axes and are left untouched, so a + single instance serves a single grid of shape `(*shape)`, a stack of + shape `(n_hists, *shape)`, or anything with extra leading axes. - def convol_imgs(log_imgs): - log_imgs = nx.logsumexp(M1[:, :, None] + log_imgs[None], axis=1) - log_imgs = nx.logsumexp(M2[:, :, None] + log_imgs.T[None], axis=1).T - return log_imgs + Parameters + ---------- + nx : Backend + backend to use for computations + log_kernels : list of array-like + one `(n_k, n_k)` matrix per grid axis, holding + :math:`-|x-y|^2/\mathrm{reg}` for that axis + log_domain : bool, optional + if True, `__call__` expects and returns log-domain arrays and uses a + stabilized log-matmul-exp; otherwise it expects and returns + exp-domain arrays and applies a plain matrix product + + .. note:: + Each 1-D factor is applied as a dense matrix product (``x @ K``) + rather than a native 1-D convolution. The kernel is Toeplitz, so a + convolution would be mathematically equivalent (checked to ~3e-15 + against ``K @ x`` on a small example), but truncating it is not a + safe optimization for a Sinkhorn iteration: the dual scalings grow + to compensate for the kernel's decay, so the transport plan stays + spread out even where the kernel itself is negligible. On this + package's own convolutional-barycenter example (64x64 images, + `reg=0.004`, :math:`\sigma \approx 2.8` pixels), truncating the + kernel below 12 standard deviations makes the iteration diverge to + NaN. Dense ``matmul`` is also simply faster here: ~4.5 ms against + ~220 ms for a depthwise ``conv2d`` with a 4-sigma-truncated kernel + (Torch 2.14, CPU, float64, 5 images, `n=512`, `reg=4e-3`). + """ - return convol_imgs + def __init__(self, nx, log_kernels, log_domain=False): + self.nx = nx + self.ndim = len(log_kernels) + self.log_domain = log_domain + self.log_kernels = log_kernels + self.kernels = [nx.exp(K) for K in log_kernels] + + def __call__(self, x): + nx = self.nx + n_batch = x.ndim - self.ndim + for axis, K in enumerate(self.kernels): + src = n_batch + axis + perm, inv_perm = _move_axis_perm(x.ndim, src) + x = nx.transpose(x, perm) + if self.log_domain: + x = _log_matmul_exp(nx, x, K) + else: + x = nx.matmul(x, K) + x = nx.transpose(x, inv_perm) + return x + + +def _grid_gaussian_kernel(nx, shape, reg, type_as, log_domain=False): + r"""Gaussian kernel :math:`\exp(-|x-y|^2/\mathrm{reg})` on a regular grid + of the given `shape`, as a :class:`_SeparableKernel` with one factor per + axis. Each axis is sampled on `nx.linspace(0, 1, n)`. + """ + log_kernels = [] + for n in shape: + t = nx.linspace(0, 1, n, type_as=type_as) + Y, X = nx.meshgrid(t, t) + log_kernels.append(-((X - Y) ** 2) / reg) + return _SeparableKernel(nx, log_kernels, log_domain=log_domain) + + +def _exact_separable_log_apply(nx, y, log_kernels): + """Reference (non-stabilized-matmul) separable log-domain kernel + application, computed with a full :any:`Backend.logsumexp` reduction + per axis instead of the shifted log-matmul-exp used by + :class:`_SeparableKernel`. + + This is a test helper only: it materialises a `(..., n, n)` + intermediate per axis and is only tractable on small grids. It is used + to check the fast, shifted form against the exact reduction; it is not + part of the public API and is not used by any solver. + """ + n_batch = y.ndim - len(log_kernels) + for axis, logK in enumerate(log_kernels): + src = n_batch + axis + perm, inv_perm = _move_axis_perm(y.ndim, src) + y = nx.transpose(y, perm) + y = nx.logsumexp(y[..., :, None] + logK, axis=-2) + y = nx.transpose(y, inv_perm) + return y def _print_report(ii, err): @@ -59,7 +140,7 @@ def _print_report(ii, err): print("{:5d}|{:8e}|".format(ii, err)) -def convolutional_barycenter2d( +def convolutional_grid_barycenter( A, reg, weights=None, @@ -71,8 +152,9 @@ def convolutional_barycenter2d( warn=True, **kwargs, ): - r"""Compute the entropic regularized wasserstein barycenter of distributions :math:`\mathbf{A}` - where :math:`\mathbf{A}` is a collection of 2D images. + r"""Compute the entropic regularized Wasserstein barycenter of distributions + :math:`\mathbf{A}` where :math:`\mathbf{A}` is a collection of histograms on a + common regular grid of arbitrary dimension (1D signals, 2D images, 3D volumes, ...). The function solves the following optimization problem: @@ -83,21 +165,24 @@ def convolutional_barycenter2d( - :math:`W_{reg}(\cdot,\cdot)` is the entropic regularized Wasserstein distance (see :py:func:`ot.bregman.sinkhorn`) - - :math:`\mathbf{a}_i` are training distributions (2D images) in the mast two dimensions - of matrix :math:`\mathbf{A}` + - :math:`\mathbf{a}_i` are training distributions on the grid, in the last + `A.ndim - 1` dimensions of matrix :math:`\mathbf{A}` - `reg` is the regularization strength scalar value - The algorithm used for solving the problem is the Sinkhorn-Knopp matrix scaling algorithm - as proposed in :ref:`[21] ` + The algorithm used for solving the problem is the Sinkhorn-Knopp matrix scaling + algorithm as proposed in :ref:`[21] `, + applied through a separable Gaussian kernel on the grid. Parameters ---------- - A : array-like, shape (n_hists, width, height) - `n` distributions (2D images) of size `width` x `height` + A : array-like, shape (n_hists, \*grid_shape) + `n_hists` distributions on a regular grid of shape `grid_shape` + (any number of grid dimensions, e.g. `(width, height)` for images or + `(width, height, depth)` for volumes) reg : float Regularization term >0 weights : array-like, shape (n_hists,) - Weights of each image on the simplex (barycentric coordinates) + Weights of each histogram on the simplex (barycentric coordinates) method : string, optional method used for the solver either 'sinkhorn' or 'sinkhorn_log' numItermax : int, optional @@ -115,13 +200,13 @@ def convolutional_barycenter2d( Returns ------- - a : array-like, shape (width, height) - 2D Wasserstein barycenter + a : array-like, shape (\*grid_shape) + Wasserstein barycenter on the grid log : dict log dictionary return only if log==True in parameters - .. _references-convolutional-barycenter-2d: + .. _references-convolutional-grid-barycenter: References ---------- @@ -135,7 +220,7 @@ def convolutional_barycenter2d( """ if method.lower() == "sinkhorn": - return _convolutional_barycenter2d( + return _convolutional_grid_barycenter( A, reg, weights=weights, @@ -147,7 +232,7 @@ def convolutional_barycenter2d( **kwargs, ) elif method.lower() == "sinkhorn_log": - return _convolutional_barycenter2d_log( + return _convolutional_grid_barycenter_log( A, reg, weights=weights, @@ -162,7 +247,7 @@ def convolutional_barycenter2d( raise ValueError("Unknown method '%s'." % method) -def _convolutional_barycenter2d( +def _convolutional_grid_barycenter( A, reg, weights=None, @@ -173,12 +258,15 @@ def _convolutional_barycenter2d( log=False, warn=True, ): - r"""Compute the entropic regularized wasserstein barycenter of distributions A - where A is a collection of 2D images. + r"""Compute the entropic regularized Wasserstein barycenter of distributions A + where A is a collection of histograms on a regular grid (e.g. unit-normalised + images or volumes). """ A = list_to_array(A) - n_hists, width, height = A.shape + n_hists = A.shape[0] + grid_shape = tuple(A.shape[1:]) + grid_ndim = len(grid_shape) nx = get_backend(A) @@ -186,26 +274,27 @@ def _convolutional_barycenter2d( weights = nx.ones((n_hists,), type_as=A) / n_hists else: assert len(weights) == n_hists + weights = nx.reshape(weights, (n_hists,) + (1,) * grid_ndim) if log: log = {"err": []} - bar = nx.ones((width, height), type_as=A) + bar = nx.ones(grid_shape, type_as=A) bar /= nx.sum(bar) U = nx.ones(A.shape, type_as=A) V = nx.ones(A.shape, type_as=A) err = 1 - # build the convolution operator - convol_imgs = _get_convol_img_fn(nx, width, height, reg, type_as=A) + # build the separable convolution kernel + kernel = _grid_gaussian_kernel(nx, grid_shape, reg, type_as=A) - KU = convol_imgs(U) + KU = kernel(U) for ii in range(numItermax): V = bar[None] / KU - KV = convol_imgs(V) + KV = kernel(V) U = A / KV - KU = convol_imgs(U) - bar = nx.exp(nx.sum(weights[:, None, None] * nx.log(KU + stabThr), axis=0)) + KU = kernel(U) + bar = nx.exp(nx.sum(weights * nx.log(KU + stabThr), axis=0)) if ii % 10 == 9: err = nx.sum(nx.std(V * KU, axis=0)) # log and verbose print @@ -228,7 +317,7 @@ def _convolutional_barycenter2d( return bar -def _convolutional_barycenter2d_log( +def _convolutional_grid_barycenter_log( A, reg, weights=None, @@ -239,47 +328,40 @@ def _convolutional_barycenter2d_log( log=False, warn=True, ): - r"""Compute the entropic regularized wasserstein barycenter of distributions A - where A is a collection of 2D images in log-domain. + r"""Compute the entropic regularized Wasserstein barycenter of distributions A + where A is a collection of histograms on a regular grid (e.g. unit-normalised + images or volumes), in log-domain. """ A = list_to_array(A) + n_hists = A.shape[0] + grid_shape = tuple(A.shape[1:]) + grid_ndim = len(grid_shape) nx = get_backend(A) - # This error is raised because we are using mutable assignment in the line - # `log_KU[k] = ...` which is not allowed in Jax and TF. - if nx.__name__ in ("jax", "tf"): - raise NotImplementedError( - "Log-domain functions are not yet implemented" - " for Jax and TF. Use numpy or torch arrays instead." - ) - - n_hists, width, height = A.shape if weights is None: weights = nx.ones((n_hists,), type_as=A) / n_hists else: assert len(weights) == n_hists + weights = nx.reshape(weights, (n_hists,) + (1,) * grid_ndim) if log: log = {"err": []} - err = 1 - # build the convolution operator - convol_img = _get_convol_img_fn(nx, width, height, reg, type_as=A, log_domain=True) + # build the separable convolution kernel + kernel = _grid_gaussian_kernel(nx, grid_shape, reg, type_as=A, log_domain=True) logA = nx.log(A + stabThr) - log_KU, G, F = nx.zeros((3, *logA.shape), type_as=A) + G = nx.zeros(logA.shape, type_as=A) err = 1 for ii in range(numItermax): - log_bar = nx.zeros((width, height), type_as=A) - for k in range(n_hists): - f = logA[k] - convol_img(G[k]) - log_KU[k] = convol_img(f) - log_bar = log_bar + weights[k] * log_KU[k] + f = logA - kernel(G) + log_KU = kernel(f) + log_bar = nx.sum(weights * log_KU, axis=0) if ii % 10 == 9: - err = nx.exp(G + log_KU).std(axis=0).sum() + err = nx.sum(nx.std(nx.exp(G + log_KU), axis=0)) # log and verbose print if log: log["err"].append(err) @@ -287,7 +369,7 @@ def _convolutional_barycenter2d_log( _print_report(ii, err) if err < stopThr: break - G = log_bar[None, :, :] - log_KU + G = log_bar[None] - log_KU else: if warn: @@ -299,7 +381,103 @@ def _convolutional_barycenter2d_log( return nx.exp(log_bar) -def convolutional_barycenter2d_debiased( +def convolutional_barycenter2d( + A, + reg, + weights=None, + method="sinkhorn", + numItermax=10000, + stopThr=1e-4, + verbose=False, + log=False, + warn=True, + **kwargs, +): + r"""Compute the entropic regularized wasserstein barycenter of distributions :math:`\mathbf{A}` + where :math:`\mathbf{A}` is a collection of 2D images. + + The function solves the following optimization problem: + + .. math:: + \mathbf{a} = \mathop{\arg \min}_\mathbf{a} \quad \sum_i W_{reg}(\mathbf{a},\mathbf{a}_i) + + where : + + - :math:`W_{reg}(\cdot,\cdot)` is the entropic regularized Wasserstein + distance (see :py:func:`ot.bregman.sinkhorn`) + - :math:`\mathbf{a}_i` are training distributions (2D images) in the mast two dimensions + of matrix :math:`\mathbf{A}` + - `reg` is the regularization strength scalar value + + The algorithm used for solving the problem is the Sinkhorn-Knopp matrix scaling algorithm + as proposed in :ref:`[21] ` + + Parameters + ---------- + A : array-like, shape (n_hists, width, height) + `n` distributions (2D images) of size `width` x `height` + reg : float + Regularization term >0 + weights : array-like, shape (n_hists,) + Weights of each image on the simplex (barycentric coordinates) + method : string, optional + method used for the solver either 'sinkhorn' or 'sinkhorn_log' + numItermax : int, optional + Max number of iterations + stopThr : float, optional + Stop threshold on error (> 0) + stabThr : float, optional + Stabilization threshold to avoid numerical precision issue + verbose : bool, optional + Print information along iterations + log : bool, optional + record log if True + warn : bool, optional + if True, raises a warning if the algorithm doesn't convergence. + + Returns + ------- + a : array-like, shape (width, height) + 2D Wasserstein barycenter + log : dict + log dictionary return only if log==True in parameters + + + .. _references-convolutional-barycenter-2d: + References + ---------- + + .. [21] Solomon, J., De Goes, F., Peyré, G., Cuturi, M., Butscher, + A., Nguyen, A. & Guibas, L. (2015). Convolutional wasserstein distances: + Efficient optimal transportation on geometric domains. ACM Transactions + on Graphics (TOG), 34(4), 66 + + .. [37] Janati, H., Cuturi, M., Gramfort, A. Proceedings of the 37th + International Conference on Machine Learning, PMLR 119:4692-4701, 2020 + """ + A = list_to_array(A) + if A.ndim != 3: + raise ValueError( + "convolutional_barycenter2d expects `A` of shape " + f"(n_hists, width, height) (A.ndim == 3), got A.ndim == {A.ndim}. " + "For grids of other dimensions, use " + "ot.bregman.convolutional_grid_barycenter." + ) + return convolutional_grid_barycenter( + A, + reg, + weights=weights, + method=method, + numItermax=numItermax, + stopThr=stopThr, + verbose=verbose, + log=log, + warn=warn, + **kwargs, + ) + + +def convolutional_grid_barycenter_debiased( A, reg, weights=None, @@ -312,7 +490,8 @@ def convolutional_barycenter2d_debiased( **kwargs, ): r"""Compute the debiased sinkhorn barycenter of distributions :math:`\mathbf{A}` - where :math:`\mathbf{A}` is a collection of 2D images. + where :math:`\mathbf{A}` is a collection of histograms on a common regular grid of + arbitrary dimension (1D signals, 2D images, 3D volumes, ...). The function solves the following optimization problem: @@ -323,21 +502,25 @@ def convolutional_barycenter2d_debiased( - :math:`S_{reg}(\cdot,\cdot)` is the debiased entropic regularized Wasserstein distance (see :py:func:`ot.bregman.barycenter_debiased`) - - :math:`\mathbf{a}_i` are training distributions (2D images) in the mast two - dimensions of matrix :math:`\mathbf{A}` + - :math:`\mathbf{a}_i` are training distributions on the grid, in the last + `A.ndim - 1` dimensions of matrix :math:`\mathbf{A}` - `reg` is the regularization strength scalar value The algorithm used for solving the problem is the debiased Sinkhorn scaling - algorithm as proposed in :ref:`[37] ` + algorithm as proposed in + :ref:`[37] `, applied through a + separable Gaussian kernel on the grid. Parameters ---------- - A : array-like, shape (n_hists, width, height) - `n` distributions (2D images) of size `width` x `height` + A : array-like, shape (n_hists, \*grid_shape) + `n_hists` distributions on a regular grid of shape `grid_shape` + (any number of grid dimensions, e.g. `(width, height)` for images or + `(width, height, depth)` for volumes) reg : float Regularization term >0 weights : array-like, shape (n_hists,) - Weights of each image on the simplex (barycentric coordinates) + Weights of each histogram on the simplex (barycentric coordinates) method : string, optional method used for the solver either 'sinkhorn' or 'sinkhorn_log' numItermax : int, optional @@ -356,13 +539,13 @@ def convolutional_barycenter2d_debiased( Returns ------- - a : array-like, shape (width, height) - 2D Wasserstein barycenter + a : array-like, shape (\*grid_shape) + Wasserstein barycenter on the grid log : dict log dictionary return only if log==True in parameters - .. _references-convolutional-barycenter2d-debiased: + .. _references-convolutional-grid-barycenter-debiased: References ---------- @@ -371,7 +554,7 @@ def convolutional_barycenter2d_debiased( """ if method.lower() == "sinkhorn": - return _convolutional_barycenter2d_debiased( + return _convolutional_grid_barycenter_debiased( A, reg, weights=weights, @@ -383,7 +566,7 @@ def convolutional_barycenter2d_debiased( **kwargs, ) elif method.lower() == "sinkhorn_log": - return _convolutional_barycenter2d_debiased_log( + return _convolutional_grid_barycenter_debiased_log( A, reg, weights=weights, @@ -398,7 +581,7 @@ def convolutional_barycenter2d_debiased( raise ValueError("Unknown method '%s'." % method) -def _convolutional_barycenter2d_debiased( +def _convolutional_grid_barycenter_debiased( A, reg, weights=None, @@ -409,10 +592,16 @@ def _convolutional_barycenter2d_debiased( log=False, warn=True, ): - r"""Compute the debiased barycenter of 2D images via sinkhorn convolutions.""" + r"""Compute the debiased barycenter of histograms on a regular grid + (e.g. unit-normalised images or volumes) via sinkhorn convolutions.""" A = list_to_array(A) - n_hists, width, height = A.shape + n_hists = A.shape[0] + grid_shape = tuple(A.shape[1:]) + grid_ndim = len(grid_shape) + grid_size = 1 + for n in grid_shape: + grid_size *= n nx = get_backend(A) @@ -420,30 +609,31 @@ def _convolutional_barycenter2d_debiased( weights = nx.ones((n_hists,), type_as=A) / n_hists else: assert len(weights) == n_hists + weights = nx.reshape(weights, (n_hists,) + (1,) * grid_ndim) if log: log = {"err": []} - bar = nx.ones((width, height), type_as=A) - bar /= width * height + bar = nx.ones(grid_shape, type_as=A) + bar /= grid_size U = nx.ones(A.shape, type_as=A) V = nx.ones(A.shape, type_as=A) - c = nx.ones((width, height), type_as=A) + c = nx.ones(grid_shape, type_as=A) err = 1 - # build the convolution operator - convol_imgs = _get_convol_img_fn(nx, width, height, reg, type_as=A) + # build the separable convolution kernel + kernel = _grid_gaussian_kernel(nx, grid_shape, reg, type_as=A) - KU = convol_imgs(U) + KU = kernel(U) for ii in range(numItermax): V = bar[None] / KU - KV = convol_imgs(V) + KV = kernel(V) U = A / KV - KU = convol_imgs(U) - bar = c * nx.exp(nx.sum(weights[:, None, None] * nx.log(KU + stabThr), axis=0)) + KU = kernel(U) + bar = c * nx.exp(nx.sum(weights * nx.log(KU + stabThr), axis=0)) for _ in range(10): - c = (c * bar / nx.squeeze(convol_imgs(c[None]))) ** 0.5 + c = (c * bar / nx.squeeze(kernel(c[None]))) ** 0.5 if ii % 10 == 9: err = nx.sum(nx.std(V * KU, axis=0)) @@ -469,7 +659,7 @@ def _convolutional_barycenter2d_debiased( return bar -def _convolutional_barycenter2d_debiased_log( +def _convolutional_grid_barycenter_debiased_log( A, reg, weights=None, @@ -480,43 +670,39 @@ def _convolutional_barycenter2d_debiased_log( log=False, warn=True, ): - r"""Compute the debiased barycenter of 2D images in log-domain.""" + r"""Compute the debiased barycenter of histograms on a regular grid + (e.g. unit-normalised images or volumes) in log-domain.""" A = list_to_array(A) - n_hists, width, height = A.shape + n_hists = A.shape[0] + grid_shape = tuple(A.shape[1:]) + grid_ndim = len(grid_shape) + nx = get_backend(A) - # This error is raised because we are using mutable assignment in the line - # `log_KU[k] = ...` which is not allowed in Jax and TF. - if nx.__name__ in ("jax", "tf"): - raise NotImplementedError( - "Log-domain functions are not yet implemented" - " for Jax and TF. Use numpy or torch arrays instead." - ) + if weights is None: weights = nx.ones((n_hists,), type_as=A) / n_hists else: - assert len(weights) == A.shape[0] + assert len(weights) == n_hists + weights = nx.reshape(weights, (n_hists,) + (1,) * grid_ndim) if log: log = {"err": []} - err = 1 - # build the convolution operator - convol_img = _get_convol_img_fn(nx, width, height, reg, type_as=A, log_domain=True) + # build the separable convolution kernel + kernel = _grid_gaussian_kernel(nx, grid_shape, reg, type_as=A, log_domain=True) logA = nx.log(A + stabThr) - log_bar, c = nx.zeros((2, width, height), type_as=A) - log_KU, G, F = nx.zeros((3, *logA.shape), type_as=A) + G = nx.zeros(logA.shape, type_as=A) + c = nx.zeros(grid_shape, type_as=A) err = 1 for ii in range(numItermax): - log_bar = nx.zeros((width, height), type_as=A) - for k in range(n_hists): - f = logA[k] - convol_img(G[k]) - log_KU[k] = convol_img(f) - log_bar = log_bar + weights[k] * log_KU[k] - log_bar += c + f = logA - kernel(G) + log_KU = kernel(f) + log_bar = nx.sum(weights * log_KU, axis=0) + c + for _ in range(10): - c = 0.5 * (c + log_bar - convol_img(c)) + c = 0.5 * (c + log_bar - kernel(c)) if ii % 10 == 9: err = nx.sum(nx.std(nx.exp(G + log_KU), axis=0)) @@ -527,7 +713,7 @@ def _convolutional_barycenter2d_debiased_log( _print_report(ii, err) if err < stopThr and ii > 20: break - G = log_bar[None, :, :] - log_KU + G = log_bar[None] - log_KU else: if warn: @@ -537,3 +723,95 @@ def _convolutional_barycenter2d_debiased_log( return nx.exp(log_bar), log else: return nx.exp(log_bar) + + +def convolutional_barycenter2d_debiased( + A, + reg, + weights=None, + method="sinkhorn", + numItermax=10000, + stopThr=1e-3, + verbose=False, + log=False, + warn=True, + **kwargs, +): + r"""Compute the debiased sinkhorn barycenter of distributions :math:`\mathbf{A}` + where :math:`\mathbf{A}` is a collection of 2D images. + + The function solves the following optimization problem: + + .. math:: + \mathbf{a} = \mathop{\arg \min}_\mathbf{a} \quad \sum_i S_{reg}(\mathbf{a},\mathbf{a}_i) + + where : + + - :math:`S_{reg}(\cdot,\cdot)` is the debiased entropic regularized Wasserstein + distance (see :py:func:`ot.bregman.barycenter_debiased`) + - :math:`\mathbf{a}_i` are training distributions (2D images) in the mast two + dimensions of matrix :math:`\mathbf{A}` + - `reg` is the regularization strength scalar value + + The algorithm used for solving the problem is the debiased Sinkhorn scaling + algorithm as proposed in :ref:`[37] ` + + Parameters + ---------- + A : array-like, shape (n_hists, width, height) + `n` distributions (2D images) of size `width` x `height` + reg : float + Regularization term >0 + weights : array-like, shape (n_hists,) + Weights of each image on the simplex (barycentric coordinates) + method : string, optional + method used for the solver either 'sinkhorn' or 'sinkhorn_log' + numItermax : int, optional + Max number of iterations + stopThr : float, optional + Stop threshold on error (> 0) + stabThr : float, optional + Stabilization threshold to avoid numerical precision issue + verbose : bool, optional + Print information along iterations + log : bool, optional + record log if True + warn : bool, optional + if True, raises a warning if the algorithm doesn't convergence. + + + Returns + ------- + a : array-like, shape (width, height) + 2D Wasserstein barycenter + log : dict + log dictionary return only if log==True in parameters + + + .. _references-convolutional-barycenter2d-debiased: + References + ---------- + + .. [37] Janati, H., Cuturi, M., Gramfort, A. Proceedings of the 37th International + Conference on Machine Learning, PMLR 119:4692-4701, 2020 + """ + A = list_to_array(A) + if A.ndim != 3: + raise ValueError( + "convolutional_barycenter2d_debiased expects `A` of shape " + f"(n_hists, width, height) (A.ndim == 3), got A.ndim == {A.ndim}. " + "For grids of other dimensions, use " + "ot.bregman.convolutional_grid_barycenter_debiased." + ) + return convolutional_grid_barycenter_debiased( + A, + reg, + weights=weights, + method=method, + numItermax=numItermax, + stopThr=stopThr, + verbose=verbose, + log=log, + warn=warn, + **kwargs, + ) diff --git a/test/test_bregman.py b/test/test_bregman.py index 17b400306..e22508fe5 100644 --- a/test/test_bregman.py +++ b/test/test_bregman.py @@ -825,22 +825,18 @@ def test_wasserstein_bary_2d(nx, method): # wasserstein reg = 1e-2 - if nx.__name__ in ("jax", "tf") and method == "sinkhorn_log": - with pytest.raises(NotImplementedError): - ot.bregman.convolutional_barycenter2d(A_nx, reg, method=method) - else: - bary_wass_np, log_np = ot.bregman.convolutional_barycenter2d( - A, reg, method=method, verbose=True, log=True - ) - bary_wass = nx.to_numpy( - ot.bregman.convolutional_barycenter2d(A_nx, reg, method=method) - ) + bary_wass_np, log_np = ot.bregman.convolutional_barycenter2d( + A, reg, method=method, verbose=True, log=True + ) + bary_wass = nx.to_numpy( + ot.bregman.convolutional_barycenter2d(A_nx, reg, method=method) + ) - np.testing.assert_allclose(1, np.sum(bary_wass), rtol=1e-3) - np.testing.assert_allclose(bary_wass, bary_wass_np, atol=1e-3) + np.testing.assert_allclose(1, np.sum(bary_wass), rtol=1e-3) + np.testing.assert_allclose(bary_wass, bary_wass_np, atol=1e-3) - # help in checking if log and verbose do not bug the function - ot.bregman.convolutional_barycenter2d(A, reg, log=True, verbose=True) + # help in checking if log and verbose do not bug the function + ot.bregman.convolutional_barycenter2d(A, reg, log=True, verbose=True) @pytest.skip_backend("tf") @@ -856,27 +852,23 @@ def test_wasserstein_bary_2d_dtype_device(nx, method): # wasserstein reg = 1e-2 - if nx.__name__ in ("jax", "tf") and method == "sinkhorn_log": - with pytest.raises(NotImplementedError): - ot.bregman.convolutional_barycenter2d(Ab, reg, method=method) - else: - # Compute the barycenter with numpy - bary_wass_np, log_np = ot.bregman.convolutional_barycenter2d( - A, reg, method=method, verbose=True, log=True - ) - # Compute the barycenter with the backend - bary_wass_b = ot.bregman.convolutional_barycenter2d(Ab, reg, method=method) - # Convert the backend result to numpy, to compare with the numpy result - bary_wass = nx.to_numpy(bary_wass_b) + # Compute the barycenter with numpy + bary_wass_np, log_np = ot.bregman.convolutional_barycenter2d( + A, reg, method=method, verbose=True, log=True + ) + # Compute the barycenter with the backend + bary_wass_b = ot.bregman.convolutional_barycenter2d(Ab, reg, method=method) + # Convert the backend result to numpy, to compare with the numpy result + bary_wass = nx.to_numpy(bary_wass_b) - np.testing.assert_allclose(1, np.sum(bary_wass), rtol=1e-3) - np.testing.assert_allclose(bary_wass, bary_wass_np, atol=1e-3) + np.testing.assert_allclose(1, np.sum(bary_wass), rtol=1e-3) + np.testing.assert_allclose(bary_wass, bary_wass_np, atol=1e-3) - # help in checking if log and verbose do not bug the function - ot.bregman.convolutional_barycenter2d(A, reg, log=True, verbose=True) + # help in checking if log and verbose do not bug the function + ot.bregman.convolutional_barycenter2d(A, reg, log=True, verbose=True) - # Test that the dtype and device are the same after the computation - nx.assert_same_dtype_device(Ab, bary_wass_b) + # Test that the dtype and device are the same after the computation + nx.assert_same_dtype_device(Ab, bary_wass_b) @pytest.mark.skipif(not tf, reason="tf not installed") @@ -894,37 +886,6 @@ def test_wasserstein_bary_2d_device_tf(method): # wasserstein reg = 1e-2 - if method == "sinkhorn_log": - with pytest.raises(NotImplementedError): - ot.bregman.convolutional_barycenter2d(Ab, reg, method=method) - else: - # Compute the barycenter with numpy - bary_wass_np, log_np = ot.bregman.convolutional_barycenter2d( - A, reg, method=method, verbose=True, log=True - ) - # Compute the barycenter with the backend - bary_wass_b = ot.bregman.convolutional_barycenter2d(Ab, reg, method=method) - # Convert the backend result to numpy, to compare with the numpy result - bary_wass = nx.to_numpy(bary_wass_b) - - np.testing.assert_allclose(1, np.sum(bary_wass), rtol=1e-3) - np.testing.assert_allclose(bary_wass, bary_wass_np, atol=1e-3) - - # help in checking if log and verbose do not bug the function - ot.bregman.convolutional_barycenter2d(A, reg, log=True, verbose=True) - - # Test that the dtype and device are the same after the computation - nx.assert_same_dtype_device(Ab, bary_wass_b) - - # Check that everything happens on the GPU - Ab = nx.from_numpy(A) - - # wasserstein - reg = 1e-2 - if method == "sinkhorn_log": - with pytest.raises(NotImplementedError): - ot.bregman.convolutional_barycenter2d(Ab, reg, method=method) - else: # Compute the barycenter with numpy bary_wass_np, log_np = ot.bregman.convolutional_barycenter2d( A, reg, method=method, verbose=True, log=True @@ -943,9 +904,32 @@ def test_wasserstein_bary_2d_device_tf(method): # Test that the dtype and device are the same after the computation nx.assert_same_dtype_device(Ab, bary_wass_b) - # Check this only if GPU is available - if len(tf.config.list_physical_devices("GPU")) > 0: - assert nx.dtype_device(bary_wass_b)[1].startswith("GPU") + # Check that everything happens on the GPU + Ab = nx.from_numpy(A) + + # wasserstein + reg = 1e-2 + # Compute the barycenter with numpy + bary_wass_np, log_np = ot.bregman.convolutional_barycenter2d( + A, reg, method=method, verbose=True, log=True + ) + # Compute the barycenter with the backend + bary_wass_b = ot.bregman.convolutional_barycenter2d(Ab, reg, method=method) + # Convert the backend result to numpy, to compare with the numpy result + bary_wass = nx.to_numpy(bary_wass_b) + + np.testing.assert_allclose(1, np.sum(bary_wass), rtol=1e-3) + np.testing.assert_allclose(bary_wass, bary_wass_np, atol=1e-3) + + # help in checking if log and verbose do not bug the function + ot.bregman.convolutional_barycenter2d(A, reg, log=True, verbose=True) + + # Test that the dtype and device are the same after the computation + nx.assert_same_dtype_device(Ab, bary_wass_b) + + # Check this only if GPU is available + if len(tf.config.list_physical_devices("GPU")) > 0: + assert nx.dtype_device(bary_wass_b)[1].startswith("GPU") @pytest.mark.parametrize("method", ["sinkhorn", "sinkhorn_log"]) @@ -957,22 +941,18 @@ def test_wasserstein_bary_2d_debiased(nx, method): # wasserstein reg = 1e-2 - if nx.__name__ in ("jax", "tf") and method == "sinkhorn_log": - with pytest.raises(NotImplementedError): - ot.bregman.convolutional_barycenter2d_debiased(A_nx, reg, method=method) - else: - bary_wass_np, log_np = ot.bregman.convolutional_barycenter2d_debiased( - A, reg, method=method, verbose=True, log=True - ) - bary_wass = nx.to_numpy( - ot.bregman.convolutional_barycenter2d_debiased(A_nx, reg, method=method) - ) + bary_wass_np, log_np = ot.bregman.convolutional_barycenter2d_debiased( + A, reg, method=method, verbose=True, log=True + ) + bary_wass = nx.to_numpy( + ot.bregman.convolutional_barycenter2d_debiased(A_nx, reg, method=method) + ) - np.testing.assert_allclose(1, np.sum(bary_wass), rtol=1e-3) - np.testing.assert_allclose(bary_wass, bary_wass_np, atol=1e-3) + np.testing.assert_allclose(1, np.sum(bary_wass), rtol=1e-3) + np.testing.assert_allclose(bary_wass, bary_wass_np, atol=1e-3) - # help in checking if log and verbose do not bug the function - ot.bregman.convolutional_barycenter2d_debiased(A, reg, log=True, verbose=True) + # help in checking if log and verbose do not bug the function + ot.bregman.convolutional_barycenter2d_debiased(A, reg, log=True, verbose=True) @pytest.skip_backend("tf") @@ -988,31 +968,25 @@ def test_wasserstein_bary_2d_debiased_dtype_device(nx, method): # wasserstein reg = 1e-2 - if nx.__name__ in ("jax", "tf") and method == "sinkhorn_log": - with pytest.raises(NotImplementedError): - ot.bregman.convolutional_barycenter2d_debiased(Ab, reg, method=method) - else: - # Compute the barycenter with numpy - bary_wass_np, log_np = ot.bregman.convolutional_barycenter2d_debiased( - A, reg, method=method, verbose=True, log=True - ) - # Compute the barycenter with the backend - bary_wass_b = ot.bregman.convolutional_barycenter2d_debiased( - Ab, reg, method=method - ) - # Convert the backend result to numpy, to compare with the numpy result - bary_wass = nx.to_numpy(bary_wass_b) + # Compute the barycenter with numpy + bary_wass_np, log_np = ot.bregman.convolutional_barycenter2d_debiased( + A, reg, method=method, verbose=True, log=True + ) + # Compute the barycenter with the backend + bary_wass_b = ot.bregman.convolutional_barycenter2d_debiased( + Ab, reg, method=method + ) + # Convert the backend result to numpy, to compare with the numpy result + bary_wass = nx.to_numpy(bary_wass_b) - np.testing.assert_allclose(1, np.sum(bary_wass), rtol=1e-3) - np.testing.assert_allclose(bary_wass, bary_wass_np, atol=1e-3) + np.testing.assert_allclose(1, np.sum(bary_wass), rtol=1e-3) + np.testing.assert_allclose(bary_wass, bary_wass_np, atol=1e-3) - # help in checking if log and verbose do not bug the function - ot.bregman.convolutional_barycenter2d_debiased( - A, reg, log=True, verbose=True - ) + # help in checking if log and verbose do not bug the function + ot.bregman.convolutional_barycenter2d_debiased(A, reg, log=True, verbose=True) - # Test that the dtype and device are the same after the computation - nx.assert_same_dtype_device(Ab, bary_wass_b) + # Test that the dtype and device are the same after the computation + nx.assert_same_dtype_device(Ab, bary_wass_b) @pytest.mark.skipif(not tf, reason="tf not installed") @@ -1030,41 +1004,6 @@ def test_wasserstein_bary_2d_debiased_device_tf(method): # wasserstein reg = 1e-2 - if method == "sinkhorn_log": - with pytest.raises(NotImplementedError): - ot.bregman.convolutional_barycenter2d_debiased(Ab, reg, method=method) - else: - # Compute the barycenter with numpy - bary_wass_np, log_np = ot.bregman.convolutional_barycenter2d_debiased( - A, reg, method=method, verbose=True, log=True - ) - # Compute the barycenter with the backend - bary_wass_b = ot.bregman.convolutional_barycenter2d_debiased( - Ab, reg, method=method - ) - # Convert the backend result to numpy, to compare with the numpy result - bary_wass = nx.to_numpy(bary_wass_b) - - np.testing.assert_allclose(1, np.sum(bary_wass), rtol=1e-3) - np.testing.assert_allclose(bary_wass, bary_wass_np, atol=1e-3) - - # help in checking if log and verbose do not bug the function - ot.bregman.convolutional_barycenter2d_debiased( - A, reg, log=True, verbose=True - ) - - # Test that the dtype and device are the same after the computation - nx.assert_same_dtype_device(Ab, bary_wass_b) - - # Check that everything happens on the GPU - Ab = nx.from_numpy(A) - - # wasserstein - reg = 1e-2 - if method == "sinkhorn_log": - with pytest.raises(NotImplementedError): - ot.bregman.convolutional_barycenter2d_debiased(Ab, reg, method=method) - else: # Compute the barycenter with numpy bary_wass_np, log_np = ot.bregman.convolutional_barycenter2d_debiased( A, reg, method=method, verbose=True, log=True @@ -1077,6 +1016,29 @@ def test_wasserstein_bary_2d_debiased_device_tf(method): bary_wass = nx.to_numpy(bary_wass_b) np.testing.assert_allclose(1, np.sum(bary_wass), rtol=1e-3) + np.testing.assert_allclose(bary_wass, bary_wass_np, atol=1e-3) + + # help in checking if log and verbose do not bug the function + ot.bregman.convolutional_barycenter2d_debiased(A, reg, log=True, verbose=True) + + # Test that the dtype and device are the same after the computation + nx.assert_same_dtype_device(Ab, bary_wass_b) + + # Check that everything happens on the GPU + Ab = nx.from_numpy(A) + + # wasserstein + reg = 1e-2 + # Compute the barycenter with numpy + bary_wass_np, log_np = ot.bregman.convolutional_barycenter2d_debiased( + A, reg, method=method, verbose=True, log=True + ) + # Compute the barycenter with the backend + bary_wass_b = ot.bregman.convolutional_barycenter2d_debiased(Ab, reg, method=method) + # Convert the backend result to numpy, to compare with the numpy result + bary_wass = nx.to_numpy(bary_wass_b) + + np.testing.assert_allclose(1, np.sum(bary_wass), rtol=1e-3) def test_unmix(nx): @@ -1447,6 +1409,159 @@ def test_convolutional_barycenter_non_square(nx): np.testing.assert_allclose(b, b_np) +@pytest.mark.parametrize( + "method, debiased", + product(["sinkhorn", "sinkhorn_log"], [False, True]), +) +def test_convolutional_grid_barycenter_1d(nx, method, debiased): + # 1D grid: the separable-kernel refactor should work for any grid ndim, + # not just 2D images. + rng = np.random.RandomState(42) + n_hists, n = 3, 40 + A = rng.rand(n_hists, n) + 0.1 + A = A / A.sum(axis=1, keepdims=True) + A_nx = nx.from_numpy(A) + reg = 5e-2 + + fun = ( + ot.bregman.convolutional_grid_barycenter_debiased + if debiased + else ot.bregman.convolutional_grid_barycenter + ) + + bar = nx.to_numpy(fun(A_nx, reg, method=method, numItermax=500)) + assert bar.shape == (n,) + assert np.all(bar >= -1e-8) + np.testing.assert_allclose(1, np.sum(bar), rtol=1e-2) + + # a barycenter with weights concentrated on one input approximately + # recovers that input + weights = np.zeros(n_hists) + weights[0] = 1.0 + weights_nx = nx.from_numpy(weights) + bar_single = nx.to_numpy( + fun(A_nx, reg, weights=weights_nx, method=method, numItermax=500) + ) + np.testing.assert_allclose(bar_single, A[0], atol=3e-2) + + +@pytest.mark.parametrize( + "method, debiased", + product(["sinkhorn", "sinkhorn_log"], [False, True]), +) +def test_convolutional_grid_barycenter_3d(nx, method, debiased): + # 3D grid (volumetric): the separable-kernel refactor should generalize + # beyond the 2D case handled by convolutional_barycenter2d. + rng = np.random.RandomState(42) + n_hists, shape = 3, (12, 13, 14) + A = rng.rand(n_hists, *shape) + 0.1 + A = A / A.sum(axis=(1, 2, 3), keepdims=True) + A_nx = nx.from_numpy(A) + reg = 1e-1 + + fun = ( + ot.bregman.convolutional_grid_barycenter_debiased + if debiased + else ot.bregman.convolutional_grid_barycenter + ) + + bar = nx.to_numpy(fun(A_nx, reg, method=method, numItermax=200)) + assert bar.shape == shape + assert np.all(bar >= -1e-6) + np.testing.assert_allclose(1, np.sum(bar), rtol=1e-2) + + # a barycenter with weights concentrated on one input approximately + # recovers that input + weights = np.zeros(n_hists) + weights[0] = 1.0 + weights_nx = nx.from_numpy(weights) + bar_single = nx.to_numpy( + fun(A_nx, reg, weights=weights_nx, method=method, numItermax=200) + ) + np.testing.assert_allclose(bar_single, A[0], atol=5e-2) + + +def test_convolutional_barycenter2d_matches_grid_barycenter(nx): + # convolutional_barycenter2d{,_debiased} are thin wrappers around + # convolutional_grid_barycenter{,_debiased} for A.ndim == 3 + rng = np.random.RandomState(0) + A = rng.rand(3, 10, 12) + 0.1 + A = A / A.sum(axis=(1, 2), keepdims=True) + A_nx = nx.from_numpy(A) + reg = 1e-2 + + b_2d = nx.to_numpy(ot.bregman.convolutional_barycenter2d(A_nx, reg)) + b_grid = nx.to_numpy(ot.bregman.convolutional_grid_barycenter(A_nx, reg)) + np.testing.assert_allclose(b_2d, b_grid) + + b_2d_deb = nx.to_numpy(ot.bregman.convolutional_barycenter2d_debiased(A_nx, reg)) + b_grid_deb = nx.to_numpy( + ot.bregman.convolutional_grid_barycenter_debiased(A_nx, reg) + ) + np.testing.assert_allclose(b_2d_deb, b_grid_deb) + + A_1d = nx.from_numpy(rng.rand(3, 10)) + with pytest.raises(ValueError): + ot.bregman.convolutional_barycenter2d(A_1d, reg) + with pytest.raises(ValueError): + ot.bregman.convolutional_barycenter2d_debiased(A_1d, reg) + + A_4d = nx.from_numpy(rng.rand(3, 4, 5, 6)) + with pytest.raises(ValueError): + ot.bregman.convolutional_barycenter2d(A_4d, reg) + with pytest.raises(ValueError): + ot.bregman.convolutional_barycenter2d_debiased(A_4d, reg) + + +def test_separable_kernel_log_domain_matches_exact_logsumexp(nx): + # The shifted log-matmul-exp used by _SeparableKernel in log-domain is a + # numerics change relative to the exact per-axis logsumexp reduction; it + # should still agree with it closely for reg large enough not to underflow. + from ot.bregman._convolutional import ( + _exact_separable_log_apply, + _grid_gaussian_kernel, + ) + + rng = np.random.RandomState(0) + shape = (18, 15) + A = rng.rand(4, *shape) + 0.1 + A = A / A.sum(axis=(1, 2), keepdims=True) + A_nx = nx.from_numpy(A) + logA = nx.log(A_nx) + + for reg in [1e-1, 1e-2, 1e-3]: + kernel = _grid_gaussian_kernel(nx, shape, reg, type_as=A_nx, log_domain=True) + fast = nx.to_numpy(kernel(logA)) + exact = nx.to_numpy(_exact_separable_log_apply(nx, logA, kernel.log_kernels)) + np.testing.assert_allclose(fast, exact, atol=1e-9) + + +def test_separable_kernel_log_domain_underflow_limitation(nx): + # Known limitation of the shifted log-matmul-exp: for an exact delta + # input and reg small enough, it underflows to -inf far from the delta + # where the exact reduction still returns a finite (very negative) value. + from ot.bregman._convolutional import ( + _exact_separable_log_apply, + _grid_gaussian_kernel, + ) + + n = 20 + delta = np.zeros((1, n, n)) + delta[0, n // 2, n // 2] = 1.0 + delta_nx = nx.from_numpy(delta) + with np.errstate(divide="ignore"): + log_delta = nx.log(delta_nx) + + reg = 1e-5 + kernel = _grid_gaussian_kernel(nx, (n, n), reg, type_as=delta_nx, log_domain=True) + fast = nx.to_numpy(kernel(log_delta)) + exact = nx.to_numpy(_exact_separable_log_apply(nx, log_delta, kernel.log_kernels)) + + far_corner = (0, 0, 0) + assert np.isneginf(fast[far_corner]) + assert np.isfinite(exact[far_corner]) + + def test_sinkhorn_warmstart(): m, n = 10, 20 a = ot.unif(m)