227 lines
10 KiB
C++
227 lines
10 KiB
C++
|
#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;
|
||
|
|
}
|
||
|
|
}
|