fcd_ks

Contents

fcd_ks#

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

Compute the Kolmogorov-Smirnov distance between two FCD value distributions.

The functional-connectivity-dynamics (FCD) matrix of each input is computed with functional_connectivity_dynamics(); the values above the n_diag-th diagonal are smoothed into a density on midpoints by a Gaussian kernel density estimate (jax.scipy.stats.gaussian_kde), and the two densities are compared with ks_distance(). Comparing FCD distributions (rather than the matrices) is the standard FCD fitting target [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) – 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 KS distance in [0, 1].

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_wasserstein

The smooth counterpart, preferred for gradient fitting.

ks_distance

Notes

Ported from brainmass.objectives.fcd_ks(...) (a factory there) with identical numerics. The KS statistic is non-smooth; use it for evaluation and fcd_wasserstein() as a gradient-fitting loss. 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.

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 (previously 192 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_ks(x, x, window_size=30, step_size=5))
0.0
Contents