brainmass.objectives.wasserstein_1d

Contents

brainmass.objectives.wasserstein_1d#

brainmass.objectives.wasserstein_1d(p, q, x)[source]#

Wasserstein-1 distance between two 1-D densities on a shared grid.

\[W_1(p, q) = \int \left| F_p(x) - F_q(x) \right| \, dx,\]

discretised on the uniform grid x. Both inputs are normalised internally (CDFs run 0..1), so p and q may be densities or unnormalised histograms; the returned value lives in the same units as x. Unlike ks_distance() this is smooth and differentiable in both inputs, making it the preferred distributional loss for gradient-based fitting. Matches scipy.stats.wasserstein_distance() as the grid is refined.

Parameters:

  • p (array) – Densities or histograms on the shared grid x.

  • q (array) – Densities or histograms on the shared grid x.

  • x (array) – Uniform evaluation grid (same length as p / q).

Returns:

Scalar Wasserstein-1 distance.

Return type:

jax.Array

See also

ks_distance

the non-smooth KS counterpart.

Contents