129 lines
4.0 KiB
Python
129 lines
4.0 KiB
Python
#!/usr/bin/env python3
|
|
"""Independent NumPy checks for the spectral-space AI-REML equations.
|
|
|
|
This file deliberately does not import the production C++ implementation. It
|
|
serves as a small, readable oracle for gradient and transformation tests.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
|
|
import numpy as np
|
|
|
|
|
|
def restricted_loglik_and_derivatives(
|
|
theta: np.ndarray,
|
|
eigenvalues: np.ndarray,
|
|
x_star: np.ndarray,
|
|
y_star: np.ndarray,
|
|
) -> tuple[float, np.ndarray, np.ndarray]:
|
|
"""Return paper-scale REML log likelihood, score, and AI matrix.
|
|
|
|
Parameters are the standard-deviation coordinates ``(sigma_e, sigma_g)``.
|
|
The function forms P explicitly because it is only used on tiny test data.
|
|
"""
|
|
|
|
sigma_e, sigma_g = (float(theta[0]), float(theta[1]))
|
|
diagonal = sigma_e**2 + sigma_g**2 * eigenvalues
|
|
if np.any(diagonal <= 0.0):
|
|
raise ValueError("H is not positive definite")
|
|
w = 1.0 / diagonal
|
|
h_inv = np.diag(w)
|
|
c = x_star.T @ h_inv @ x_star
|
|
c_inv = np.linalg.inv(c)
|
|
p = h_inv - h_inv @ x_star @ c_inv @ x_star.T @ h_inv
|
|
py = p @ y_star
|
|
|
|
loglik = -0.5 * (
|
|
np.linalg.slogdet(c)[1]
|
|
+ np.log(diagonal).sum()
|
|
+ float(y_star @ py)
|
|
)
|
|
|
|
score_e = -sigma_e * (np.trace(p) - float(py @ py))
|
|
score_g = -sigma_g * (
|
|
np.trace(p @ np.diag(eigenvalues))
|
|
- float(py @ (eigenvalues * py))
|
|
)
|
|
|
|
p_py = p @ py
|
|
lambda_py = eigenvalues * py
|
|
p_lambda_py = p @ lambda_py
|
|
ai_ee = 2.0 * sigma_e**2 * float(py @ p_py)
|
|
ai_eg = 2.0 * sigma_e * sigma_g * float(py @ p_lambda_py)
|
|
ai_gg = 2.0 * sigma_g**2 * float(lambda_py @ p_lambda_py)
|
|
ai = np.asarray([[ai_ee, ai_eg], [ai_eg, ai_gg]], dtype=np.float64)
|
|
return loglik, np.asarray([score_e, score_g]), ai
|
|
|
|
|
|
def finite_difference_gradient(
|
|
theta: np.ndarray,
|
|
eigenvalues: np.ndarray,
|
|
x_star: np.ndarray,
|
|
y_star: np.ndarray,
|
|
step: float = 1e-6,
|
|
) -> np.ndarray:
|
|
result = np.empty(2, dtype=np.float64)
|
|
for index in range(2):
|
|
delta = np.zeros(2, dtype=np.float64)
|
|
delta[index] = step
|
|
high = restricted_loglik_and_derivatives(
|
|
theta + delta, eigenvalues, x_star, y_star
|
|
)[0]
|
|
low = restricted_loglik_and_derivatives(
|
|
theta - delta, eigenvalues, x_star, y_star
|
|
)[0]
|
|
result[index] = (high - low) / (2.0 * step)
|
|
return result
|
|
|
|
|
|
def test_score_matches_finite_difference() -> None:
|
|
rng = np.random.default_rng(70123)
|
|
n = 17
|
|
p = 4
|
|
eigenvalues = np.linspace(0.05, 2.1, n)
|
|
x_star = np.column_stack((np.ones(n), rng.normal(size=(n, p - 1))))
|
|
y_star = rng.normal(size=n)
|
|
theta = np.asarray([0.83, 0.57])
|
|
_, score, ai = restricted_loglik_and_derivatives(
|
|
theta, eigenvalues, x_star, y_star
|
|
)
|
|
numerical = finite_difference_gradient(theta, eigenvalues, x_star, y_star)
|
|
np.testing.assert_allclose(score, numerical, rtol=2e-6, atol=2e-6)
|
|
np.testing.assert_allclose(ai, ai.T, rtol=0.0, atol=1e-12)
|
|
assert np.linalg.eigvalsh(ai).min() > 0.0
|
|
|
|
|
|
def test_orthogonal_transformation_preserves_reml() -> None:
|
|
rng = np.random.default_rng(70124)
|
|
n = 15
|
|
q, _ = np.linalg.qr(rng.normal(size=(n, n)))
|
|
eigenvalues = np.linspace(0.1, 1.8, n)
|
|
grm = q @ np.diag(eigenvalues) @ q.T
|
|
x = np.column_stack((np.ones(n), rng.normal(size=(n, 3))))
|
|
y = rng.normal(size=n)
|
|
theta = np.asarray([0.76, 0.62])
|
|
sigma_e, sigma_g = theta
|
|
|
|
h = sigma_g**2 * grm + sigma_e**2 * np.eye(n)
|
|
h_inv = np.linalg.inv(h)
|
|
c = x.T @ h_inv @ x
|
|
p = h_inv - h_inv @ x @ np.linalg.inv(c) @ x.T @ h_inv
|
|
original = -0.5 * (
|
|
np.linalg.slogdet(c)[1]
|
|
+ np.linalg.slogdet(h)[1]
|
|
+ float(y @ p @ y)
|
|
)
|
|
|
|
transformed = restricted_loglik_and_derivatives(
|
|
theta, eigenvalues, q.T @ x, q.T @ y
|
|
)[0]
|
|
assert math.isclose(original, transformed, rel_tol=2e-12, abs_tol=2e-12)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_score_matches_finite_difference()
|
|
test_orthogonal_transformation_preserves_reml()
|
|
print("PASS: spectral AI-REML reference checks")
|