Skip to content

[MRG] Use the weights w in ot.dist for all metrics - #884

Open
raashish1601 wants to merge 4 commits into
PythonOT:masterfrom
raashish1601:fix/dist-weights
Open

raashish1601 wants to merge 4 commits into
PythonOT:masterfrom
raashish1601:fix/dist-weights

Conversation

@raashish1601

Copy link
Copy Markdown
Contributor

Types of changes

  • Bug fix

Motivation and context / Related issue

ot.dist silently ignored w for every metric except cityblock and minkowski, so it returned the unweighted distance:

import numpy as np, ot
from scipy.spatial.distance import cdist
x, y, w = np.array([[0., 0.]]), np.array([[1., 1.]]), np.array([1., 2.])
ot.dist(x, y, metric="sqeuclidean", w=w)   # [[2.]]
cdist(x, y, metric="sqeuclidean", w=w)     # [[3.]]

This affected the backend implementations of sqeuclidean, euclidean, cosine and correlation, and the numpy fallback to cdist, which only passed w for minkowski/wminkowski. With backend="scipy" the weights were used, so the result depended on the backend. test_dist already lists these metrics as "those that support weights" but did not check the values. #859 fixed the same problem for cityblock.

Changes:

  • sqeuclidean/euclidean: scale both point sets by sqrt(w) before euclidean_distances.
  • cosine: same scaling. correlation: center with the weighted mean, then the same scaling. This matches the scipy definitions.
  • numpy fallback: pass w to cdist whenever it is given (scipy still raises for metrics that do not accept weights).

Unweighted results are unchanged.

How has this been tested (if it applies)

Added test_dist_weighted_vs_cdist (numpy, compared with scipy.spatial.distance.cdist for sqeuclidean, euclidean, cosine, correlation, braycurtis and canberra) and test_dist_weighted_backends (all backends, for the four backend metrics). All new cases fail on master. pytest test/test_utils.py passes (numpy and torch backends installed locally), and ruff 0.5.2 lint and format are clean.

PR checklist

  • I have read the CONTRIBUTING document.
  • The documentation is up-to-date with the changes I made (check build artifacts).
  • All tests passed, and additional code has been covered with new tests.
  • I have added the PR and Issue fix to the RELEASES.md file.

@codecov

codecov Bot commented Oct 12, 2026 •

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 97.06%. Comparing base (e1c7481) to head (04b0b85).
⚠️ Report is 1 commits behind head on master.

Additional details and impacted files
@@           Coverage Diff           @@
##           master     #884   +/-   ##
=======================================
  Coverage   97.06%   97.06%           
=======================================
  Files         129      129           
  Lines       26509    26531   +22     
=======================================
+ Hits        25730    25753   +23     
+ Misses        779      778    -1     
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@cedricvincentcuaz cedricvincentcuaz left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Dear @raashish1601,
Thank you for your PR. Could you take into account my comments before i merge in the master branch ?
Best,
Cédric

Comment thread ot/utils.py
# sum_k w_k (x1_k - x2_k)^2 is the squared distance between the
# samples scaled by sqrt(w)
x1 = x1 * nx.sqrt(w)[None, :]
x2 = x2 * nx.sqrt(w)[None, :]

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As euclidean_distances includes a test if X is Y, which tests whether both variables point to the same Python object - it is problematic here to define two new variables x1 and x2.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You should instead preserve the function behaviour by doing:

sqrt_w = nx.sqrt(w)
if x1 is x2:
x1 = x1 * sqrt_w[None, :]
x2 = x1
else:
x1 = x1 * sqrt_w[None, :]
x2 = x2 * sqrt_w[None, :]

Comment thread ot/utils.py
if not squared:
c = nx.sqrt(c)

if X is Y:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This test should actually come sooner to avoid redundant computations as in this case b2=a2

Comment thread ot/utils.py

"""
if nx is None:
nx = get_backend(x1, x2)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should add w in this backend test. get_backend can handle w=None so no need for any case differentiation here

Comment thread ot/utils.py
x1 = x1 - nx.mean(x1, axis=1)[:, None]
x2 = x2 - nx.mean(x2, axis=1)[:, None]
else:
x1 = x1 - (nx.dot(x1, w) / nx.sum(w))[:, None]

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nx.sum(w) should be computed only once and tested for zero value in which case we can raise a ValueError

Comment thread ot/utils.py
x1 = x1 - (nx.dot(x1, w) / nx.sum(w))[:, None]
x2 = x2 - (nx.dot(x2, w) / nx.sum(w))[:, None]
if w is not None:
x1 = x1 * nx.sqrt(w)[None, :]

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

compute sqrt only once

Comment thread test/test_utils.py
@pytest.mark.parametrize(
"metric", ["sqeuclidean", "euclidean", "cosine", "correlation"]
)
def test_dist_weighted_backends(nx, metric):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

these 2 tests can be merged, with different cases depending on the backend if nx.__name__ == "jax": fixing a lower tolerance than the other backends.

Comment thread RELEASES.md

#### Closed issues

- Fix `ot.dist` ignoring the weights `w` for the `sqeuclidean`, `euclidean`, `cosine` and `correlation` metrics and for the other metrics computed with `scipy.spatial.distance.cdist` (PR #884)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

move to the end of the itemize please to try to follow PR order

@cedricvincentcuaz cedricvincentcuaz self-assigned this Oct 13, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants