File size: 1,661 Bytes
c30b460 | 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 43 44 45 | from __future__ import annotations
import numpy as np
from .spectrum import sample_covariance, symmetric_matrix_sqrt
def sample_covariance_sqrt_shrinkage(
eigenvalues: np.ndarray, n_samples: int, trials: int, seed: int
) -> np.ndarray:
"""Estimate E[u_i^T sample_cov^{1/2} u_i] in the population eigenbasis."""
eig = np.asarray(eigenvalues, dtype=np.float64)
rng = np.random.default_rng(seed)
accum = np.zeros_like(eig)
for _ in range(trials):
cov = sample_covariance(eig, n_samples, rng)
sqrt_cov = symmetric_matrix_sqrt(cov)
accum += np.diag(sqrt_cov)
return accum / float(trials)
def _cross_split_mse(eigenvalues: np.ndarray, n_samples: int, trials: int, rng: np.random.Generator) -> float:
eig = np.asarray(eigenvalues, dtype=np.float64)
mses = []
for _ in range(trials):
seed_vector = rng.standard_normal(eig.size)
left = symmetric_matrix_sqrt(sample_covariance(eig, n_samples, rng)) @ seed_vector
right = symmetric_matrix_sqrt(sample_covariance(eig, n_samples, rng)) @ seed_vector
mses.append(float(np.mean((left - right) ** 2)))
return float(np.mean(mses))
def sampling_map_cross_split_variance(
eigenvalues: np.ndarray, n_samples: int, trials: int, seed: int
) -> dict[str, float]:
rng = np.random.default_rng(seed)
base = _cross_split_mse(eigenvalues, n_samples, trials, rng)
larger = _cross_split_mse(eigenvalues, n_samples * 4, trials, rng)
return {
"cross_split_mse": base,
"larger_n_cross_split_mse": larger,
"n_samples": float(n_samples),
"larger_n_samples": float(n_samples * 4),
}
|