Module wasserstein

Wasserstein (optimal transport) distance between distributions.

Measures the “cost” of transforming one probability distribution into another, respecting the geometry of the underlying space. Unlike L2 distance between histograms, Wasserstein accounts for which bins are close vs far.

Why Wasserstein for CVX

Region distributions at two timestamps are histograms over K regions. L2 treats all regions as equally distant. Wasserstein uses the actual distances between region centroids — shifting mass between neighboring regions costs less than between distant ones.

Implementations

• sliced_wasserstein: Fast approximation via random 1D projections. O(K × n_proj × K log K).
• wasserstein_1d: Exact W₁ on 1D distributions. O(K log K).
• [emd_1d]: Earth Mover’s Distance on sorted 1D values.

References

• Villani, C. (2008). Optimal Transport: Old and New. Springer.
• Bonneel, N. et al. (2015). Sliced and Radon Wasserstein barycenters. JMIV.

Functions

sliced_wasserstein

Sliced Wasserstein distance between two distributions in K dimensions.

wasserstein_1d

Exact Wasserstein-1 (Earth Mover’s Distance) between two 1D distributions.

wasserstein_drift

Compute Wasserstein drift between two region distributions.