JiaruiLi/SpectraREML
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add reusable SpectraREML batch AI-REML engine

Browse Source
This commit is contained in:
Jiarui Li
2026-08-09 14:13:38 +08:00
parent 777eeb288c
commit da971a2ca7
22 changed files with 4736 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

130
tests/cpp/test_grm_io.cpp Normal file
View File

@@ -0,0 +1,130 @@
#include "spectra_reml/grm.hpp"
#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstdlib>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <stdexcept>
#include <string>
#include <vector>
namespace {
void require(bool condition, const std::string& message) {
if (!condition) {
throw std::runtime_error(message);
}
}
void require_near(double actual, double expected, double tolerance,
const std::string& message) {
if (std::abs(actual - expected) > tolerance) {
throw std::runtime_error(message + ": actual=" +
std::to_string(actual) +
", expected=" + std::to_string(expected));
}
}
class TemporaryDirectory {
public:
TemporaryDirectory() {
const auto seed =
std::chrono::high_resolution_clock::now().time_since_epoch().count();
path_ = std::filesystem::temp_directory_path() /
("spectra_reml_grm_test_" + std::to_string(seed));
std::filesystem::create_directories(path_);
}
~TemporaryDirectory() {
std::error_code ignored;
std::filesystem::remove_all(path_, ignored);
}
const std::filesystem::path& path() const noexcept { return path_; }
private:
std::filesystem::path path_;
};
void write_floats(const std::filesystem::path& path,
const std::vector<float>& values) {
std::ofstream output(path, std::ios::binary | std::ios::trunc);
output.write(reinterpret_cast<const char*>(values.data()),
static_cast<std::streamsize>(values.size() * sizeof(float)));
if (!output) {
throw std::runtime_error("failed to create GRM fixture");
}
}
void test_read_and_decompose() {
TemporaryDirectory temporary;
const auto grm_path = temporary.path() / "fixture.grm.bin";
const auto id_path = temporary.path() / "fixture.grm.id";
// Symmetric positive-definite matrix, packed as GCTA lower triangle.
const std::vector<float> packed = {1.2F, 0.2F, 1.1F,
-0.1F, 0.15F, 0.9F};
write_floats(grm_path, packed);
{
std::ofstream ids(id_path);
ids << "F1 I1\nF2 I2\nF3 I3\n";
}
spectra::reml::validate_grm_id_count(id_path, 3);
const auto matrix = spectra::reml::read_gcta_grm_lower_triangle(grm_path, 3);
require_near(matrix(0, 0), 1.2, 1e-7, "wrong GRM diagonal");
require_near(matrix(1, 0), 0.2, 1e-7, "wrong GRM lower triangle");
require_near(matrix(0, 1), 0.2, 1e-7, "GRM was not symmetrized");
require_near(matrix(2, 0), -0.1, 1e-7, "wrong GRM packed order");
require_near(matrix(1, 2), 0.15, 1e-7, "wrong GRM upper triangle");
const auto spectral =
spectra::reml::read_and_diagonalize_gcta_grm(grm_path, 3);
require(spectral.eigenvalues.size() == 3, "wrong eigenvalue count");
require(std::is_sorted(spectral.eigenvalues.begin(),
spectral.eigenvalues.end()),
"LAPACK eigenvalues are not ascending");
require(spectral.minimum_eigenvalue > 0.0,
"positive-definite fixture has non-positive eigenvalue");
// Reconstruct U diag(lambda) U' and compare to the original matrix.
for (std::size_t row = 0; row < 3; ++row) {
for (std::size_t col = 0; col < 3; ++col) {
double reconstructed = 0.0;
for (std::size_t axis = 0; axis < 3; ++axis) {
reconstructed += spectral.eigenvectors(row, axis) *
spectral.eigenvalues[axis] *
spectral.eigenvectors(col, axis);
}
require_near(reconstructed, matrix(row, col), 2e-12,
"eigendecomposition does not reconstruct GRM");
}
}
}
void test_wrong_file_size_is_rejected() {
TemporaryDirectory temporary;
const auto grm_path = temporary.path() / "short.grm.bin";
write_floats(grm_path, {1.0F, 0.0F});
bool rejected = false;
try {
(void)spectra::reml::read_gcta_grm_lower_triangle(grm_path, 2);
} catch (const std::runtime_error&) {
rejected = true;
}
require(rejected, "short GRM file was not rejected");
}
} // namespace
int main() {
try {
test_read_and_decompose();
test_wrong_file_size_is_rejected();
std::cout << "test_grm_io: PASS\n";
return EXIT_SUCCESS;
} catch (const std::exception& exception) {
std::cerr << "test_grm_io: FAIL: " << exception.what() << '\n';
return EXIT_FAILURE;
}
}

