diff options
-rw-r--r-- | ot/bregman.py | 8 | ||||
-rw-r--r-- | test/test_bregman.py | 7 |
2 files changed, 14 insertions, 1 deletions
diff --git a/ot/bregman.py b/ot/bregman.py index 2707b7c..d5e3563 100644 --- a/ot/bregman.py +++ b/ot/bregman.py @@ -9,6 +9,7 @@ Bregman projections for regularized OT # Titouan Vayer <titouan.vayer@irisa.fr> # Hicham Janati <hicham.janati@inria.fr> # Mokhtar Z. Alaya <mokhtarzahdi.alaya@gmail.com> +# Alexander Tong <alexander.tong@yale.edu> # # License: MIT License @@ -1346,12 +1347,17 @@ def convolutional_barycenter2d(A, reg, weights=None, numItermax=10000, err = 1 # build the convolution operator + # this is equivalent to blurring on horizontal then vertical directions t = np.linspace(0, 1, A.shape[1]) [Y, X] = np.meshgrid(t, t) xi1 = np.exp(-(X - Y)**2 / reg) + t = np.linspace(0, 1, A.shape[2]) + [Y, X] = np.meshgrid(t, t) + xi2 = np.exp(-(X - Y)**2 / reg) + def K(x): - return np.dot(np.dot(xi1, x), xi1) + return np.dot(np.dot(xi1, x), xi2) while (err > stopThr and cpt < numItermax): diff --git a/test/test_bregman.py b/test/test_bregman.py index f54ba9f..ec4388d 100644 --- a/test/test_bregman.py +++ b/test/test_bregman.py @@ -351,3 +351,10 @@ def test_screenkhorn(): # check marginals np.testing.assert_allclose(G_sink.sum(0), G_screen.sum(0), atol=1e-02) np.testing.assert_allclose(G_sink.sum(1), G_screen.sum(1), atol=1e-02) + + +def test_convolutional_barycenter_non_square(): + # test for image with height not equal width + A = np.ones((2, 2, 3)) / (2 * 3) + b = ot.bregman.convolutional_barycenter2d(A, 1e-03) + np.testing.assert_allclose(np.ones((2, 3)) / (2 * 3), b, atol=1e-02) |