incerto.shift.wasserstein_distance

incerto.shift.wasserstein_distance#

incerto.shift.wasserstein_distance(x, y, p=2.0, max_iter=100)[source]#

Wasserstein distance (Earth Mover’s Distance) between two empirical distributions.

Computes the p-Wasserstein distance between two samples using Sinkhorn algorithm for optimal transport. For p=1, this is the classic Earth Mover’s Distance.

Parameters:

  • x (Tensor) – Source samples of shape (n, d)

  • y (Tensor) – Target samples of shape (m, d)

  • p (float) – Order of the Wasserstein distance (default: 2.0 for W2 distance)

  • max_iter (int) – Maximum iterations for Sinkhorn algorithm

Return type:

float

Returns:

Wasserstein distance between the two distributions

Reference:

Cuturi, “Sinkhorn Distances: Lightspeed Computation of Optimal Transport” (NeurIPS 2013)

Example

>>> source_features = model(source_data)
>>> target_features = model(target_data)
>>> distance = wasserstein_distance(source_features, target_features)
Contents