From 903320194895e294cb6726c1c6d80b9dae912c92 Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Fri, 14 Aug 2026 16:54:57 +0800 Subject: [PATCH] Handle boundary and line-search fixed-effect inference --- README.md | 10 +- docs/FORMAT.md | 35 ++++--- include/spectra_reml/fixed_effects.hpp | 8 +- include/spectra_reml/types.hpp | 9 +- python/spectra_reml.py | 6 +- src/batch_io.cpp | 6 +- src/fixed_effects.cpp | 135 ++++++++++++++++++++++++- src/reml_core.cpp | 8 +- tests/cpp/test_reml_synthetic.cpp | 42 ++++++-- tests/python/test_spectra_reml_cli.py | 2 +- 10 files changed, 220 insertions(+), 41 deletions(-) diff --git a/README.md b/README.md index 88471a7..e501368 100644 --- a/README.md +++ b/README.md @@ -151,14 +151,20 @@ Optional named contrasts are supplied by `--contrast-matrix`, task-specific covariate columns receive zero weights automatically. Finalized results contain contrast estimates, adjusted standard errors, statistics, numerator and denominator degrees of freedom, and p-values as JSON arrays. -This release uses manifest v2, run-signature v2, block v3, and finalized-output -v2 contracts. It intentionally does not read older contracts. +This release uses manifest v2, run-signature v3, block v4, and finalized-output +v3 contracts. It intentionally does not read older block, run-signature, or +finalized-output contracts. Fixed-effect inference defaults to Satterthwaite. Select Kenward-Roger or turn inference off with `--fixed-effect-test kenward-roger` or `--fixed-effect-test none`. Coefficient-wise standard errors, statistics, denominator degrees of freedom, and p-values are exported as JSON arrays; tasks with extra covariates also report their joint F test in the summary columns. +An accepted `sigma_g2=0` boundary is refitted as ordinary least squares without +the GRM term and is reported as `boundary_ols`. A retained +`line_search_failed` iterate continues through the selected fixed-effect +inference and is reported as `line_search_conditional` while preserving the +optimizer status. ## Recovery and provenance diff --git a/docs/FORMAT.md b/docs/FORMAT.md index 2607db7..8c161a1 100644 --- a/docs/FORMAT.md +++ b/docs/FORMAT.md @@ -165,7 +165,7 @@ The covariance array uses the row-wise packed lower triangle: The `.complete` marker is written last and contains tab-separated key/value rows: ```text -format spectra-reml-block-v3 +format spectra-reml-block-v4 block 0 tasks 256 beta_elements 4096 @@ -198,34 +198,39 @@ numerical_error The four `fixed_*` arrays have the same offsets and coefficient order as `beta`. They contain standard errors, statistics, denominator degrees of -freedom, and p-values. With Satterthwaite inference, a coefficient statistic is -a signed t statistic. With Kenward-Roger inference, it is an F statistic with -one numerator degree of freedom. For every task with extra covariates, the -summary also contains an F test of the joint null that all task-specific fixed -effects are zero. +freedom, and p-values. With ordinary least-squares or Satterthwaite inference, +a coefficient statistic is a signed t statistic. With Kenward-Roger inference, +it is an F statistic with one numerator degree of freedom. For every task with +extra covariates, the summary also contains a test of the joint null that all +task-specific fixed effects are zero. The six `contrast_*` arrays share `contrast_test_offset` and follow the -contrast metadata order. A one-row Satterthwaite contrast reports a signed t -statistic; Kenward-Roger reports an F statistic with one numerator degree of -freedom. Both methods retain the signed `L beta` estimate and their adjusted -standard error. Negative `contrast_test_offset` means that no valid contrast -tests were emitted for that task. +contrast metadata order. A one-row ordinary least-squares or Satterthwaite +contrast reports a signed t statistic; Kenward-Roger reports an F statistic +with one numerator degree of freedom. All methods retain the signed `L beta` +estimate and their standard error. Negative `contrast_test_offset` means that +no valid contrast tests were emitted for that task. `fixed_test_status` is one of: ```text not_requested ok -boundary_conditional +boundary_ols +line_search_conditional fit_not_converged invalid_contrast information_singular numerical_error ``` -At `converged_boundary`, inference conditions on the accepted active set -`sigma_g2=0`; `fixed_test_status` is `boundary_conditional` and only residual -variance uncertainty contributes to the small-sample adjustment. +At `converged_boundary`, the random term is removed and the fixed model is +refitted as ordinary least squares. `fixed_test_method` is +`ordinary-least-squares`, `fixed_test_status` is `boundary_ols`, and the +denominator degrees of freedom are `n-rank(X)`. At `line_search_failed`, the +last retained REML iterate is tested with the requested Satterthwaite or +Kenward-Roger method and `fixed_test_status` is +`line_search_conditional`; the optimizer flag remains visible in `status`. Phenotypes are scaled internally by their task-specific OLS residual RMS before optimization. Reported fixed effects, fixed-effect covariance, variance diff --git a/include/spectra_reml/fixed_effects.hpp b/include/spectra_reml/fixed_effects.hpp index c0cdbe1..c445b2c 100644 --- a/include/spectra_reml/fixed_effects.hpp +++ b/include/spectra_reml/fixed_effects.hpp @@ -13,9 +13,11 @@ struct FixedEffectHypothesis { std::vector rhs; }; -// Computes coefficient-wise tests and any supplied general linear hypotheses -// at an already fitted REML solution. Inputs must use the same (possibly GRM- -// rotated) coordinate system and phenotype scale as the supplied fit. +// Computes coefficient-wise tests and any supplied general linear hypotheses. +// At an interior or retained line-search iterate it uses the requested REML +// small-sample method. A sigma_g2=0 boundary fit is refitted and tested as +// ordinary least squares without a GRM term. Inputs must use the same +// orthogonally rotated coordinate system and phenotype scale as the fit. [[nodiscard]] FixedEffectInferenceResult infer_fixed_effects_spectral( const std::vector& y_star, const ColMajorMatrix& x_star, const std::vector& eigenvalues, const RemlResult& fit, diff --git a/include/spectra_reml/types.hpp b/include/spectra_reml/types.hpp index 7b1af02..8de7fc7 100644 --- a/include/spectra_reml/types.hpp +++ b/include/spectra_reml/types.hpp @@ -66,6 +66,7 @@ enum class FitStatus { enum class FixedEffectTestMethod { none, + ordinary_least_squares, satterthwaite, kenward_roger }; @@ -75,7 +76,8 @@ enum class FixedEffectTestMethod { enum class FixedEffectInferenceStatus { not_requested, ok, - boundary_conditional, + boundary_ols, + line_search_conditional, fit_not_converged, invalid_contrast, information_singular, @@ -99,8 +101,9 @@ struct FixedEffectInferenceResult { FixedEffectTestMethod method = FixedEffectTestMethod::none; FixedEffectInferenceStatus status = FixedEffectInferenceStatus::not_requested; - // One test per beta, in design-matrix column order. Satterthwaite reports - // a signed t statistic; Kenward-Roger reports an F statistic with 1 NumDF. + // One test per beta, in design-matrix column order. OLS and + // Satterthwaite report signed t statistics; Kenward-Roger reports an F + // statistic with 1 NumDF. std::vector coefficient_tests; // Optional general linear hypotheses requested by the caller. std::vector hypothesis_tests; diff --git a/python/spectra_reml.py b/python/spectra_reml.py index 14cbf04..fa1b805 100644 --- a/python/spectra_reml.py +++ b/python/spectra_reml.py @@ -32,9 +32,9 @@ except ImportError as exc: # pragma: no cover MANIFEST_FORMAT = "spectra-reml-manifest-v2" -BLOCK_FORMAT = "spectra-reml-block-v3" -RUN_SIGNATURE_FORMAT = "spectra-reml-run-signature-v2" -FINALIZE_FORMAT = "spectra-reml-finalize-v2" +BLOCK_FORMAT = "spectra-reml-block-v4" +RUN_SIGNATURE_FORMAT = "spectra-reml-run-signature-v3" +FINALIZE_FORMAT = "spectra-reml-finalize-v3" TASK_HEADER = ( "task_index", "task_id", diff --git a/src/batch_io.cpp b/src/batch_io.cpp index feb1a25..61a4d5b 100644 --- a/src/batch_io.cpp +++ b/src/batch_io.cpp @@ -34,7 +34,7 @@ namespace spectra::reml { namespace { -constexpr const char* kOutputFormatVersion = "spectra-reml-block-v3"; +constexpr const char* kOutputFormatVersion = "spectra-reml-block-v4"; std::vector split_tab(const std::string& line); std::uint64_t parse_u64(const std::string& text, const char* field, @@ -442,7 +442,9 @@ void write_block_atomic(const std::filesystem::path& output_directory, const auto& inference = fit.fixed_effect_inference; if ((inference.status == FixedEffectInferenceStatus::ok || inference.status == - FixedEffectInferenceStatus::boundary_conditional) && + FixedEffectInferenceStatus::boundary_ols || + inference.status == + FixedEffectInferenceStatus::line_search_conditional) && inference.coefficient_tests.size() == fit.beta.size()) { fixed_test_offsets[index] = static_cast(fixed_se_values.size()); diff --git a/src/fixed_effects.cpp b/src/fixed_effects.cpp index 8aa3632..4ba5ffb 100644 --- a/src/fixed_effects.cpp +++ b/src/fixed_effects.cpp @@ -395,6 +395,110 @@ double contrast_variance(const std::vector& contrast, return quadratic_form(contrast, covariance, p); } +struct OlsWorkspace { + std::size_t n = 0; + std::size_t p = 0; + double denominator_df = std::numeric_limits::quiet_NaN(); + std::vector beta; + Matrix covariance; +}; + +OlsWorkspace make_ols_workspace(const std::vector& y, + const ColMajorMatrix& x) { + OlsWorkspace workspace; + workspace.n = x.rows(); + workspace.p = x.cols(); + if (workspace.n <= workspace.p) { + throw std::runtime_error( + "ordinary least squares requires positive residual degrees of freedom"); + } + const std::vector unit_weights(workspace.n, 1.0); + const Matrix normal = weighted_cross_product(x, unit_weights); + Matrix inverse_normal; + std::string error; + if (!invert_positive_definite(normal, workspace.p, inverse_normal, error)) { + throw std::runtime_error( + "ordinary least-squares design is rank deficient: " + error); + } + std::vector rhs(workspace.p, 0.0); + for (std::size_t col = 0; col < workspace.p; ++col) { + for (std::size_t row = 0; row < workspace.n; ++row) { + rhs[col] += x(row, col) * y[row]; + } + } + workspace.beta.assign(workspace.p, 0.0); + for (std::size_t row = 0; row < workspace.p; ++row) { + for (std::size_t col = 0; col < workspace.p; ++col) { + workspace.beta[row] += + at(inverse_normal, workspace.p, row, col) * rhs[col]; + } + } + double rss = 0.0; + for (std::size_t row = 0; row < workspace.n; ++row) { + double fitted = 0.0; + for (std::size_t col = 0; col < workspace.p; ++col) { + fitted += x(row, col) * workspace.beta[col]; + } + const double residual = y[row] - fitted; + rss += residual * residual; + } + workspace.denominator_df = + static_cast(workspace.n - workspace.p); + const double residual_variance = rss / workspace.denominator_df; + if (!(residual_variance > 0.0) || !std::isfinite(residual_variance)) { + throw std::runtime_error( + "ordinary least-squares residual variance is non-positive"); + } + workspace.covariance = std::move(inverse_normal); + for (double& value : workspace.covariance) { + value *= residual_variance; + } + return workspace; +} + +FixedEffectTestResult ols_test( + const FixedEffectHypothesis& hypothesis, const OlsWorkspace& workspace, + double rank_tolerance_relative) { + const ReducedHypothesis reduced = reduce_hypothesis( + hypothesis, workspace.covariance, workspace.p, + rank_tolerance_relative); + FixedEffectTestResult result; + result.valid = true; + result.numerator_df = reduced.rank; + result.denominator_df = workspace.denominator_df; + double sum_t_squared = 0.0; + for (std::size_t row = 0; row < reduced.rank; ++row) { + double raw_estimate = 0.0; + for (std::size_t col = 0; col < workspace.p; ++col) { + raw_estimate += + reduced.contrast(row, col) * workspace.beta[col]; + } + const double difference = raw_estimate - reduced.rhs[row]; + const double variance = reduced.eigenvalues[row]; + if (!(variance > 0.0) || !std::isfinite(variance)) { + throw std::runtime_error( + "ordinary least-squares contrast variance is non-positive"); + } + sum_t_squared += difference * difference / variance; + if (reduced.rank == 1) { + result.estimate = raw_estimate; + result.standard_error = std::sqrt(variance); + result.statistic = difference / result.standard_error; + } + } + if (reduced.rank == 1) { + result.p_value = student_t_two_sided_p( + result.statistic, result.denominator_df); + } else { + result.statistic = + sum_t_squared / static_cast(reduced.rank); + result.p_value = f_upper_tail( + result.statistic, static_cast(reduced.rank), + result.denominator_df); + } + return result; +} + double satterthwaite_df(const std::vector& contrast, const InferenceWorkspace& workspace, double variance) { @@ -632,9 +736,12 @@ FixedEffectInferenceResult infer_fixed_effects_spectral( return result; } if (fit.status != FitStatus::converged && - fit.status != FitStatus::converged_boundary) { + fit.status != FitStatus::converged_boundary && + fit.status != FitStatus::line_search_failed) { result.status = FixedEffectInferenceStatus::fit_not_converged; - result.error = "fixed-effect tests require a converged REML fit"; + result.error = + "fixed-effect tests require an interior fit, a zero-GRM boundary " + "fit, or a retained line-search iterate"; return result; } if (y_star.size() != x_star.rows() || eigenvalues.size() != x_star.rows() || @@ -644,6 +751,26 @@ FixedEffectInferenceResult infer_fixed_effects_spectral( return result; } try { + if (fit.status == FitStatus::converged_boundary) { + result.method = FixedEffectTestMethod::ordinary_least_squares; + const OlsWorkspace workspace = make_ols_workspace(y_star, x_star); + result.coefficient_tests.reserve(x_star.cols()); + for (std::size_t coefficient = 0; coefficient < x_star.cols(); + ++coefficient) { + FixedEffectHypothesis hypothesis; + hypothesis.contrast = ColMajorMatrix(1, x_star.cols()); + hypothesis.contrast(0, coefficient) = 1.0; + result.coefficient_tests.push_back(ols_test( + hypothesis, workspace, rank_tolerance_relative)); + } + result.hypothesis_tests.reserve(hypotheses.size()); + for (const auto& hypothesis : hypotheses) { + result.hypothesis_tests.push_back(ols_test( + hypothesis, workspace, rank_tolerance_relative)); + } + result.status = FixedEffectInferenceStatus::boundary_ols; + return result; + } const InferenceWorkspace workspace = make_workspace( y_star, x_star, eigenvalues, fit, method, covariance_floor_relative); @@ -663,8 +790,8 @@ FixedEffectInferenceResult infer_fixed_effects_spectral( hypothesis, fit.beta, workspace, method, rank_tolerance_relative)); } - result.status = workspace.boundary - ? FixedEffectInferenceStatus::boundary_conditional + result.status = fit.status == FitStatus::line_search_failed + ? FixedEffectInferenceStatus::line_search_conditional : FixedEffectInferenceStatus::ok; } catch (const std::invalid_argument& exception) { result.status = FixedEffectInferenceStatus::invalid_contrast; diff --git a/src/reml_core.cpp b/src/reml_core.cpp index 7aa4ba4..449bee3 100644 --- a/src/reml_core.cpp +++ b/src/reml_core.cpp @@ -738,6 +738,8 @@ const char* to_string(FixedEffectTestMethod method) noexcept { switch (method) { case FixedEffectTestMethod::none: return "none"; + case FixedEffectTestMethod::ordinary_least_squares: + return "ordinary-least-squares"; case FixedEffectTestMethod::satterthwaite: return "satterthwaite"; case FixedEffectTestMethod::kenward_roger: @@ -752,8 +754,10 @@ const char* to_string(FixedEffectInferenceStatus status) noexcept { return "not_requested"; case FixedEffectInferenceStatus::ok: return "ok"; - case FixedEffectInferenceStatus::boundary_conditional: - return "boundary_conditional"; + case FixedEffectInferenceStatus::boundary_ols: + return "boundary_ols"; + case FixedEffectInferenceStatus::line_search_conditional: + return "line_search_conditional"; case FixedEffectInferenceStatus::fit_not_converged: return "fit_not_converged"; case FixedEffectInferenceStatus::invalid_contrast: diff --git a/tests/cpp/test_reml_synthetic.cpp b/tests/cpp/test_reml_synthetic.cpp index 4011803..9ddae34 100644 --- a/tests/cpp/test_reml_synthetic.cpp +++ b/tests/cpp/test_reml_synthetic.cpp @@ -195,6 +195,30 @@ void test_ai_reml_fit_improves_likelihood() { kr.hypothesis_tests.front().valid && kr.hypothesis_tests.front().numerator_df == 2, "interior KR joint test is invalid"); + + auto line_search_flagged = fitted; + line_search_flagged.status = + spectra::reml::FitStatus::line_search_failed; + line_search_flagged.error = "synthetic retained line-search iterate"; + const auto flagged_satterthwaite = + spectra::reml::infer_fixed_effects_spectral( + fixture.y, fixture.x, fixture.lambda, line_search_flagged, + spectra::reml::FixedEffectTestMethod::satterthwaite, {joint}); + require(flagged_satterthwaite.status == + spectra::reml::FixedEffectInferenceStatus::line_search_conditional, + "line-search Satterthwaite inference was not retained and flagged"); + require(flagged_satterthwaite.hypothesis_tests.size() == 1 && + flagged_satterthwaite.hypothesis_tests.front().valid, + "line-search Satterthwaite hypothesis test is invalid"); + const auto flagged_kr = spectra::reml::infer_fixed_effects_spectral( + fixture.y, fixture.x, fixture.lambda, line_search_flagged, + spectra::reml::FixedEffectTestMethod::kenward_roger, {joint}); + require(flagged_kr.status == + spectra::reml::FixedEffectInferenceStatus::line_search_conditional, + "line-search KR inference was not retained and flagged"); + require(flagged_kr.hypothesis_tests.size() == 1 && + flagged_kr.hypothesis_tests.front().valid, + "line-search KR hypothesis test is invalid"); } void test_probability_distributions_against_r() { @@ -372,8 +396,11 @@ void test_residual_only_kkt_boundary() { y, x, lambda, fitted, spectra::reml::FixedEffectTestMethod::satterthwaite); require(satterthwaite.status == - spectra::reml::FixedEffectInferenceStatus::boundary_conditional, - "Satterthwaite boundary inference did not report conditional status"); + spectra::reml::FixedEffectInferenceStatus::boundary_ols, + "Satterthwaite request at the boundary did not fall back to OLS"); + require(satterthwaite.method == + spectra::reml::FixedEffectTestMethod::ordinary_least_squares, + "boundary inference did not report ordinary least squares"); require(satterthwaite.coefficient_tests.size() == 1, "Satterthwaite coefficient test is missing"); const auto& satt = satterthwaite.coefficient_tests.front(); @@ -389,8 +416,11 @@ void test_residual_only_kkt_boundary() { y, x, lambda, fitted, spectra::reml::FixedEffectTestMethod::kenward_roger); require(kr.status == - spectra::reml::FixedEffectInferenceStatus::boundary_conditional, - "KR boundary inference did not report conditional status"); + spectra::reml::FixedEffectInferenceStatus::boundary_ols, + "KR request at the boundary did not fall back to OLS"); + require(kr.method == + spectra::reml::FixedEffectTestMethod::ordinary_least_squares, + "boundary KR request did not report ordinary least squares"); require(kr.coefficient_tests.size() == 1, "KR coefficient test is missing"); const auto& kr_test = kr.coefficient_tests.front(); @@ -399,8 +429,8 @@ void test_residual_only_kkt_boundary() { "KR did not recover OLS residual df"); require_near(kr_test.standard_error, expected_se, 2e-11, "KR did not recover OLS standard error"); - require_near(kr_test.statistic, expected_t * expected_t, 2e-10, - "KR did not recover OLS F statistic"); + require_near(kr_test.statistic, expected_t, 2e-10, + "boundary fallback did not recover OLS t statistic"); } } // namespace diff --git a/tests/python/test_spectra_reml_cli.py b/tests/python/test_spectra_reml_cli.py index 45023d1..1ab4a21 100644 --- a/tests/python/test_spectra_reml_cli.py +++ b/tests/python/test_spectra_reml_cli.py @@ -86,7 +86,7 @@ class SpectraRemlCliTests(unittest.TestCase): "0\ttrait_a\tconverged\t2\t0\t0\t0\t1\t1\t0.5\t-1\t4\t8\t1e-8" "\tsatterthwaite\tok\t0\t{}\t0\tnan\tnan\tnan\t\t\n" "1\ttrait_b\tconverged_boundary\t3\t1\t2\t3\t0\t1\t0\t-2\t3\t6\t1e-9" - "\tsatterthwaite\tboundary_conditional\t2\t{}\t1\t8\t2\t0.2\t\t\n".format( + "\tordinary-least-squares\tboundary_ols\t2\t{}\t1\t8\t2\t0.2\t\t\n".format( first_contrast_offset, second_contrast_offset ), encoding="utf-8",