change API bw

Author: clbonet

ot/backend.py

@@ -1364,7 +1364,6 @@ def solve(self, a, b):
 
     def trace(self, a):
         return np.einsum("...ii", a)
-        # return np.trace(a)
 
     def inv(self, a):
         return scipy.linalg.inv(a)
@@ -1777,8 +1776,7 @@ def solve(self, a, b):
         return jnp.linalg.solve(a, b)
 
     def trace(self, a):
-        return jnp.einsum("...ii", a)
-        # return jnp.trace(a)
+        return jnp.diagonal(a, axis1=-2, axis2=-1).sum(-1)
 
     def inv(self, a):
         return jnp.linalg.inv(a)
@@ -2311,8 +2309,7 @@ def solve(self, a, b):
         return torch.linalg.solve(a, b)
 
     def trace(self, a):
-        return torch.einsum("...ii", a)
-        # return torch.trace(a)
+        return torch.diagonal(a, dim1=-2, dim2=-1).sum(-1)
 
     def inv(self, a):
         return torch.linalg.inv(a)
@@ -2726,8 +2723,7 @@ def solve(self, a, b):
         return cp.linalg.solve(a, b)
 
     def trace(self, a):
-        return cp.einsum("..ii", a)
-        # return cp.trace(a)
+        return cp.trace(a, axis1=-2, axis2=-1)
 
     def inv(self, a):
         return cp.linalg.inv(a)
@@ -3163,8 +3159,7 @@ def solve(self, a, b):
         return tf.linalg.solve(a, b)
 
     def trace(self, a):
-        return tf.einsum("...ii", a)
-        # return tf.linalg.trace(a)
+        return tf.linalg.trace(a)
 
     def inv(self, a):
         return tf.linalg.inv(a)

ot/gaussian.py

@@ -200,7 +200,7 @@ def empirical_bures_wasserstein_mapping(
         return A, b
 
 
-def bures_distance(Cs, Ct, log=False, nx=None):
+def bures_distance(Cs, Ct, paired=False, log=False, nx=None):
     r"""Return Bures distance.
 
     The function computes the Bures distance between :math:`\mu_s=\mathcal{N}(0,\Sigma_s)` and :math:`\mu_t=\mathcal{N}(0,\Sigma_t)`,
@@ -215,6 +215,8 @@ def bures_distance(Cs, Ct, log=False, nx=None):
         covariance of the source distribution
     Ct : array-like (d,d) or (m,d,d)
         covariance of the target distribution
+    paired: bool, optional
+        if True and n==m, return the paired distances and crossed distance otherwise
     log : bool, optional
         record log if True
     nx : module, optional
@@ -223,7 +225,7 @@ def bures_distance(Cs, Ct, log=False, nx=None):
 
     Returns
     -------
-    W : float if Cs and Cd of shape (d,d), array-like (n,) if Cs of shape (n,d,d), Ct  of shape (d,d), array-like (m,) if Cs of shape (d,d) and Ct of shape (m,d,d), array-like (n,m) if Cs of shape (n,d,d) and mt of shape (m,d,d)
+    W : float if Cs and Cd of shape (d,d), array-like (n,m) if Cs of shape (n,d,d) and Ct of shape (m,d,d), array-like (n,) if Cs and Ct of shape (n, d, d) and paired is True
         Bures Wasserstein distance
     log : dict
         log dictionary return only if log==True in parameters
@@ -247,18 +249,18 @@ def bures_distance(Cs, Ct, log=False, nx=None):
     if len(Cs.shape) == 2 and len(Ct.shape) == 2:
         # Return float
         bw2 = nx.trace(Cs + Ct - 2 * nx.sqrtm(dots(Cs12, Ct, Cs12)))
-    elif len(Cs.shape) == 2:
-        # Return shape (m,)
-        M = nx.einsum("ij, mjk, kl -> mil", Cs12, Ct, Cs12)
-        bw2 = nx.trace(Cs[None] + Ct - 2 * nx.sqrtm(M))
-    elif len(Ct.shape) == 2:
-        # Return shape (n,)
-        M = nx.einsum("nij, jk, nkl -> nil", Cs12, Ct, Cs12)
-        bw2 = nx.trace(Cs + Ct[None] - 2 * nx.sqrtm(M))
     else:
-        # Return shape (n,m)
-        M = nx.einsum("nij, mjk, nkl -> nmil", Cs12, Ct, Cs12)
-        bw2 = nx.trace(Cs[:, None] + Ct[None] - 2 * nx.sqrtm(M))
+        assert (
+            len(Cs.shape) == 3 and len(Ct.shape) == 3
+        ), "Both Cs and Ct should be batched"
+        if paired and len(Cs) == len(Ct):
+            # Return shape (n,)
+            M = nx.einsum("nij, njk, nkl -> nil", Cs12, Ct, Cs12)
+            bw2 = nx.trace(Cs + Ct - 2 * nx.sqrtm(M))
+        else:
+            # Return shape (n,m)
+            M = nx.einsum("nij, mjk, nkl -> nmil", Cs12, Ct, Cs12)
+            bw2 = nx.trace(Cs[:, None] + Ct[None] - 2 * nx.sqrtm(M))
 
     W = nx.sqrt(nx.maximum(bw2, 0))
 
@@ -270,7 +272,7 @@ def bures_distance(Cs, Ct, log=False, nx=None):
         return W
 
 
-def bures_wasserstein_distance(ms, mt, Cs, Ct, log=False):
+def bures_wasserstein_distance(ms, mt, Cs, Ct, paired=False, log=False):
     r"""Return Bures Wasserstein distance between samples.
 
     The function computes the Bures-Wasserstein distance between :math:`\mu_s=\mathcal{N}(m_s,\Sigma_s)` and :math:`\mu_t=\mathcal{N}(m_t,\Sigma_t)`,
@@ -294,12 +296,14 @@ def bures_wasserstein_distance(ms, mt, Cs, Ct, log=False):
         covariance of the source distribution
     Ct : array-like (d,d) or (m,d,d)
         covariance of the target distribution
+    paired: bool, optional
+        if True and n==m, return the paired distances and crossed distance otherwise
     log : bool, optional
         record log if True
 
     Returns
     -------
-    W : float if ms and md of shape (d,), array-like (n,) if ms of shape (n,d), mt  of shape (d,), array-like (m,) if ms of shape (d,) and mt of shape (m,d), array-like (n,m) if ms of shape (n,d) and mt of shape (m,d)
+    W : float if ms and md of shape (d,), array-like (n,m) if ms of shape (n,d) and mt of shape (m,d), array-like (n,) if ms and mt of shape (n,d) and paired is True
         Bures Wasserstein distance
     log : dict
         log dictionary return only if log==True in parameters
@@ -328,23 +332,24 @@ def bures_wasserstein_distance(ms, mt, Cs, Ct, log=False):
     ), "All Gaussian must have the same dimension"
 
     if log:
-        bw, log_dict = bures_distance(Cs, Ct, log=log, nx=nx)
+        bw, log_dict = bures_distance(Cs, Ct, paired=paired, log=log, nx=nx)
         Cs12 = log_dict["Cs12"]
     else:
-        bw = bures_distance(Cs, Ct, nx=nx)
+        bw = bures_distance(Cs, Ct, paired=paired, nx=nx)
 
     if len(ms.shape) == 1 and len(mt.shape) == 1:
         # Return float
         squared_dist_m = nx.norm(ms - mt) ** 2
-    elif len(ms.shape) == 1:
-        # Return shape (m,)
-        squared_dist_m = nx.norm(ms[None] - mt, axis=-1) ** 2
-    elif len(mt.shape) == 1:
-        # Return shape (n,)
-        squared_dist_m = nx.norm(ms - mt[None], axis=-1) ** 2
     else:
-        # Return shape (n,m)
-        squared_dist_m = nx.norm(ms[:, None] - mt[None], axis=-1) ** 2
+        assert (
+            len(ms.shape) == 2 and len(mt.shape) == 2
+        ), "Both ms and mt should be batched"
+        if paired and len(ms.shape) == len(mt.shape):
+            # Return shape (n,)
+            squared_dist_m = nx.norm(ms - mt, axis=-1) ** 2
+        else:
+            # Return shape (n,m)
+            squared_dist_m = nx.norm(ms[:, None] - mt[None], axis=-1) ** 2
 
     W = nx.sqrt(nx.maximum(squared_dist_m + bw**2, 0))
 
@@ -882,12 +887,14 @@ def empirical_bures_wasserstein_barycenter(
         nx.dot((X[i] * w[i]).T, X[i]) / nx.sum(w[i]) + reg * nx.eye(d[i], type_as=X[i])
         for i in range(k)
     ]
-    m = nx.stack(m, axis=0)
+    m = nx.stack(m, axis=0)[:, 0]
     C = nx.stack(C, axis=0)
+
     if log:
         mb, Cb, log = bures_wasserstein_barycenter(
             m, C, weights=weights, num_iter=num_iter, eps=eps, log=log
         )