226
tests/cpp/test_reml_synthetic.cpp Normal file
View File

@@ -0,0 +1,226 @@
#include "spectra_reml/reml.hpp"
#include <algorithm>
#include <cmath>
#include <cstdlib>
#include <iostream>
#include <stdexcept>
#include <string>
#include <vector>
namespace {
void require(bool condition, const std::string& message) {
if (!condition) {
throw std::runtime_error(message);
}
}
void require_near(double actual, double expected, double relative_tolerance,
const std::string& message) {
const double scale = std::max({1.0, std::abs(actual), std::abs(expected)});
if (!std::isfinite(actual) || !std::isfinite(expected) ||
std::abs(actual - expected) > relative_tolerance * scale) {
throw std::runtime_error(message + ": actual=" +
std::to_string(actual) +
", expected=" + std::to_string(expected));
}
}
struct Fixture {
std::vector<double> y;
spectra::reml::ColMajorMatrix x;
std::vector<double> lambda;
Fixture() : x(18, 3) {
lambda = {0.12, 0.18, 0.25, 0.31, 0.39, 0.48, 0.58, 0.69, 0.81,
0.95, 1.10, 1.28, 1.47, 1.69, 1.94, 2.22, 2.55, 2.91};
y.resize(lambda.size());
for (std::size_t i = 0; i < lambda.size(); ++i) {
const double z = (static_cast<double>(i) - 8.5) / 5.0;
x(i, 0) = 1.0;
x(i, 1) = z;
x(i, 2) = (i % 3 == 0 ? -0.8 : (i % 3 == 1 ? 0.1 : 0.7));
const double deterministic_noise =
std::sin(1.7 * static_cast<double>(i) + 0.3) *
std::sqrt(0.35 + 0.55 * lambda[i]) +
0.22 * std::cos(0.37 * static_cast<double>(i));
y[i] = 0.4 - 0.25 * x(i, 1) + 0.15 * x(i, 2) +
deterministic_noise;
}
}
};
void test_score_matches_finite_difference() {
Fixture fixture;
constexpr double sigma_e = 0.73;
constexpr double sigma_g = 0.46;
const auto evaluated = spectra::reml::evaluate_reml_spectral(
fixture.y, fixture.x, fixture.lambda, sigma_e, sigma_g, true, true);
require(evaluated.valid, "reference REML evaluation failed: " +
evaluated.error);
require(evaluated.beta.size() == fixture.x.cols(),
"beta has wrong length");
require(evaluated.beta_covariance.size() ==
fixture.x.cols() * fixture.x.cols(),
"beta covariance has wrong size");
const double epsilon_e = 2e-6 * std::max(1.0, std::abs(sigma_e));
const double epsilon_g = 2e-6 * std::max(1.0, std::abs(sigma_g));
const auto e_plus = spectra::reml::evaluate_reml_spectral(
fixture.y, fixture.x, fixture.lambda, sigma_e + epsilon_e, sigma_g,
false, false);
const auto e_minus = spectra::reml::evaluate_reml_spectral(
fixture.y, fixture.x, fixture.lambda, sigma_e - epsilon_e, sigma_g,
false, false);
const auto g_plus = spectra::reml::evaluate_reml_spectral(
fixture.y, fixture.x, fixture.lambda, sigma_e, sigma_g + epsilon_g,
false, false);
const auto g_minus = spectra::reml::evaluate_reml_spectral(
fixture.y, fixture.x, fixture.lambda, sigma_e, sigma_g - epsilon_g,
false, false);
require(e_plus.valid && e_minus.valid && g_plus.valid && g_minus.valid,
"finite-difference REML evaluation failed");
const double finite_difference_e =
(e_plus.log_likelihood - e_minus.log_likelihood) / (2.0 * epsilon_e);
const double finite_difference_g =
(g_plus.log_likelihood - g_minus.log_likelihood) / (2.0 * epsilon_g);
require_near(evaluated.gradient_e, finite_difference_e, 2e-6,
"sigma_e score disagrees with finite difference");
require_near(evaluated.gradient_g, finite_difference_g, 2e-6,
"sigma_g score disagrees with finite difference");
require(evaluated.ai_ee > 0.0 && evaluated.ai_gg > 0.0,
"AI diagonal must be positive");
require(evaluated.ai_ee * evaluated.ai_gg -
evaluated.ai_eg * evaluated.ai_eg >=
-1e-10,
"AI matrix must be positive semidefinite");
}
void test_signed_parameterization_and_row_permutation() {
Fixture fixture;
const auto positive = spectra::reml::evaluate_reml_spectral(
fixture.y, fixture.x, fixture.lambda, 0.61, 0.52, true, false);
const auto negative = spectra::reml::evaluate_reml_spectral(
fixture.y, fixture.x, fixture.lambda, -0.61, -0.52, true, false);
require(positive.valid && negative.valid, "signed evaluations failed");
require_near(positive.log_likelihood, negative.log_likelihood, 1e-13,
"likelihood must depend on squared standard deviations");
require_near(positive.gradient_e, -negative.gradient_e, 1e-12,
"sigma_e gradient must change sign");
require_near(positive.gradient_g, -negative.gradient_g, 1e-12,
"sigma_g gradient must change sign");
std::vector<double> permuted_y(fixture.y.rbegin(), fixture.y.rend());
std::vector<double> permuted_lambda(fixture.lambda.rbegin(),
fixture.lambda.rend());
spectra::reml::ColMajorMatrix permuted_x(fixture.x.rows(), fixture.x.cols());
for (std::size_t i = 0; i < fixture.x.rows(); ++i) {
for (std::size_t j = 0; j < fixture.x.cols(); ++j) {
permuted_x(i, j) = fixture.x(fixture.x.rows() - 1 - i, j);
}
}
const auto permuted = spectra::reml::evaluate_reml_spectral(
permuted_y, permuted_x, permuted_lambda, 0.61, 0.52, true, false);
require(permuted.valid, "permuted evaluation failed");
require_near(positive.log_likelihood, permuted.log_likelihood, 2e-13,
"orthogonal-axis permutation changed REML likelihood");
require_near(positive.gradient_e, permuted.gradient_e, 2e-12,
"orthogonal-axis permutation changed sigma_e score");
require_near(positive.gradient_g, permuted.gradient_g, 2e-12,
"orthogonal-axis permutation changed sigma_g score");
}
void test_ai_reml_fit_improves_likelihood() {
Fixture fixture;
spectra::reml::RemlOptions options;
options.initial_sigma_e = 0.9;
options.initial_sigma_g = 0.9;
options.max_iterations = 150;
options.gradient_absolute_tolerance = 1e-8;
options.gradient_relative_tolerance = 1e-9;
const auto initial = spectra::reml::evaluate_reml_spectral(
fixture.y, fixture.x, fixture.lambda, options.initial_sigma_e,
options.initial_sigma_g, true, false);
require(initial.valid, "initial fit evaluation failed");
const auto fitted = spectra::reml::fit_ai_reml_spectral(
fixture.y, fixture.x, fixture.lambda, options);
require(fitted.has_estimates(),
"AI-REML returned no estimates: " + fitted.error);
require(fitted.status == spectra::reml::FitStatus::converged,
std::string("AI-REML did not converge: ") +
spectra::reml::to_string(fitted.status) + " " + fitted.error);
require(fitted.log_likelihood >= initial.log_likelihood - 1e-10,
"AI-REML decreased the likelihood");
require(fitted.gradient_inf_norm < 2e-5,
"AI-REML final score is too large");
require(fitted.beta.size() == fixture.x.cols(),
"fitted beta has wrong length");
require(fitted.beta_covariance_packed_lower.size() ==
fixture.x.cols() * (fixture.x.cols() + 1) / 2,
"packed beta covariance has wrong length");
require_near(fitted.sigma_e2, fitted.sigma_e * fitted.sigma_e, 1e-14,
"reported sigma_e2 is inconsistent");
require_near(fitted.sigma_g2, fitted.sigma_g * fitted.sigma_g, 1e-14,
"reported sigma_g2 is inconsistent");
}
void test_residual_only_kkt_boundary() {
constexpr std::size_t n = 14;
spectra::reml::ColMajorMatrix x(n, 1);
std::vector<double> lambda(n);
std::vector<double> y(n, 0.0);
for (std::size_t i = 0; i < n; ++i) {
x(i, 0) = 1.0;
lambda[i] = 0.1 * static_cast<double>(i + 1);
}
// The residual energy is confined to the two smallest GRM eigenvalues,
// making the one-sided score for v_g at zero strictly negative.
y[0] = 1.0;
y[1] = -1.0;
const double expected_sigma_e2 = 2.0 / static_cast<double>(n - 1);
const auto boundary = spectra::reml::evaluate_reml_spectral(
y, x, lambda, std::sqrt(expected_sigma_e2), 0.0, false, true);
require(boundary.valid, "residual-only boundary evaluation failed: " +
boundary.error);
require(boundary.variance_score_g < -1e-3,
"fixture does not have a negative one-sided genetic score");
spectra::reml::RemlOptions options;
options.max_iterations = 100;
options.gradient_absolute_tolerance = 1e-10;
options.gradient_relative_tolerance = 1e-11;
const auto fitted =
spectra::reml::fit_ai_reml_spectral(y, x, lambda, options);
require(fitted.has_estimates(),
"boundary AI-REML returned no estimates: " + fitted.error);
require(fitted.status == spectra::reml::FitStatus::converged_boundary,
std::string("residual-only optimum was not accepted as a KKT boundary: ") +
spectra::reml::to_string(fitted.status) + " " + fitted.error);
require(fitted.iterations <= 10,
"near-boundary KKT check did not terminate promptly");
require_near(fitted.sigma_g2, 0.0, 1e-15,
"genetic boundary must report exactly zero variance");
require_near(fitted.h2, 0.0, 1e-15,
"genetic boundary must report exactly zero h2");
require_near(fitted.sigma_e2, expected_sigma_e2, 2e-12,
"residual-only variance is not RSS/(n-p)");
require_near(fitted.log_likelihood, boundary.log_likelihood, 2e-12,
"reported boundary likelihood is inconsistent");
}
} // namespace
int main() {
try {
test_score_matches_finite_difference();
test_signed_parameterization_and_row_permutation();
test_ai_reml_fit_improves_likelihood();
test_residual_only_kkt_boundary();
std::cout << "test_reml_synthetic: PASS\n";
return EXIT_SUCCESS;
} catch (const std::exception& exception) {
std::cerr << "test_reml_synthetic: FAIL: " << exception.what() << '\n';
return EXIT_FAILURE;
}
}

