boltzkit.evaluation.sample_based.wasserstein

Implementation of different Wasserstein metrics.

Functions

get_euclidean_wasserstein_1_2(X1, X2[, ...])

Compute empirical Wasserstein-1 (W1) and Wasserstein-2 (W2) distances between two point clouds using exact optimal transport.

get_torus_wasserstein(angles0, angles1)

Compute the 2-Wasserstein distance between two sets of angular samples on a torus (periodic domain).
boltzkit.evaluation.sample_based.wasserstein.get_torus_wasserstein(angles0: ndarray, angles1: ndarray) → float[source]

Compute the 2-Wasserstein distance between two sets of angular samples on a torus (periodic domain).

Parameters:

  • angles0 (np.ndarray) – First set of angular samples with shape (n_samples_0, d), where d is the number of angular dimensions.

  • angles1 (np.ndarray) – Second set of angular samples with shape (n_samples_1, d).

Returns:

The 2-Wasserstein distance between the empirical distributions defined by angles0 and angles1.

Return type:

float

Notes

  • Distances are computed using the wrapped (circular) metric: min(|Δθ|, 2π - |Δθ|).

  • The squared Euclidean distance on the torus is used as the ground cost.

  • Uniform weights are assumed for both empirical distributions.

  • The optimal transport problem is solved using POT’s emd2.

boltzkit.evaluation.sample_based.wasserstein.get_euclidean_wasserstein_1_2(X1: ndarray, X2: ndarray, weights1: ndarray | None = None, num_iter_max: int = 1000000000, include_w1: bool = True, include_w2: bool = True)[source]

Compute empirical Wasserstein-1 (W1) and Wasserstein-2 (W2) distances between two point clouds using exact optimal transport.

Parameters:

  • X1 (np.ndarray of shape (N, d)) – First point cloud with N samples in d dimensions.

  • X2 (np.ndarray of shape (M, d)) – Second point cloud with M samples in d dimensions.

  • weights1 (np.ndarray of shape (N,), optional) – Non-negative weights for samples in X1. If None, uniform weights are used. The second point cloud always uses uniform weights.

  • num_iter_max (int, default=1e6) – Maximum number of iterations for the optimal transport solver.

  • include_w1 (bool, default=True) – If True, compute and return Wasserstein-1 distance.

  • include_w2 (bool, default=True) – If True, compute and return Wasserstein-2 distance.

Returns:

  • W1 (float or None) – Wasserstein-1 distance using Euclidean ground cost. Returns None if include_w1=False.

  • W2 (float or None) – Wasserstein-2 distance using squared Euclidean ground cost. Returns None if include_w2=False.

Notes

  • Exact optimal transport is computed using ot.emd2.

  • W1 is computed as:

    W1 = emd2(a, b, ||x - y||)

  • W2 is computed as:

    W2 = sqrt(emd2(a, b, ||x - y||^2))

  • Computational complexity is O(NM) in memory and typically super-cubic in time for exact OT.