incerto.shift.wasserstein_distance

incerto.shift.wasserstein_distance#

incerto.shift.wasserstein_distance(x, y, p=2.0, max_iter=100, epsilon=None)[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

  • epsilon (Optional[float]) – Entropy regularization parameter for Sinkhorn. If None (default), automatically scaled based on the cost matrix median to avoid numerical issues. Smaller values give more accurate results but may cause underflow.

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