163
tests/python/test_spectra_reml_cli.py Normal file
View File

@@ -0,0 +1,163 @@
from __future__ import annotations
import csv
import gzip
import importlib.util
import json
import tempfile
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest import mock
import numpy as np
SCRIPT = Path(__file__).resolve().parents[2] / "python" / "spectra_reml.py"
SPEC = importlib.util.spec_from_file_location("spectra_reml", SCRIPT)
assert SPEC is not None and SPEC.loader is not None
cli = importlib.util.module_from_spec(SPEC)
SPEC.loader.exec_module(cli)
class SpectraRemlCliTests(unittest.TestCase):
def fixture(self, root: Path) -> tuple[Path, dict]:
n, p0 = 3, 2
np.zeros(n * (n + 1) // 2, dtype="<f4").tofile(root / "grm.bin")
(root / "grm.id").write_text("F1 I1\nF2 I2\nF3 I3\n", encoding="utf-8")
np.asarray([[1, -1], [1, 0], [1, 1]], dtype="<f8").tofile(root / "base.bin")
np.zeros((2, n), dtype="<f8").tofile(root / "phenotypes.bin")
np.asarray([[0, 1, 2]], dtype="<f4").tofile(root / "extra.bin")
(root / "tasks.tsv").write_text(
"task_index\ttask_id\tphenotype_row\tn_extra_covariates\n"
"0\ttrait_a\t0\t0\n"
"1\ttrait_b\t1\t1\n",
encoding="utf-8",
)
np.asarray([0, 0, 1], dtype="<i8").tofile(root / "offsets.bin")
np.asarray([0], dtype="<i4").tofile(root / "indices.bin")
manifest_path = root / "manifest.json"
args = cli.build_parser().parse_args(
[
"make-manifest",
"--manifest", str(manifest_path),
"--grm-bin", str(root / "grm.bin"),
"--grm-id", str(root / "grm.id"),
"--base-x", str(root / "base.bin"),
"--phenotypes", str(root / "phenotypes.bin"),
"--extra-covariates", str(root / "extra.bin"),
"--tasks", str(root / "tasks.tsv"),
"--extra-offsets", str(root / "offsets.bin"),
"--extra-indices", str(root / "indices.bin"),
"--output-dir", str(root / "blocks"),
"--n-samples", "3",
"--n-base-covariates", "2",
"--n-phenotype-rows", "2",
"--n-extra-covariate-rows", "1",
]
)
cli.make_manifest(args)
return manifest_path, cli.validate_manifest(manifest_path)
@staticmethod
def write_blocks(root: Path) -> None:
root.mkdir(parents=True, exist_ok=True)
(root / "block_000000.summary.tsv").write_text(
"\t".join(cli.SUMMARY_HEADER) + "\n"
"0\ttrait_a\tconverged\t2\t0\t0\t0\t1\t1\t0.5\t-1\t4\t8\t1e-8\t\n"
"1\ttrait_b\tconverged_boundary\t3\t1\t2\t3\t0\t1\t0\t-2\t3\t6\t1e-9\t\n",
encoding="utf-8",
)
np.asarray([1, 2, 3, 4, 5], dtype="<f8").tofile(root / "block_000000.beta.f64.bin")
np.arange(1, 10, dtype="<f8").tofile(root / "block_000000.cov.f64.bin")
(root / "block_000000.complete").write_text(
"format\t{}\nblock\t0\ntasks\t2\nbeta_elements\t5\ncov_elements\t9\n".format(cli.BLOCK_FORMAT),
encoding="utf-8",
)
def test_manifest_and_generic_command(self) -> None:
with tempfile.TemporaryDirectory() as temporary:
root = Path(temporary)
_, manifest = self.fixture(root)
args = cli.build_parser().parse_args(
["run", "--manifest", str(root / "manifest.json"), "--engine", "spectra_reml"]
)
command = cli.build_engine_command(Path("spectra_reml"), manifest, args)
for option in (
"--phenotypes",
"--extra-covariates",
"--extra-offsets",
"--extra-indices",
"--n-phenotype-rows",
"--n-extra-covariate-rows",
):
self.assertEqual(command.count(option), 1)
def test_signature_recovery_and_finalize(self) -> None:
with tempfile.TemporaryDirectory() as temporary:
root = Path(temporary)
manifest_path, manifest = self.fixture(root)
engine = root / "spectra_reml"
engine.write_bytes(b"synthetic engine")
args = cli.build_parser().parse_args(
["run", "--manifest", str(manifest_path), "--engine", str(engine), "--resume"]
)
with mock.patch.object(cli.subprocess, "run", return_value=SimpleNamespace(returncode=0)):
cli.run_engine(manifest, args)
output = Path(manifest["paths"]["output_directory"])
signature = cli.read_signature(output / "run.signature.json")
self.assertEqual(signature["format"], cli.RUN_SIGNATURE_FORMAT)
self.write_blocks(output)
changed = cli.build_parser().parse_args(
[
"run", "--manifest", str(manifest_path), "--engine", str(engine),
"--resume", "--max-iterations", "101",
]
)
with self.assertRaisesRegex(RuntimeError, "signature differs"):
cli.run_engine(manifest, changed)
result = root / "results.tsv.gz"
cli.finalize(manifest, result)
with gzip.open(result, "rt", encoding="utf-8", newline="") as handle:
rows = list(csv.DictReader(handle, delimiter="\t"))
self.assertEqual([row["task_id"] for row in rows], ["trait_a", "trait_b"])
self.assertEqual(json.loads(rows[0]["beta_json"]), [1.0, 2.0])
self.assertEqual(len(json.loads(rows[1]["covariance_packed_lower_json"])), 6)
altered = dict(manifest)
altered["created_utc"] = "changed"
with self.assertRaisesRegex(RuntimeError, "does not belong"):
cli.finalize(altered, root / "unsafe.tsv")
def test_force_invalidates_before_cleanup_and_dry_run_is_read_only(self) -> None:
with tempfile.TemporaryDirectory() as temporary:
root = Path(temporary)
manifest_path, manifest = self.fixture(root)
engine = root / "spectra_reml"
engine.write_bytes(b"engine")
output = Path(manifest["paths"]["output_directory"])
self.write_blocks(output)
resume_args = cli.build_parser().parse_args(
["run", "--manifest", str(manifest_path), "--engine", str(engine), "--resume"]
)
cli.atomic_json(output / "run.signature.json", cli.make_signature(manifest, engine, resume_args))
before = sorted(path.name for path in output.iterdir())
dry = cli.build_parser().parse_args(
["run", "--manifest", str(manifest_path), "--engine", "missing", "--force", "--dry-run"]
)
cli.run_engine(manifest, dry)
self.assertEqual(before, sorted(path.name for path in output.iterdir()))
force = cli.build_parser().parse_args(
["run", "--manifest", str(manifest_path), "--engine", str(engine), "--force"]
)
with mock.patch.object(cli, "block_files", side_effect=RuntimeError("cleanup interrupted")):
with self.assertRaisesRegex(RuntimeError, "cleanup interrupted"):
cli.run_engine(manifest, force)
self.assertFalse((output / "run.signature.json").exists())
if __name__ == "__main__":
unittest.main()

128
tests/test_math_reference.py Normal file
View File

@@ -0,0 +1,128 @@
#!/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")