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

Improve phenotype scaling and safeguarded line search

Browse Source
This commit is contained in:
Jiarui Li
2026-08-09 15:05:40 +08:00
parent da971a2ca7
commit 76dd8b1379
8 changed files with 226 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
README.md
View File

@@ -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

11
docs/FORMAT.md
View File

@@ -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:

2
include/spectra_reml/types.hpp
View File

@@ -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;

4
python/spectra_reml.py
View File

@@ -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)

14
src/main.cpp
View File

@@ -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 =

157
src/reml_core.cpp
View File

@@ -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<double>::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<double>::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<double>::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<double>(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<double>& 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<double> 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<double>& 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;

44
tests/cpp/test_reml_synthetic.cpp
View File

@@ -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<double> 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<double>(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;

2
tests/python/test_spectra_reml_cli.py
View File

@@ -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)