| """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") |
|
|
