File size: 1,624 Bytes
41dd7e0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 | """Distributional metrics in state space.
``energy_distance`` is the validation metric of the paper's model selection (a V-statistic: ``md(X, X)`` and
``md(Y, Y)`` average over all pairs including self-pairs; ``2 md(X, Y) - md(X, X) - md(Y, Y)``).
"""
from __future__ import annotations
import numpy as np
def _mean_pairwise_dist(A: np.ndarray, B: np.ndarray) -> float:
return np.mean(np.linalg.norm(A[:, None] - B[None, :], axis=-1))
def energy_distance(X: np.ndarray, Y: np.ndarray) -> float:
md = _mean_pairwise_dist
return float(2 * md(X, Y) - md(X, X) - md(Y, Y))
def energy_distance_cached_self(X: np.ndarray, Y: np.ndarray, md_yy: float) -> float:
"""``energy_distance`` with the truth-self term ``md(Y, Y)`` supplied (it does not depend on X)."""
md = _mean_pairwise_dist
return float(2 * md(X, Y) - md(X, X) - float(md_yy))
def energy_distance_u(X: np.ndarray, Y: np.ndarray) -> float:
"""The unbiased U-statistic: the self terms exclude the diagonal, ``ED_U = ED_V - md(X,X)/(M-1) - md(Y,Y)/(N-1)``."""
md = _mean_pairwise_dist
mxx, myy = md(X, X), md(Y, Y)
M, N = X.shape[0], Y.shape[0]
return float(2 * md(X, Y) - mxx - myy - (mxx / (M - 1) if M > 1 else 0.0) - (myy / (N - 1) if N > 1 else 0.0))
def effective_rank(X: np.ndarray) -> float:
"""Effective rank of the sample covariance, tr(S)^2 / ||S||_F^2; NaN for fewer than 2 samples."""
X = np.asarray(X, dtype=np.float64)
if X.shape[0] < 2:
return float("nan")
S = np.cov(X.T)
fro2 = float((S * S).sum())
return float(S.trace() ** 2 / fro2) if fro2 > 0 else float("nan")
|