wasserstein_1d

Contents

wasserstein_1d#

class braintools.fit.losses.wasserstein_1d(p, q, x)#

Compute the Wasserstein-1 distance between two densities on a shared uniform grid.

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

where \(F_p\) and \(F_q\) are the cumulative sums of p and q normalised to end at 1, so unnormalised histograms are accepted.

Parameters:

  • p (Any) – 1-D density or histogram on the grid x.

  • q (Any) – Density on the same grid (same length, compatible units).

  • x (Any) – Uniform evaluation grid of the same length, at least 2 points. Only its spacing x[1] - x[0] is used.

Returns:

Scalar distance, in the units of x (a Quantity if x is one).

Return type:

Any

Raises:

ValueError – If p is not 1-D, the shapes differ, or x has fewer than 2 points.

See also

ks_distance

The non-smooth counterpart.

fcd_wasserstein

Wasserstein distance between two FCD distributions.

Notes

Ported from brainmass.objectives.wasserstein_1d with identical numerics. It converges to scipy.stats.wasserstein_distance() as the grid is refined. Works under jax.jit and jax.vmap and is differentiable in p and q.

Examples

>>> import jax.numpy as jnp
>>> from braintools.fit import losses
>>> x = jnp.arange(5.)
>>> p = jnp.array([1., 0., 0., 0., 0.])
>>> q = jnp.array([0., 0., 1., 0., 0.])
>>> float(losses.wasserstein_1d(p, q, x))
2.0
Contents