[MRG] Warn when the convolutional barycenter kernel underflows (closes #458)

RELEASES.md

@@ -15,6 +15,7 @@
 
 #### Closed issues
 
+- Warn when the convolution kernel in `ot.bregman.convolutional_barycenter2d` and `convolutional_barycenter2d_debiased` underflows at small `reg`. Past that point mass can no longer cross the image and the barycenter collapses towards the arithmetic mean of the inputs, which looks like an over-diffuse result rather than an error; `method='sinkhorn_log'` is exact in that regime (PR #873, Issue #458)
 - 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)

ot/bregman/_convolutional.py

@@ -10,6 +10,8 @@
 
 import warnings
 
+import numpy as np
+
 from ..backend import get_backend
 from ..utils import list_to_array
 
@@ -20,7 +22,39 @@
 )
 
 
-def _get_convol_img_fn(nx, width, height, reg, type_as, log_domain=False):
+def _warn_if_kernel_underflows(nx, reg, type_as, stacklevel):
+    """Warn when exp(M) is about to lose the transport to underflow.
+
+    The Sinkhorn iterations form products and ratios of kernel entries, so the
+    usable exponent range is roughly half that of the float type. Past that the
+    convolution can no longer move mass across the image and the barycenter
+    degenerates towards the arithmetic mean of the inputs, which looks like an
+    over-diffuse result rather than an error.
+
+    The grid is always ``linspace(0, 1, n)``, so the most negative exponent is
+    ``-1 / reg`` whatever the image size, and no reduction over the kernel is
+    needed.
+    """
+    try:
+        dtype = nx.to_numpy(nx.zeros((1,), type_as=type_as)).dtype
+        tiny = np.finfo(dtype).tiny
+    except (TypeError, ValueError):  # pragma: no cover - exotic dtypes
+        return
+    min_exponent = -1.0 / reg
+    # half the exponent range, i.e. the exponent of sqrt(tiny)
+    safe_exponent = np.log(tiny) / 2
+    if min_exponent < safe_exponent:
+        warnings.warn(
+            f"reg={reg:g} is small enough that the convolution kernel "
+            f"underflows: its smallest exponent is {min_exponent:.0f} against a "
+            f"usable limit of {safe_exponent:.0f}. The result will be too "
+            "diffuse, and more iterations will not help. Use "
+            "method='sinkhorn_log' for this regularization.",
+            stacklevel=stacklevel,
+        )
+
+
+def _get_convol_img_fn(nx, width, height, reg, type_as, log_domain=False, stacklevel=4):
     """Return the convolution operator for 2D images.
 
     The function constructed is equivalent to blurring on horizontal then vertical directions."""
@@ -34,6 +68,7 @@ def _get_convol_img_fn(nx, width, height, reg, type_as, log_domain=False):
 
     # If normal domain is selected, we can use M1 and M2 to compute the convolution
     if not log_domain:
+        _warn_if_kernel_underflows(nx, reg, type_as, stacklevel=stacklevel + 1)
         K1, K2 = nx.exp(M1), nx.exp(M2)
 
         def convol_imgs(imgs):

test/test_bregman.py

@@ -1434,13 +1434,74 @@ def test_screenkhorn(nx):
     np.testing.assert_allclose(G_sink.sum(1), G_screen.sum(1), atol=1e-02)
 
 
+def test_convolutional_barycenter_kernel_underflow_warns():
+    """Small reg silently loses the transport (issue #458).
+
+    exp(-(x-y)**2 / reg) underflows for distant pixels, mass can no longer
+    cross the image, and the barycenter collapses towards the arithmetic mean
+    of the inputs, which reads as an over-diffuse result rather than an error.
+    """
+    rng = np.random.RandomState(0)
+    n, sigma, sep = 32, 0.06, 0.15
+    t = np.linspace(0, 1, n)
+    X, Y = np.meshgrid(t, t, indexing="ij")
+
+    def gauss(cx):
+        g = np.exp(-((X - cx) ** 2 + (Y - 0.5) ** 2) / (2 * sigma**2))
+        return g / g.sum()
+
+    A = np.stack([gauss(0.5 - sep), gauss(0.5 + sep)])
+
+    with pytest.warns(UserWarning, match="underflow"):
+        ot.bregman.convolutional_barycenter2d_debiased(A, 1e-04)
+
+    with pytest.warns(UserWarning, match="underflow"):
+        ot.bregman.convolutional_barycenter2d(A, 1e-04)
+
+    # a usable kernel must stay silent, and so must the log-domain solver
+    with warnings.catch_warnings():
+        warnings.simplefilter("error", UserWarning)
+        ot.bregman.convolutional_barycenter2d_debiased(A, 1e-02)
+        ot.bregman.convolutional_barycenter2d_debiased(A, 1e-04, method="sinkhorn_log")
+
+
+def test_convolutional_barycenter_debiased_preserves_width():
+    """The debiased barycenter of two equal-width Gaussians keeps that width.
+
+    Janati et al. 2020. The log-domain solver gets this right at every reg;
+    the default one only where its kernel has not underflowed.
+    """
+    n, sigma, sep = 32, 0.06, 0.15
+    t = np.linspace(0, 1, n)
+    X, Y = np.meshgrid(t, t, indexing="ij")
+
+    def gauss(cx):
+        g = np.exp(-((X - cx) ** 2 + (Y - 0.5) ** 2) / (2 * sigma**2))
+        return g / g.sum()
+
+    A = np.stack([gauss(0.5 - sep), gauss(0.5 + sep)])
+
+    def width(img):
+        px = img.sum(axis=1)
+        mx = (px * t).sum()
+        return np.sqrt(((t - mx) ** 2 * px).sum())
+
+    bar = ot.bregman.convolutional_barycenter2d_debiased(
+        A, 1e-03, method="sinkhorn_log"
+    )
+    np.testing.assert_allclose(width(bar), width(A[0]), rtol=0.05)
+
+
 def test_convolutional_barycenter_non_square(nx):
     # test for image with height not equal width
     A = np.ones((2, 2, 3)) / (2 * 3)
     A_nx = nx.from_numpy(A)
 
-    b_np = ot.bregman.convolutional_barycenter2d(A, 1e-03)
-    b = nx.to_numpy(ot.bregman.convolutional_barycenter2d(A_nx, 1e-03))
+    # reg=1e-3 underflows the convolution kernel on a unit grid, which does not
+    # affect a uniform image but does emit a warning; 1e-2 exercises the same
+    # non-square code path with a usable kernel
+    b_np = ot.bregman.convolutional_barycenter2d(A, 1e-02)
+    b = nx.to_numpy(ot.bregman.convolutional_barycenter2d(A_nx, 1e-02))
 
     np.testing.assert_allclose(np.ones((2, 3)) / (2 * 3), b, atol=1e-02)
     np.testing.assert_allclose(np.ones((2, 3)) / (2 * 3), b, atol=1e-02)