ks_distance

Contents

ks_distance#

class braintools.fit.losses.ks_distance(p, q)#

Compute the Kolmogorov-Smirnov statistic between two densities on a shared grid.

Each cumulative sum is normalised to a proper CDF before the supremum is taken,

\[D = \max_k \left| F_p(x_k) - F_q(x_k) \right|,\]

so the result lies in [0, 1] independently of the bin width, and is comparable with scipy.stats.ks_2samp().

Parameters:

  • p (Any) – 1-D density or (unnormalised) histogram on an ordered grid.

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

Returns:

Scalar KS statistic in [0, 1].

Return type:

Array

Raises:

ValueError – If p is not 1-D or the shapes differ.

See also

wasserstein_1d

The smooth counterpart, better suited to gradient fitting.

fcd_ks

KS distance between two FCD distributions.

Notes

Ported from brainmass.objectives.ks_distance with identical numerics. The max makes the statistic non-smooth: its gradient is non-zero only at the arg-max. Prefer wasserstein_1d() as a fitting loss and use KS for evaluation. Works under jax.jit and jax.vmap.

Examples

>>> import jax.numpy as jnp
>>> from braintools.fit import losses
>>> float(losses.ks_distance(jnp.array([1., 0., 0., 0.]), jnp.array([0., 0., 0., 1.])))
1.0
>>> float(losses.ks_distance(jnp.array([1., 1., 0.]), jnp.array([2., 2., 0.])))
0.0
Contents