Improve phenotype scaling and safeguarded line search

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

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:

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;

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)

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 =

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<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";
+                return return_fallback(
-                break;
+                    "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;
+        return return_fallback(
-        if (zoom_result.error.empty()) {
+            "strong-Wolfe zoom exceeded its evaluation limit");
-            zoom_result.error = "strong-Wolfe zoom exceeded its evaluation limit";
-        }
-        return zoom_result;
     };
 
     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 =
+            return return_fallback(
-                "strong-Wolfe curvature condition not reached at maximum step";
+                "strong-Wolfe curvature condition not reached at maximum step");
-            return result;
         }
         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 return_fallback(
-    return result;
+        "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 auto original_problem =
-        const InitialGuess guess = initial_guess(problem, options);
+            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;

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

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)