[MRG] Use the weights w in ot.dist for the cityblock metric

RELEASES.md

@@ -12,6 +12,7 @@
 
 #### Closed issues
 
+- 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)
 - Fix device placement in `ot.batch.bregman_projection_batch` so `ot.solve_batch(..., method="sinkhorn")` no longer crashes on GPU when the torch default device is CPU (PR #851)
 - Preserve input dtype and device for expected sliced plans, avoid materializing dense distance matrices for sparse plans, and fix weighted sparse-distance ordering (PR #846, Issue #845)

ot/utils.py

@@ -511,13 +511,24 @@ def dist(
     elif metric == "euclidean":
         return euclidean_distances(x1, x2, squared=False, nx=nx)
     elif metric == "cityblock":
-        if use_tensor:
-            return nx.sum(nx.abs(x1[:, None, :] - x2[None, :, :]), axis=2)
+        if w is None:
+            if use_tensor:
+                return nx.sum(nx.abs(x1[:, None, :] - x2[None, :, :]), axis=2)
+            else:
+                M = 0.0
+                for i in range(x1.shape[1]):
+                    M += nx.abs(x1[:, i][:, None] - x2[:, i][None, :])
+                return M
         else:
-            M = 0.0
-            for i in range(x1.shape[1]):
-                M += nx.abs(x1[:, i][:, None] - x2[:, i][None, :])
-            return M
+            if use_tensor:
+                return nx.sum(
+                    w[None, None, :] * nx.abs(x1[:, None, :] - x2[None, :, :]), axis=2
+                )
+            else:
+                M = 0.0
+                for i in range(x1.shape[1]):
+                    M += w[i] * nx.abs(x1[:, i][:, None] - x2[:, i][None, :])
+                return M
     elif metric == "minkowski":
         if w is None:
             if use_tensor:

test/test_utils.py

@@ -307,6 +307,19 @@ def test_dist():
         ot.dist(x, x, metric="fakeone")
 
 
+@pytest.mark.parametrize("use_tensor", [True, False])
+def test_dist_weighted_cityblock(use_tensor):
+    rng = np.random.RandomState(0)
+    x1 = rng.randn(5, 3)
+    x2 = rng.randn(4, 3)
+    w = rng.rand(3)
+
+    expected = scipy.spatial.distance.cdist(x1, x2, metric="cityblock", w=w)
+    D = ot.dist(x1, x2, metric="cityblock", w=w, use_tensor=use_tensor)
+
+    np.testing.assert_allclose(D, expected, atol=1e-12)
+
+
 def test_sparse_ot_dist_uses_pair_weights():
     x1 = np.array([[0.0], [1.0]])
     x2 = np.array([[0.0], [5.0]])