fcd_wasserstein

Contents

fcd_wasserstein#

class braintools.fit.losses.fcd_wasserstein(prediction, target, *, window_size=30, step_size=5, midpoints=None, bw_method=None, n_diag=1)#

Compute the Wasserstein-1 distance between two FCD value distributions.

Like fcd_ks(), but the two kernel density estimates of the FCD values are compared with wasserstein_1d() on midpoints, which is smooth in the inputs and so the recommended FCD loss for gradient fitting [1].

Parameters:

  • prediction (Any) – Model output of shape (n_time, n_regions), n_regions >= 2.

  • target (Any) – Data of shape (n_time', n_regions'). FCD is scale-invariant, so units are not compared.

  • window_size (int) – Sliding-window length in samples (static).

  • step_size (int) – Sliding-window stride in samples (static).

  • midpoints (Any | None) – Uniform grid of FCD values on which the densities are evaluated. Defaults to 100 points on [-0.99, 0.99].

  • bw_method (Any | None) – KDE bandwidth, forwarded to jax.scipy.stats.gaussian_kde (default: Scott’s rule).

  • n_diag (int) – Diagonal offset of the upper triangle; 1 excludes the unit diagonal.

Returns:

Scalar distance (in FCD-correlation units).

Return type:

Array

Raises:

ValueError – If an input is not 2-D with at least 2 regions, or its FCD has fewer than 2 values above the diagonal offset.

See also

fcd_ks, wasserstein_1d

Notes

Ported from brainmass.objectives.fcd_wasserstein(...) (a factory there) with identical numerics. If all FCD values of an input are equal, the KDE is degenerate and the result is nan. Works under jax.jit and jax.vmap with static window options, and is differentiable.

Memory. The FCD matrix of W windows has W**2 entries and the KDE sums W(W-1)/2 kernel terms for each of the M midpoints. The W(W-1)/2 x M kernel matrix is never formed whole: midpoints are processed in blocks of at most about 2**20 terms (4 MiB in float32), and each block is recomputed on the backward pass rather than stored, so temporaries grow as O(W**2) rather than O(W**2 M). For example, 20 regions, T = 5000, window_size=30, step_size=5 (995 windows) need about 24 MiB of temporaries forward and 36 MiB for the gradient (previously 192 MiB and 588 MiB). Under vmap (e.g. over the starts of a fit) the temporaries scale with the batch size.

References

Examples

>>> import numpy as np
>>> import jax.numpy as jnp
>>> from braintools.fit import losses
>>> x = jnp.asarray(np.random.default_rng(0).standard_normal((200, 6)))
>>> float(losses.fcd_wasserstein(x, x, window_size=30, step_size=5))
0.0
Contents