+
         return mb, Cb, log
     else:
         mb, Cb = bures_wasserstein_barycenter(

test/test_gaussian.py

@@ -124,6 +124,7 @@ def test_bures_wasserstein_distance_batch(nx):
 
     Wb = ot.gaussian.bures_wasserstein_distance(m[0, 0], m[1, 0], C[0], C[1], log=False)
 
+    # Test cross vs 1
     Wb2 = ot.gaussian.bures_wasserstein_distance(
         m[0, 0][None], m[1, 0][None], C[0][None], C[1][None]
     )
@@ -135,25 +136,34 @@ def test_bures_wasserstein_distance_batch(nx):
     np.testing.assert_allclose(0, nx.to_numpy(Wb2[1, 0]), atol=1e-5)
     np.testing.assert_equal(Wb2.shape, (2, 1))
 
-    Wb2 = ot.gaussian.bures_wasserstein_distance(
-        m[0, 0][None], m[1, 0], C[0][None], C[1]
-    )
-    np.testing.assert_allclose(nx.to_numpy(Wb), nx.to_numpy(Wb2[0]), atol=1e-5)
-    np.testing.assert_equal(Wb2.shape, (1,))
-
-    Wb2 = ot.gaussian.bures_wasserstein_distance(
-        m[0, 0], m[1, 0][None], C[0], C[1][None]
-    )
-    np.testing.assert_allclose(nx.to_numpy(Wb), nx.to_numpy(Wb2[0]), atol=1e-5)
-    np.testing.assert_equal(Wb2.shape, (1,))
-
     Wb2 = ot.gaussian.bures_wasserstein_distance(m[:, 0], m[:, 0], C, C)
     np.testing.assert_allclose(nx.to_numpy(Wb), nx.to_numpy(Wb2[1, 0]), atol=1e-5)
     np.testing.assert_allclose(nx.to_numpy(Wb), nx.to_numpy(Wb2[0, 1]), atol=1e-5)
     np.testing.assert_allclose(0, nx.to_numpy(Wb2[0, 0]), atol=1e-5)
     np.testing.assert_allclose(0, nx.to_numpy(Wb2[1, 1]), atol=1e-5)
     np.testing.assert_equal(Wb2.shape, (2, 2))
 
+    # Test paired
+    Wb3 = ot.gaussian.bures_wasserstein_distance(m[:, 0], m[:, 0], C, C, paired=True)
+    np.testing.assert_allclose(0, nx.to_numpy(Wb3[0]), atol=1e-5)
+    np.testing.assert_allclose(0, nx.to_numpy(Wb3[1]), atol=1e-5)
+
+    m_rev = np.zeros((k, 2))
+    C_rev = np.zeros((k, 2, 2))
+    m_rev[0] = m[1, 0]
+    m_rev[1] = m[0, 0]
+    C_rev[0] = C[1]
+    C_rev[1] = C[0]
+
+    Wb3 = ot.gaussian.bures_wasserstein_distance(m_rev, m[:, 0], C_rev, C, paired=True)
+    np.testing.assert_allclose(nx.to_numpy(Wb2)[0, 1], nx.to_numpy(Wb3)[0], atol=1e-5)
+    np.testing.assert_allclose(nx.to_numpy(Wb2)[0, 1], nx.to_numpy(Wb3)[0], atol=1e-5)
+
+    with pytest.raises(AssertionError):
+        Wb3 = ot.gaussian.bures_wasserstein_distance(m[0, 0], m[:, 0], C[0], C)
+        np.testing.assert_allclose(0, nx.to_numpy(Wb3[0]), atol=1e-5)
+        np.testing.assert_allclose(0, nx.to_numpy(Wb3[1]), atol=1e-5)
+
 
 @pytest.mark.parametrize("bias", [True, False])
 def test_empirical_bures_wasserstein_distance(nx, bias):
@@ -205,7 +215,7 @@ def test_bures_wasserstein_barycenter(nx, method):
     m = np.array(m)
     C = np.array(C)
     X = nx.from_numpy(*X)
-    m = nx.from_numpy(m)
+    m = nx.from_numpy(m)[:, 0]
     C = nx.from_numpy(C)
 
     mblog, Cblog, log = ot.gaussian.bures_wasserstein_barycenter(
@@ -256,7 +266,7 @@ def test_fixedpoint_vs_gradientdescent_bures_wasserstein_barycenter(nx):
     m = np.array(m)
     C = np.array(C)
     X = nx.from_numpy(*X)
-    m = nx.from_numpy(m)
+    m = nx.from_numpy(m)[:, 0]
     C = nx.from_numpy(C)
 
     mb, Cb = ot.gaussian.bures_wasserstein_barycenter(
@@ -297,7 +307,7 @@ def test_stochastic_gd_bures_wasserstein_barycenter(nx, method):
     m = np.array(m)
     C = np.array(C)
     X = nx.from_numpy(*X)
-    m = nx.from_numpy(m)
+    m = nx.from_numpy(m)[:, 0]
     C = nx.from_numpy(C)
 
     mb, Cb = ot.gaussian.bures_wasserstein_barycenter(
@@ -346,7 +356,7 @@ def test_not_implemented_method(nx):
     m = np.array(m)
     C = np.array(C)
     X = nx.from_numpy(*X)
-    m = nx.from_numpy(m)
+    m = nx.from_numpy(m)[:, 0]
     C = nx.from_numpy(C)
 
     not_implemented = "new_method"