diff --git a/RELEASES.md b/RELEASES.md index d2b3acdc1..a83c8ba2d 100644 --- a/RELEASES.md +++ b/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) diff --git a/ot/utils.py b/ot/utils.py index bb7dc1884..0c1cd735d 100644 --- a/ot/utils.py +++ b/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: diff --git a/test/test_utils.py b/test/test_utils.py index 23f1ca96a..9c93d520a 100644 --- a/test/test_utils.py +++ b/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]])