diff --git a/README.md b/README.md index a33303f..68508ad 100644 --- a/README.md +++ b/README.md @@ -17,6 +17,11 @@ The engine diagonalizes the GRM once, rotates the common design and every unique - C++17 numerical core using oneMKL or OpenBLAS/LAPACKE. - GRM eigendecomposition once per task set. - AI-REML with a strong-Wolfe line search; no EM updates. +- Per-task OLS-residual phenotype scaling inside the numerical core, with all + estimates and the restricted likelihood restored to the input phenotype units. +- Safeguarded quadratic zoom interpolation (central 96% of the bracket), + bisection fallback, and a last-valid-improving-point fallback when strict Wolfe + curvature cannot be reached because of numerical roundoff. - Signed standard-deviation parameterization with a separate one-sided KKT check for the \(\sigma_g^2=0\) boundary. - No explicit dense \(P\) matrix. - Variable numbers of extra fixed-effect covariates per task through a CSR index. @@ -25,6 +30,13 @@ The engine diagonalizes the GRM once, rotates the common design and every unique The numerical core has no knowledge of cohorts, molecular assay types, or domain-specific variable names. +The line search starts at 1, expands by 1.618 up to the configured maximum +step, and uses quadratic interpolation only when its stationary point lies at +least 2% away from both bracket endpoints. An invalid interpolation falls back +to bisection. If the evaluation limit is reached, the search accepts the last +finite covariance-valid point that improved the likelihood; it reports +`line_search_failed` only when no such point exists. + ## Repository layout ```text diff --git a/docs/FORMAT.md b/docs/FORMAT.md index 79c6fa8..0fe356f 100644 --- a/docs/FORMAT.md +++ b/docs/FORMAT.md @@ -150,6 +150,17 @@ numerical_error `converged_boundary` is a successful residual-only solution accepted after the one-sided variance-component score and likelihood checks. +Phenotypes are scaled internally by their task-specific OLS residual RMS before +optimization. Reported fixed effects, fixed-effect covariance, variance +components, and restricted likelihood are transformed back to the original +input phenotype units. Input binary files are never modified. + +The strong-Wolfe zoom uses safeguarded quadratic interpolation with a 2% +endpoint margin and falls back to bisection. At an evaluation limit or a +collapsed bracket, the last valid likelihood-improving point encountered is +accepted. `line_search_failed` therefore means that the search found no valid +point that improved the starting likelihood. + ## Generic finalized output The Python CLI exports one TSV row per task. It includes the full summary plus: diff --git a/include/spectra_reml/types.hpp b/include/spectra_reml/types.hpp index a6a3505..64ec3bb 100644 --- a/include/spectra_reml/types.hpp +++ b/include/spectra_reml/types.hpp @@ -76,6 +76,8 @@ struct RemlOptions { double wolfe_c2 = 0.9; double initial_line_search_step = 1.0; double maximum_line_search_step = 64.0; + double line_search_expansion_factor = 1.618; + double zoom_safeguard_fraction = 0.02; double covariance_floor_relative = 1e-12; double rank_tolerance_relative = 1e-10; double ai_ridge_relative = 1e-10; diff --git a/python/spectra_reml.py b/python/spectra_reml.py index 69420aa..05048f1 100644 --- a/python/spectra_reml.py +++ b/python/spectra_reml.py @@ -378,6 +378,8 @@ def build_engine_command(engine: Path, manifest: Mapping[str, Any], args: argpar ("wolfe_c2", "--wolfe-c2"), ("initial_step", "--initial-step"), ("maximum_step", "--maximum-step"), + ("line_search_expansion", "--line-search-expansion"), + ("zoom_safeguard", "--zoom-safeguard"), ("rank_tol", "--rank-tol"), ("covariance_floor", "--covariance-floor"), ("boundary_h2_trigger", "--boundary-h2-trigger"), @@ -685,6 +687,8 @@ def add_run_options(parser: argparse.ArgumentParser) -> None: parser.add_argument("--wolfe-c2", type=float, default=0.9) parser.add_argument("--initial-step", type=float, default=1.0) parser.add_argument("--maximum-step", type=float, default=64.0) + parser.add_argument("--line-search-expansion", type=float, default=1.618) + parser.add_argument("--zoom-safeguard", type=float, default=0.02) parser.add_argument("--rank-tol", type=float, default=1e-10) parser.add_argument("--covariance-floor", type=float, default=1e-12) parser.add_argument("--boundary-h2-trigger", type=float, default=1e-3) diff --git a/src/main.cpp b/src/main.cpp index 0c8b4e3..0c2328b 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -55,6 +55,8 @@ void print_usage(std::ostream& output) { << " --wolfe-c2 X default 0.9\n" << " --initial-step X default 1\n" << " --maximum-step X default 64\n" + << " --line-search-expansion X default 1.618\n" + << " --zoom-safeguard X bracket fraction, default 0.02\n" << " --initial-sigma-e X override OLS-based initial value\n" << " --initial-sigma-g X override OLS-based initial value\n" << " --rank-tol X default 1e-10\n" @@ -188,6 +190,8 @@ void reject_unknown_options(const Arguments& arguments) { "--wolfe-c2", "--initial-step", "--maximum-step", + "--line-search-expansion", + "--zoom-safeguard", "--initial-sigma-e", "--initial-sigma-g", "--rank-tol", @@ -301,6 +305,16 @@ int main(int argc, char** argv) { options.reml.maximum_line_search_step = parse_double(value, "--maximum-step", true); }); + set_if_present(arguments, "--line-search-expansion", + [&](const std::string& value) { + options.reml.line_search_expansion_factor = + parse_double(value, "--line-search-expansion", true); + }); + set_if_present(arguments, "--zoom-safeguard", + [&](const std::string& value) { + options.reml.zoom_safeguard_fraction = + parse_double(value, "--zoom-safeguard", true, true); + }); set_if_present(arguments, "--initial-sigma-e", [&](const std::string& value) { options.reml.initial_sigma_e = diff --git a/src/reml_core.cpp b/src/reml_core.cpp index 8b23d55..9ef1973 100644 --- a/src/reml_core.cpp +++ b/src/reml_core.cpp @@ -326,6 +326,7 @@ struct InitialGuess { double sigma_e = 0.0; double sigma_g = 0.0; double residual_mean_square = 0.0; + double phenotype_scale = 1.0; std::string error; }; @@ -379,6 +380,10 @@ InitialGuess initial_guess(const ValidatedProblem& problem, const double sigma_floor = std::max(1e-12, std::sqrt(std::numeric_limits::epsilon()) * std::max(1.0, phenotype_rms)); + const double residual_scale = std::sqrt(residual_mean_square); + if (std::isfinite(residual_scale) && residual_scale > sigma_floor) { + result.phenotype_scale = residual_scale; + } const double paper_initial = std::max(sigma_floor, std::sqrt(0.5 * residual_mean_square)); result.sigma_e = std::isfinite(options.initial_sigma_e) @@ -432,6 +437,7 @@ bool ai_direction(const RemlEvaluation& evaluation, struct LineSearchResult { bool success = false; + bool used_fallback = false; double alpha = 0.0; RemlEvaluation evaluation; std::size_t evaluations = 0; @@ -467,26 +473,90 @@ LineSearchResult strong_wolfe_line_search( options.wolfe_c1 * alpha * derivative_zero; }; + // Preserve the latest finite, covariance-valid improving point. If + // roundoff prevents the curvature condition from being met after the + // bracket has already localized the optimum, this point is safer than + // discarding the whole AI iteration. + bool have_fallback = false; + double fallback_alpha = 0.0; + RemlEvaluation fallback_evaluation; + auto remember_fallback = [&](double alpha, + const RemlEvaluation& evaluation) { + if (alpha > 0.0 && evaluation.valid && + evaluation.log_likelihood > initial.log_likelihood) { + have_fallback = true; + fallback_alpha = alpha; + fallback_evaluation = evaluation; + } + }; + auto return_fallback = [&](const std::string& reason) { + LineSearchResult fallback; + fallback.evaluations = result.evaluations; + fallback.error = reason; + if (have_fallback) { + fallback.success = true; + fallback.used_fallback = true; + fallback.alpha = fallback_alpha; + fallback.evaluation = fallback_evaluation; + fallback.error += "; accepted last valid improving point"; + } + return fallback; + }; + auto zoom = [&](double alpha_lo, double alpha_hi, - RemlEvaluation evaluation_lo) -> LineSearchResult { + RemlEvaluation evaluation_lo, + RemlEvaluation evaluation_hi) -> LineSearchResult { LineSearchResult zoom_result; // Keep the global evaluation counter in result; copy it at return. for (std::size_t iteration = 0; iteration < options.line_search_max_zoom_iterations && result.evaluations < options.line_search_max_evaluations; ++iteration) { - const double alpha = 0.5 * (alpha_lo + alpha_hi); + const double interval = alpha_hi - alpha_lo; + const double lower = std::min(alpha_lo, alpha_hi); + const double upper = std::max(alpha_lo, alpha_hi); + const double width = upper - lower; const double interval_scale = std::max({1.0, std::abs(alpha_lo), std::abs(alpha_hi)}); - if (std::abs(alpha_hi - alpha_lo) <= + if (width <= 8.0 * std::numeric_limits::epsilon() * interval_scale) { - zoom_result.error = "strong-Wolfe zoom interval collapsed"; - break; + return return_fallback( + "strong-Wolfe zoom interval collapsed"); + } + + // Fit the quadratic passing through f(lo), f(hi), and f'(lo). + // Accept its stationary point only inside the central 96% of the + // bracket. Otherwise use bisection, which always contracts it. + double alpha = 0.5 * (alpha_lo + alpha_hi); + const double derivative_lo = + directional_derivative(evaluation_lo, direction); + const double denominator = interval * interval; + if (evaluation_lo.valid && evaluation_hi.valid && + std::isfinite(derivative_lo) && denominator > 0.0) { + const double quadratic = + (evaluation_hi.log_likelihood - + evaluation_lo.log_likelihood - + interval * derivative_lo) / + denominator; + if (std::isfinite(quadratic) && + quadratic < -std::numeric_limits::epsilon()) { + const double candidate = + alpha_lo - derivative_lo / (2.0 * quadratic); + const double margin = + options.zoom_safeguard_fraction * width; + if (std::isfinite(candidate) && + candidate >= lower + margin && + candidate <= upper - margin) { + alpha = candidate; + } + } } RemlEvaluation trial = evaluate_alpha(alpha); + remember_fallback(alpha, trial); if (!trial.valid || trial.log_likelihood < armijo_bound(alpha) || trial.log_likelihood <= evaluation_lo.log_likelihood) { alpha_hi = alpha; + evaluation_hi = std::move(trial); continue; } const double derivative = @@ -501,15 +571,13 @@ LineSearchResult strong_wolfe_line_search( } if (derivative * (alpha_hi - alpha_lo) <= 0.0) { alpha_hi = alpha_lo; + evaluation_hi = evaluation_lo; } alpha_lo = alpha; evaluation_lo = std::move(trial); } - zoom_result.evaluations = result.evaluations; - if (zoom_result.error.empty()) { - zoom_result.error = "strong-Wolfe zoom exceeded its evaluation limit"; - } - return zoom_result; + return return_fallback( + "strong-Wolfe zoom exceeded its evaluation limit"); }; double alpha_previous = 0.0; @@ -520,10 +588,12 @@ LineSearchResult strong_wolfe_line_search( result.evaluations < options.line_search_max_evaluations; ++iteration) { RemlEvaluation trial = evaluate_alpha(alpha); + remember_fallback(alpha, trial); if (!trial.valid || trial.log_likelihood < armijo_bound(alpha) || (iteration > 0 && trial.log_likelihood <= evaluation_previous.log_likelihood)) { - return zoom(alpha_previous, alpha, std::move(evaluation_previous)); + return zoom(alpha_previous, alpha, + std::move(evaluation_previous), std::move(trial)); } const double derivative = directional_derivative(trial, direction); if (std::abs(derivative) <= @@ -534,19 +604,20 @@ LineSearchResult strong_wolfe_line_search( return result; } if (derivative <= 0.0) { - return zoom(alpha, alpha_previous, std::move(trial)); + return zoom(alpha, alpha_previous, std::move(trial), + std::move(evaluation_previous)); } if (alpha >= options.maximum_line_search_step) { - result.error = - "strong-Wolfe curvature condition not reached at maximum step"; - return result; + return return_fallback( + "strong-Wolfe curvature condition not reached at maximum step"); } alpha_previous = alpha; evaluation_previous = std::move(trial); - alpha = std::min(2.0 * alpha, options.maximum_line_search_step); + alpha = std::min(options.line_search_expansion_factor * alpha, + options.maximum_line_search_step); } - result.error = "strong-Wolfe line search exceeded its evaluation limit"; - return result; + return return_fallback( + "strong-Wolfe line search exceeded its evaluation limit"); } void finalize_result(RemlResult& result, const ValidatedProblem& problem, @@ -583,6 +654,33 @@ void finalize_result(RemlResult& result, const ValidatedProblem& problem, } } +void restore_phenotype_scale(RemlResult& result, double phenotype_scale, + std::size_t residual_degrees_of_freedom) { + if (!(phenotype_scale > 0.0) || !std::isfinite(phenotype_scale) || + phenotype_scale == 1.0) { + return; + } + const double variance_scale = phenotype_scale * phenotype_scale; + result.sigma_e *= phenotype_scale; + result.sigma_g *= phenotype_scale; + result.sigma_e2 = result.sigma_e * result.sigma_e; + result.sigma_g2 = result.sigma_g * result.sigma_g; + if (std::isfinite(result.log_likelihood)) { + result.log_likelihood -= + static_cast(residual_degrees_of_freedom) * + std::log(phenotype_scale); + } + if (std::isfinite(result.gradient_inf_norm)) { + result.gradient_inf_norm /= phenotype_scale; + } + for (double& value : result.beta) { + value *= phenotype_scale; + } + for (double& value : result.beta_covariance_packed_lower) { + value *= variance_scale; + } +} + void validate_options(const RemlOptions& options) { if (options.max_iterations == 0) { throw std::invalid_argument("max_iterations must be positive"); @@ -596,6 +694,12 @@ void validate_options(const RemlOptions& options) { options.initial_line_search_step)) { throw std::invalid_argument("invalid line-search step bounds"); } + if (!(options.line_search_expansion_factor > 1.0) || + !(options.zoom_safeguard_fraction >= 0.0 && + options.zoom_safeguard_fraction < 0.5)) { + throw std::invalid_argument( + "line-search expansion must exceed 1 and zoom safeguard must be in [0,0.5)"); + } if (!(options.covariance_floor_relative >= 0.0) || !(options.rank_tolerance_relative >= 0.0) || !(options.boundary_score_tolerance >= 0.0) || @@ -654,13 +758,25 @@ RemlResult fit_ai_reml_spectral(const std::vector& y_star, RemlResult result; try { validate_options(options); - const auto problem = validate_problem(y_star, x_star, eigenvalues); - const InitialGuess guess = initial_guess(problem, options); + const auto original_problem = + validate_problem(y_star, x_star, eigenvalues); + InitialGuess guess = initial_guess(original_problem, options); if (!guess.valid) { result.status = FitStatus::rank_deficient; result.error = guess.error; return result; } + const double phenotype_scale = guess.phenotype_scale; + std::vector normalized_y(y_star.size()); + for (std::size_t i = 0; i < y_star.size(); ++i) { + normalized_y[i] = y_star[i] / phenotype_scale; + } + const auto problem = + validate_problem(normalized_y, x_star, eigenvalues); + guess.sigma_e /= phenotype_scale; + guess.sigma_g /= phenotype_scale; + guess.residual_mean_square /= + phenotype_scale * phenotype_scale; double sigma_e = guess.sigma_e; double sigma_g = guess.sigma_g; RemlEvaluation current = @@ -827,6 +943,7 @@ RemlResult fit_ai_reml_spectral(const std::vector& y_star, (void)accept_genetic_boundary(current); } finalize_result(result, problem, sigma_e, sigma_g, options); + restore_phenotype_scale(result, phenotype_scale, problem.n - problem.p); return result; } catch (const std::invalid_argument& exception) { result.status = FitStatus::invalid_input; diff --git a/tests/cpp/test_reml_synthetic.cpp b/tests/cpp/test_reml_synthetic.cpp index a8a265a..bc57f54 100644 --- a/tests/cpp/test_reml_synthetic.cpp +++ b/tests/cpp/test_reml_synthetic.cpp @@ -165,6 +165,49 @@ void test_ai_reml_fit_improves_likelihood() { "reported sigma_g2 is inconsistent"); } +void test_internal_phenotype_scaling_restores_original_units() { + Fixture fixture; + constexpr double scale = 1.0e6; + std::vector scaled_y = fixture.y; + for (double& value : scaled_y) { + value *= scale; + } + spectra::reml::RemlOptions options; + options.max_iterations = 150; + const auto original = spectra::reml::fit_ai_reml_spectral( + fixture.y, fixture.x, fixture.lambda, options); + const auto scaled = spectra::reml::fit_ai_reml_spectral( + scaled_y, fixture.x, fixture.lambda, options); + require(original.has_estimates() && scaled.has_estimates(), + "scale-invariance fixture returned no estimates"); + require(original.status == scaled.status, + "phenotype scaling changed the fit status"); + require_near(scaled.sigma_e2, original.sigma_e2 * scale * scale, 2e-10, + "sigma_e2 was not restored to phenotype units"); + require_near(scaled.sigma_g2, original.sigma_g2 * scale * scale, 2e-10, + "sigma_g2 was not restored to phenotype units"); + require_near(scaled.h2, original.h2, 2e-11, + "phenotype scaling changed h2"); + for (std::size_t i = 0; i < original.beta.size(); ++i) { + require_near(scaled.beta[i], original.beta[i] * scale, 2e-10, + "beta was not restored to phenotype units"); + } + for (std::size_t i = 0; + i < original.beta_covariance_packed_lower.size(); ++i) { + require_near(scaled.beta_covariance_packed_lower[i], + original.beta_covariance_packed_lower[i] * scale * scale, + 3e-10, + "beta covariance was not restored to phenotype units"); + } + const double expected_log_likelihood_shift = + -static_cast(fixture.y.size() - fixture.x.cols()) * + std::log(scale); + require_near(scaled.log_likelihood, + original.log_likelihood + expected_log_likelihood_shift, + 2e-11, + "REML likelihood was not restored to phenotype units"); +} + void test_residual_only_kkt_boundary() { constexpr std::size_t n = 14; spectra::reml::ColMajorMatrix x(n, 1); @@ -216,6 +259,7 @@ int main() { test_score_matches_finite_difference(); test_signed_parameterization_and_row_permutation(); test_ai_reml_fit_improves_likelihood(); + test_internal_phenotype_scaling_restores_original_units(); test_residual_only_kkt_boundary(); std::cout << "test_reml_synthetic: PASS\n"; return EXIT_SUCCESS; diff --git a/tests/python/test_spectra_reml_cli.py b/tests/python/test_spectra_reml_cli.py index 3c36488..93e5526 100644 --- a/tests/python/test_spectra_reml_cli.py +++ b/tests/python/test_spectra_reml_cli.py @@ -90,6 +90,8 @@ class SpectraRemlCliTests(unittest.TestCase): "--extra-indices", "--n-phenotype-rows", "--n-extra-covariate-rows", + "--line-search-expansion", + "--zoom-safeguard", ): self.assertEqual(command.count(option), 1)