From 8557daecac5fc13e313de1839908e1e915b87fb5 Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Thu, 13 Aug 2026 07:02:10 +0800 Subject: [PATCH] Add Satterthwaite and Kenward-Roger fixed-effect tests --- CMakeLists.txt | 2 + README.md | 8 + docs/FORMAT.md | 59 ++- include/spectra_reml/distributions.hpp | 13 + include/spectra_reml/fixed_effects.hpp | 27 + include/spectra_reml/types.hpp | 51 ++ python/spectra_reml.py | 64 ++- src/batch_io.cpp | 88 +++- src/distributions.cpp | 121 +++++ src/fixed_effects.cpp | 685 +++++++++++++++++++++++++ src/main.cpp | 19 + src/reml_core.cpp | 32 ++ tests/cpp/test_reml_synthetic.cpp | 159 +++++- tests/python/test_spectra_reml_cli.py | 22 +- 14 files changed, 1330 insertions(+), 20 deletions(-) create mode 100644 include/spectra_reml/distributions.hpp create mode 100644 include/spectra_reml/fixed_effects.hpp create mode 100644 src/distributions.cpp create mode 100644 src/fixed_effects.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index f795b43..19fc34e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -54,7 +54,9 @@ endif() add_library(spectra_reml_core STATIC src/linalg.cpp src/grm.cpp + src/distributions.cpp src/reml_core.cpp + src/fixed_effects.cpp src/batch_io.cpp ) add_library(SpectraREML::core ALIAS spectra_reml_core) diff --git a/README.md b/README.md index 68508ad..3759e3d 100644 --- a/README.md +++ b/README.md @@ -25,6 +25,8 @@ The engine diagonalizes the GRM once, rotates the common design and every unique - 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. +- Satterthwaite t/F fixed-effect inference by default, with optional + Kenward-Roger F tests and a joint test of task-specific covariates. - Atomic block output and safe resume/force semantics. - Generic Python CLI for manifest creation, validation, execution, provenance signatures, status, and result export. @@ -139,6 +141,12 @@ python python/spectra_reml.py finalize \ `results.tsv.gz` retains the complete per-task summary and stores the fixed-effect vector and row-wise packed lower covariance as JSON arrays. +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. + ## Recovery and provenance Each block is written as four files, with `.complete` renamed last. The Python layer adds `run.signature.json`, which binds the canonical manifest, engine SHA-256, numerical options, thread settings, and block size. diff --git a/docs/FORMAT.md b/docs/FORMAT.md index 0fe356f..92d1786 100644 --- a/docs/FORMAT.md +++ b/docs/FORMAT.md @@ -92,6 +92,10 @@ For block number `KKKKKK`: block_KKKKKK.summary.tsv block_KKKKKK.beta.f64.bin block_KKKKKK.cov.f64.bin +block_KKKKKK.fixed_se.f64.bin +block_KKKKKK.fixed_stat.f64.bin +block_KKKKKK.fixed_ddf.f64.bin +block_KKKKKK.fixed_p.f64.bin block_KKKKKK.complete ``` @@ -112,10 +116,22 @@ logL iterations line_search_steps grad_inf +fixed_test_method +fixed_test_status +fixed_test_offset +extra_joint_num_df +extra_joint_den_df +extra_joint_f +extra_joint_p +fixed_test_error error ``` -`beta_offset` and `cov_offset` count `float64` elements, not bytes. A negative offset indicates that no estimates were emitted for that task. +`beta_offset`, `cov_offset`, and `fixed_test_offset` count `float64` +elements, not bytes. A negative beta/covariance offset means that no estimates +were emitted. A negative fixed-test offset means that no valid coefficient-wise +fixed-effect tests were emitted; successful REML estimates are retained even +when inference fails. The covariance array uses the row-wise packed lower triangle: @@ -126,11 +142,12 @@ 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-v1 +format spectra-reml-block-v2 block 0 tasks 256 beta_elements 4096 cov_elements 34816 +fixed_test_elements 4096 ``` Consumers must ignore blocks without `.complete`. @@ -150,6 +167,35 @@ numerical_error `converged_boundary` is a successful residual-only solution accepted after the one-sided variance-component score and likelihood checks. +## Fixed-effect inference + +`--fixed-effect-test` selects `satterthwaite` (the default), +`kenward-roger`, or `none`. + +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. + +`fixed_test_status` is one of: + +```text +not_requested +ok +boundary_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. + 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 @@ -168,6 +214,13 @@ The Python CLI exports one TSV row per task. It includes the full summary plus: ```text beta_json covariance_packed_lower_json +fixed_effect_se_json +fixed_effect_statistic_json +fixed_effect_denominator_df_json +fixed_effect_p_value_json ``` -Project-specific software can attach coefficient names and derive contrasts without changing the numerical engine. +Project-specific software can attach coefficient names. The public C++ +inference API also accepts general linear hypotheses `L beta = rhs`; the batch +format currently emits coefficient-wise tests and the joint extra-covariate +test. diff --git a/include/spectra_reml/distributions.hpp b/include/spectra_reml/distributions.hpp new file mode 100644 index 0000000..ac82029 --- /dev/null +++ b/include/spectra_reml/distributions.hpp @@ -0,0 +1,13 @@ +#pragma once + +namespace spectra::reml { + +// Two-sided Student-t tail probability and upper F tail probability. +// Both accept non-integer positive degrees of freedom. +[[nodiscard]] double student_t_two_sided_p(double statistic, + double degrees_of_freedom); +[[nodiscard]] double f_upper_tail(double statistic, + double numerator_degrees_of_freedom, + double denominator_degrees_of_freedom); + +} // namespace spectra::reml diff --git a/include/spectra_reml/fixed_effects.hpp b/include/spectra_reml/fixed_effects.hpp new file mode 100644 index 0000000..c0cdbe1 --- /dev/null +++ b/include/spectra_reml/fixed_effects.hpp @@ -0,0 +1,27 @@ +#pragma once + +#include "spectra_reml/types.hpp" + +#include + +namespace spectra::reml { + +struct FixedEffectHypothesis { + // Rows are restrictions and columns correspond to beta in design order. + ColMajorMatrix contrast; + // Empty means a zero right-hand side. + 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. +[[nodiscard]] FixedEffectInferenceResult infer_fixed_effects_spectral( + const std::vector& y_star, const ColMajorMatrix& x_star, + const std::vector& eigenvalues, const RemlResult& fit, + FixedEffectTestMethod method, + const std::vector& hypotheses = {}, + double rank_tolerance_relative = 1e-10, + double covariance_floor_relative = 1e-12); + +} // namespace spectra::reml diff --git a/include/spectra_reml/types.hpp b/include/spectra_reml/types.hpp index 64ec3bb..788dbef 100644 --- a/include/spectra_reml/types.hpp +++ b/include/spectra_reml/types.hpp @@ -64,6 +64,49 @@ enum class FitStatus { [[nodiscard]] const char* to_string(FitStatus status) noexcept; +enum class FixedEffectTestMethod { + none, + satterthwaite, + kenward_roger +}; + +[[nodiscard]] const char* to_string(FixedEffectTestMethod method) noexcept; + +enum class FixedEffectInferenceStatus { + not_requested, + ok, + boundary_conditional, + fit_not_converged, + invalid_contrast, + information_singular, + numerical_error +}; + +[[nodiscard]] const char* to_string(FixedEffectInferenceStatus status) noexcept; + +struct FixedEffectTestResult { + bool valid = false; + std::size_t numerator_df = 0; + double denominator_df = std::numeric_limits::quiet_NaN(); + double statistic = std::numeric_limits::quiet_NaN(); + double p_value = std::numeric_limits::quiet_NaN(); + double estimate = std::numeric_limits::quiet_NaN(); + double standard_error = std::numeric_limits::quiet_NaN(); + std::string error; +}; + +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. + std::vector coefficient_tests; + // Optional general linear hypotheses requested by the caller. + std::vector hypothesis_tests; + std::string error; +}; + struct RemlOptions { std::size_t max_iterations = 100; std::size_t line_search_max_evaluations = 48; @@ -95,6 +138,12 @@ struct RemlOptions { // sqrt(OLS residual mean square / 2). double initial_sigma_e = std::numeric_limits::quiet_NaN(); double initial_sigma_g = std::numeric_limits::quiet_NaN(); + + // Batch fixed-effect inference. The numerical fit API itself remains + // usable without inference; run_task_batch invokes the inference engine at + // the final REML estimate and defaults to Satterthwaite t/F tests. + FixedEffectTestMethod fixed_effect_test = + FixedEffectTestMethod::satterthwaite; }; struct RemlResult { @@ -113,6 +162,8 @@ struct RemlResult { // Row-wise packed lower triangle: // (0,0), (1,0), (1,1), (2,0), (2,1), (2,2), ... std::vector beta_covariance_packed_lower; + FixedEffectInferenceResult fixed_effect_inference; + FixedEffectTestResult extra_fixed_effect_joint_test; std::string error; [[nodiscard]] bool has_estimates() const noexcept { diff --git a/python/spectra_reml.py b/python/spectra_reml.py index 05048f1..ccb9cc0 100644 --- a/python/spectra_reml.py +++ b/python/spectra_reml.py @@ -32,7 +32,7 @@ except ImportError as exc: # pragma: no cover MANIFEST_FORMAT = "spectra-reml-manifest-v1" -BLOCK_FORMAT = "spectra-reml-block-v1" +BLOCK_FORMAT = "spectra-reml-block-v2" RUN_SIGNATURE_FORMAT = "spectra-reml-run-signature-v1" FINALIZE_FORMAT = "spectra-reml-finalize-v1" TASK_HEADER = ( @@ -56,6 +56,14 @@ SUMMARY_HEADER = ( "iterations", "line_search_steps", "grad_inf", + "fixed_test_method", + "fixed_test_status", + "fixed_test_offset", + "extra_joint_num_df", + "extra_joint_den_df", + "extra_joint_f", + "extra_joint_p", + "fixed_test_error", "error", ) SUCCESS_STATUSES = frozenset(("converged", "converged_boundary")) @@ -63,6 +71,10 @@ BLOCK_PATTERNS = ( "block_*.summary.tsv", "block_*.beta.f64.bin", "block_*.cov.f64.bin", + "block_*.fixed_se.f64.bin", + "block_*.fixed_stat.f64.bin", + "block_*.fixed_ddf.f64.bin", + "block_*.fixed_p.f64.bin", "block_*.complete", ) @@ -361,6 +373,7 @@ def build_engine_command(engine: Path, manifest: Mapping[str, Any], args: argpar "--n-extra-covariate-rows", str(dims["extra_covariate_row_count"]), "--block-size", str(args.block_size), "--threads", str(args.threads), + "--fixed-effect-test", args.fixed_effect_test, ] if paths["grm_id"] is not None: command.extend(("--grm-id", str(paths["grm_id"]))) @@ -554,19 +567,28 @@ def read_completed(manifest: Mapping[str, Any]) -> list[dict[str, Any]]: summary_path = output / (stem + ".summary.tsv") beta_path = output / (stem + ".beta.f64.bin") cov_path = output / (stem + ".cov.f64.bin") + fixed_se_path = output / (stem + ".fixed_se.f64.bin") + fixed_stat_path = output / (stem + ".fixed_stat.f64.bin") + fixed_ddf_path = output / (stem + ".fixed_ddf.f64.bin") + fixed_p_path = output / (stem + ".fixed_p.f64.bin") marker_values = read_marker(marker) try: declared_tasks = int(marker_values["tasks"]) beta_elements = int(marker_values["beta_elements"]) cov_elements = int(marker_values["cov_elements"]) + fixed_test_elements = int(marker_values["fixed_test_elements"]) except (KeyError, ValueError) as exc: fail("Malformed completion counts in {}: {}".format(marker, exc)) if marker_values.get("format") != BLOCK_FORMAT or min( - declared_tasks, beta_elements, cov_elements + declared_tasks, beta_elements, cov_elements, fixed_test_elements ) < 0: fail("Invalid completion marker: {}".format(marker)) require_size(beta_path, beta_elements * 8, "block beta") require_size(cov_path, cov_elements * 8, "block covariance") + require_size(fixed_se_path, fixed_test_elements * 8, "fixed-effect SE") + require_size(fixed_stat_path, fixed_test_elements * 8, "fixed-effect statistic") + require_size(fixed_ddf_path, fixed_test_elements * 8, "fixed-effect denominator df") + require_size(fixed_p_path, fixed_test_elements * 8, "fixed-effect p-value") with summary_path.open("r", encoding="utf-8-sig", newline="") as handle: reader = csv.DictReader(handle, delimiter="\t") if tuple(reader.fieldnames or ()) != SUMMARY_HEADER: @@ -576,14 +598,20 @@ def read_completed(manifest: Mapping[str, Any]) -> list[dict[str, Any]]: fail("Summary row count disagrees with marker: {}".format(marker)) beta = np.fromfile(beta_path, dtype="= task_count: fail("Duplicate/out-of-range task index: {}".format(index)) p = int(row["n_fixed"]) beta_offset, cov_offset = int(row["beta_offset"]), int(row["cov_offset"]) + fixed_test_offset = int(row["fixed_test_offset"]) if beta_offset < 0 or cov_offset < 0: if beta_offset != cov_offset: fail("Only one output offset is negative for task {}".format(index)) @@ -599,11 +627,30 @@ def read_completed(manifest: Mapping[str, Any]) -> list[dict[str, Any]]: cov_values = cov[cov_offset:cov_offset + packed].tolist() next_beta_offset += p next_cov_offset += packed + if fixed_test_offset < 0: + fixed_se_values: list[float] = [] + fixed_stat_values: list[float] = [] + fixed_ddf_values: list[float] = [] + fixed_p_values: list[float] = [] + else: + if (fixed_test_offset != next_fixed_test_offset or + fixed_test_offset + p > fixed_se.size): + fail("Invalid fixed-effect test offset for task {}".format(index)) + fixed_se_values = fixed_se[fixed_test_offset:fixed_test_offset + p].tolist() + fixed_stat_values = fixed_stat[fixed_test_offset:fixed_test_offset + p].tolist() + fixed_ddf_values = fixed_ddf[fixed_test_offset:fixed_test_offset + p].tolist() + fixed_p_values = fixed_p[fixed_test_offset:fixed_test_offset + p].tolist() + next_fixed_test_offset += p result: dict[str, Any] = dict(row) result["beta_json"] = json.dumps(beta_values, separators=(",", ":")) result["covariance_packed_lower_json"] = json.dumps(cov_values, separators=(",", ":")) + result["fixed_effect_se_json"] = json.dumps(fixed_se_values, separators=(",", ":")) + result["fixed_effect_statistic_json"] = json.dumps(fixed_stat_values, separators=(",", ":")) + result["fixed_effect_denominator_df_json"] = json.dumps(fixed_ddf_values, separators=(",", ":")) + result["fixed_effect_p_value_json"] = json.dumps(fixed_p_values, separators=(",", ":")) results[index] = result - if next_beta_offset != beta_elements or next_cov_offset != cov_elements: + if (next_beta_offset != beta_elements or next_cov_offset != cov_elements or + next_fixed_test_offset != fixed_test_elements): fail("Block binary arrays contain unused elements: {}".format(stem)) missing = sorted(set(range(task_count)).difference(results)) if missing: @@ -618,7 +665,11 @@ def finalize(manifest: Mapping[str, Any], output: Path) -> Path: output.parent.mkdir(parents=True, exist_ok=True) temporary = output.with_name(output.name + ".tmp.{}".format(os.getpid())) opener = gzip.open if output.suffix == ".gz" else open - fields = list(SUMMARY_HEADER) + ["beta_json", "covariance_packed_lower_json"] + fields = list(SUMMARY_HEADER) + [ + "beta_json", "covariance_packed_lower_json", "fixed_effect_se_json", + "fixed_effect_statistic_json", "fixed_effect_denominator_df_json", + "fixed_effect_p_value_json", + ] try: with opener(temporary, "wt", encoding="utf-8", newline="") as handle: writer = csv.DictWriter(handle, fieldnames=fields, delimiter="\t", lineterminator="\n") @@ -675,6 +726,11 @@ def add_run_options(parser: argparse.ArgumentParser) -> None: parser.add_argument("--engine", default="spectra_reml") parser.add_argument("--threads", type=positive_integer, default=1) parser.add_argument("--blas-threads", type=positive_integer, default=1) + parser.add_argument( + "--fixed-effect-test", + choices=("satterthwaite", "kenward-roger", "none"), + default="satterthwaite", + ) parser.add_argument("--block-size", type=positive_integer, default=256) parser.add_argument("--max-iterations", type=positive_integer, default=100) parser.add_argument("--line-search-max-evals", type=positive_integer, default=48) diff --git a/src/batch_io.cpp b/src/batch_io.cpp index f4fae28..110c98d 100644 --- a/src/batch_io.cpp +++ b/src/batch_io.cpp @@ -1,5 +1,7 @@ #include "spectra_reml/batch_io.hpp" +#include "spectra_reml/fixed_effects.hpp" + #include "spectra_reml/linalg.hpp" #include "spectra_reml/reml.hpp" @@ -32,7 +34,7 @@ namespace spectra::reml { namespace { -constexpr const char* kOutputFormatVersion = "spectra-reml-block-v1"; +constexpr const char* kOutputFormatVersion = "spectra-reml-block-v2"; std::size_t checked_product(std::size_t left, std::size_t right, const char* description) { @@ -288,9 +290,15 @@ void write_block_atomic(const std::filesystem::path& output_directory, const auto summary_path = output_directory / (stem + ".summary.tsv"); const auto beta_path = output_directory / (stem + ".beta.f64.bin"); const auto covariance_path = output_directory / (stem + ".cov.f64.bin"); + const auto fixed_se_path = output_directory / (stem + ".fixed_se.f64.bin"); + const auto fixed_stat_path = + output_directory / (stem + ".fixed_stat.f64.bin"); + const auto fixed_ddf_path = output_directory / (stem + ".fixed_ddf.f64.bin"); + const auto fixed_p_path = output_directory / (stem + ".fixed_p.f64.bin"); const auto complete_path = output_directory / (stem + ".complete"); - const std::array final_paths = { - summary_path, beta_path, covariance_path, complete_path}; + const std::array final_paths = { + summary_path, beta_path, covariance_path, fixed_se_path, + fixed_stat_path, fixed_ddf_path, fixed_p_path, complete_path}; for (const auto& path : final_paths) { if (std::filesystem::exists(path)) { if (!overwrite) { @@ -302,8 +310,13 @@ void write_block_atomic(const std::filesystem::path& output_directory, std::vector beta_values; std::vector covariance_values; + std::vector fixed_se_values; + std::vector fixed_stat_values; + std::vector fixed_ddf_values; + std::vector fixed_p_values; std::vector beta_offsets(results.size(), -1); std::vector covariance_offsets(results.size(), -1); + std::vector fixed_test_offsets(results.size(), -1); for (std::size_t index = 0; index < results.size(); ++index) { const auto& fit = results[index].fit; if (!fit.has_estimates()) { @@ -316,17 +329,40 @@ void write_block_atomic(const std::filesystem::path& output_directory, covariance_values.insert(covariance_values.end(), fit.beta_covariance_packed_lower.begin(), fit.beta_covariance_packed_lower.end()); + const auto& inference = fit.fixed_effect_inference; + if ((inference.status == FixedEffectInferenceStatus::ok || + inference.status == + FixedEffectInferenceStatus::boundary_conditional) && + inference.coefficient_tests.size() == fit.beta.size()) { + fixed_test_offsets[index] = + static_cast(fixed_se_values.size()); + for (const auto& test : inference.coefficient_tests) { + fixed_se_values.push_back(test.standard_error); + fixed_stat_values.push_back(test.statistic); + fixed_ddf_values.push_back(test.denominator_df); + fixed_p_values.push_back(test.p_value); + } + } } const auto summary_temp = temporary_path(summary_path); const auto beta_temp = temporary_path(beta_path); const auto covariance_temp = temporary_path(covariance_path); + const auto fixed_se_temp = temporary_path(fixed_se_path); + const auto fixed_stat_temp = temporary_path(fixed_stat_path); + const auto fixed_ddf_temp = temporary_path(fixed_ddf_path); + const auto fixed_p_temp = temporary_path(fixed_p_path); const auto complete_temp = temporary_path(complete_path); - const std::array temporary_paths = { - summary_temp, beta_temp, covariance_temp, complete_temp}; + const std::array temporary_paths = { + summary_temp, beta_temp, covariance_temp, fixed_se_temp, + fixed_stat_temp, fixed_ddf_temp, fixed_p_temp, complete_temp}; try { write_binary_doubles(beta_temp, beta_values); write_binary_doubles(covariance_temp, covariance_values); + write_binary_doubles(fixed_se_temp, fixed_se_values); + write_binary_doubles(fixed_stat_temp, fixed_stat_values); + write_binary_doubles(fixed_ddf_temp, fixed_ddf_values); + write_binary_doubles(fixed_p_temp, fixed_p_values); { std::ofstream summary(summary_temp, std::ios::trunc); if (!summary) { @@ -336,7 +372,10 @@ void write_block_atomic(const std::filesystem::path& output_directory, summary << "task_index\ttask_id\tstatus\tn_fixed\tn_extra_covariates\tbeta_offset" "\tcov_offset\tsigma_g2\tsigma_e2\th2\tlogL\titerations" - "\tline_search_steps\tgrad_inf\terror\n"; + "\tline_search_steps\tgrad_inf\tfixed_test_method" + "\tfixed_test_status\tfixed_test_offset" + "\textra_joint_num_df\textra_joint_den_df" + "\textra_joint_f\textra_joint_p\tfixed_test_error\terror\n"; summary << std::setprecision(17); for (std::size_t index = 0; index < results.size(); ++index) { const auto& item = results[index]; @@ -354,6 +393,14 @@ void write_block_atomic(const std::filesystem::path& output_directory, << fit.h2 << '\t' << fit.log_likelihood << '\t' << fit.iterations << '\t' << fit.line_search_evaluations << '\t' << fit.gradient_inf_norm << '\t' + << to_string(fit.fixed_effect_inference.method) << '\t' + << to_string(fit.fixed_effect_inference.status) << '\t' + << fixed_test_offsets[index] << '\t' + << fit.extra_fixed_effect_joint_test.numerator_df << '\t' + << fit.extra_fixed_effect_joint_test.denominator_df << '\t' + << fit.extra_fixed_effect_joint_test.statistic << '\t' + << fit.extra_fixed_effect_joint_test.p_value << '\t' + << sanitize_tsv(fit.fixed_effect_inference.error) << '\t' << sanitize_tsv(fit.error) << '\n'; } summary.flush(); @@ -372,7 +419,9 @@ void write_block_atomic(const std::filesystem::path& output_directory, << "block\t" << block_index << '\n' << "tasks\t" << results.size() << '\n' << "beta_elements\t" << beta_values.size() << '\n' - << "cov_elements\t" << covariance_values.size() << '\n'; + << "cov_elements\t" << covariance_values.size() << '\n' + << "fixed_test_elements\t" << fixed_se_values.size() + << '\n'; complete.flush(); if (!complete) { throw std::runtime_error("failed writing completion marker " + @@ -381,6 +430,10 @@ void write_block_atomic(const std::filesystem::path& output_directory, } rename_checked(beta_temp, beta_path); rename_checked(covariance_temp, covariance_path); + rename_checked(fixed_se_temp, fixed_se_path); + rename_checked(fixed_stat_temp, fixed_stat_path); + rename_checked(fixed_ddf_temp, fixed_ddf_path); + rename_checked(fixed_p_temp, fixed_p_path); rename_checked(summary_temp, summary_path); rename_checked(complete_temp, complete_path); } catch (...) { @@ -696,6 +749,27 @@ void run_task_batch(const BatchInputPaths& paths, (local + 1) * dimensions.sample_count); item.fit = fit_ai_reml_spectral( phenotype, design, spectral.eigenvalues, options.reml); + std::vector hypotheses; + if (!task.extra_covariate_rows.empty()) { + FixedEffectHypothesis extra_joint; + extra_joint.contrast = ColMajorMatrix( + task.extra_covariate_rows.size(), fixed_count); + for (std::size_t extra = 0; + extra < task.extra_covariate_rows.size(); ++extra) { + extra_joint.contrast( + extra, dimensions.base_covariate_count + extra) = 1.0; + } + hypotheses.push_back(std::move(extra_joint)); + } + item.fit.fixed_effect_inference = infer_fixed_effects_spectral( + phenotype, design, spectral.eigenvalues, item.fit, + options.reml.fixed_effect_test, hypotheses, + options.reml.rank_tolerance_relative, + options.reml.covariance_floor_relative); + if (!item.fit.fixed_effect_inference.hypothesis_tests.empty()) { + item.fit.extra_fixed_effect_joint_test = + item.fit.fixed_effect_inference.hypothesis_tests.front(); + } } catch (const std::exception& exception) { item.fit.status = FitStatus::numerical_error; item.fit.error = exception.what(); diff --git a/src/distributions.cpp b/src/distributions.cpp new file mode 100644 index 0000000..83f9c4d --- /dev/null +++ b/src/distributions.cpp @@ -0,0 +1,121 @@ +#include "spectra_reml/distributions.hpp" + +#include +#include +#include +#include + +namespace spectra::reml { +namespace { + +double beta_continued_fraction(double a, double b, double x) { + constexpr int kMaximumIterations = 512; + constexpr double kTolerance = 8.0 * std::numeric_limits::epsilon(); + constexpr double kFloor = std::numeric_limits::min() / + std::numeric_limits::epsilon(); + const double qab = a + b; + const double qap = a + 1.0; + const double qam = a - 1.0; + double c = 1.0; + double d = 1.0 - qab * x / qap; + if (std::abs(d) < kFloor) { + d = kFloor; + } + d = 1.0 / d; + double result = d; + for (int iteration = 1; iteration <= kMaximumIterations; ++iteration) { + const double m = static_cast(iteration); + const double m2 = 2.0 * m; + double aa = m * (b - m) * x / ((qam + m2) * (a + m2)); + d = 1.0 + aa * d; + if (std::abs(d) < kFloor) { + d = kFloor; + } + c = 1.0 + aa / c; + if (std::abs(c) < kFloor) { + c = kFloor; + } + d = 1.0 / d; + result *= d * c; + + aa = -(a + m) * (qab + m) * x / + ((a + m2) * (qap + m2)); + d = 1.0 + aa * d; + if (std::abs(d) < kFloor) { + d = kFloor; + } + c = 1.0 + aa / c; + if (std::abs(c) < kFloor) { + c = kFloor; + } + d = 1.0 / d; + const double delta = d * c; + result *= delta; + if (std::abs(delta - 1.0) <= kTolerance) { + return result; + } + } + throw std::runtime_error("incomplete beta continued fraction did not converge"); +} + +double regularized_incomplete_beta(double x, double a, double b) { + if (!(a > 0.0) || !(b > 0.0) || !std::isfinite(a) || + !std::isfinite(b) || !(x >= 0.0 && x <= 1.0)) { + throw std::invalid_argument("invalid incomplete beta arguments"); + } + if (x == 0.0) { + return 0.0; + } + if (x == 1.0) { + return 1.0; + } + const double log_factor = std::lgamma(a + b) - std::lgamma(a) - + std::lgamma(b) + a * std::log(x) + + b * std::log1p(-x); + const double factor = std::exp(log_factor); + double result = 0.0; + if (x < (a + 1.0) / (a + b + 2.0)) { + result = factor * beta_continued_fraction(a, b, x) / a; + } else { + result = 1.0 - + factor * beta_continued_fraction(b, a, 1.0 - x) / b; + } + return std::clamp(result, 0.0, 1.0); +} + +} // namespace + +double student_t_two_sided_p(double statistic, double degrees_of_freedom) { + if (!(degrees_of_freedom > 0.0) || + !std::isfinite(degrees_of_freedom) || std::isnan(statistic)) { + throw std::invalid_argument("invalid Student-t arguments"); + } + if (std::isinf(statistic)) { + return 0.0; + } + const double square = statistic * statistic; + const double x = degrees_of_freedom / (degrees_of_freedom + square); + return regularized_incomplete_beta(x, 0.5 * degrees_of_freedom, 0.5); +} + +double f_upper_tail(double statistic, double numerator_degrees_of_freedom, + double denominator_degrees_of_freedom) { + if (!(statistic >= 0.0) || std::isnan(statistic) || + !(numerator_degrees_of_freedom > 0.0) || + !(denominator_degrees_of_freedom > 0.0) || + !std::isfinite(numerator_degrees_of_freedom) || + !std::isfinite(denominator_degrees_of_freedom)) { + throw std::invalid_argument("invalid F-distribution arguments"); + } + if (std::isinf(statistic)) { + return 0.0; + } + const double x = denominator_degrees_of_freedom / + (denominator_degrees_of_freedom + + numerator_degrees_of_freedom * statistic); + return regularized_incomplete_beta( + x, 0.5 * denominator_degrees_of_freedom, + 0.5 * numerator_degrees_of_freedom); +} + +} // namespace spectra::reml diff --git a/src/fixed_effects.cpp b/src/fixed_effects.cpp new file mode 100644 index 0000000..8aa3632 --- /dev/null +++ b/src/fixed_effects.cpp @@ -0,0 +1,685 @@ +#include "spectra_reml/fixed_effects.hpp" + +#include "spectra_reml/distributions.hpp" +#include "spectra_reml/linalg.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace spectra::reml { +namespace { + +using Matrix = std::vector; + +double& at(Matrix& matrix, std::size_t order, std::size_t row, + std::size_t col) { + return matrix[row + col * order]; +} + +double at(const Matrix& matrix, std::size_t order, std::size_t row, + std::size_t col) { + return matrix[row + col * order]; +} + +Matrix multiply(const Matrix& left, const Matrix& right, std::size_t order) { + Matrix result(order * order, 0.0); + for (std::size_t col = 0; col < order; ++col) { + for (std::size_t inner = 0; inner < order; ++inner) { + const double right_value = at(right, order, inner, col); + for (std::size_t row = 0; row < order; ++row) { + at(result, order, row, col) += + at(left, order, row, inner) * right_value; + } + } + } + return result; +} + +Matrix transpose(const Matrix& input, std::size_t order) { + Matrix result(order * order, 0.0); + for (std::size_t col = 0; col < order; ++col) { + for (std::size_t row = 0; row < order; ++row) { + at(result, order, row, col) = at(input, order, col, row); + } + } + return result; +} + +double trace(const Matrix& matrix, std::size_t order) { + double result = 0.0; + for (std::size_t i = 0; i < order; ++i) { + result += at(matrix, order, i, i); + } + return result; +} + +double element_sum_product(const Matrix& left, const Matrix& right) { + double result = 0.0; + for (std::size_t i = 0; i < left.size(); ++i) { + result += left[i] * right[i]; + } + return result; +} + +Matrix weighted_cross_product(const ColMajorMatrix& x, + const std::vector& weights) { + const std::size_t p = x.cols(); + Matrix result(p * p, 0.0); + for (std::size_t i = 0; i < x.rows(); ++i) { + for (std::size_t row = 0; row < p; ++row) { + for (std::size_t col = 0; col <= row; ++col) { + at(result, p, row, col) += + weights[i] * x(i, row) * x(i, col); + } + } + } + for (std::size_t col = 0; col < p; ++col) { + for (std::size_t row = 0; row < col; ++row) { + at(result, p, row, col) = at(result, p, col, row); + } + } + return result; +} + +double quadratic_form(const std::vector& vector, + const Matrix& matrix, std::size_t order) { + double result = 0.0; + for (std::size_t col = 0; col < order; ++col) { + for (std::size_t row = 0; row < order; ++row) { + result += vector[row] * at(matrix, order, row, col) * vector[col]; + } + } + return result; +} + +bool invert_positive_definite(const Matrix& matrix, std::size_t order, + Matrix& inverse, std::string& error) { + Matrix factor = matrix; + if (!cholesky_factor_in_place(factor, order, &error)) { + return false; + } + return cholesky_inverse(factor, order, inverse, &error); +} + +struct InferenceWorkspace { + std::size_t n = 0; + std::size_t p = 0; + std::size_t parameter_count = 0; + bool boundary = false; + Matrix covariance; + std::vector covariance_gradient; + std::vector kr_p; + Matrix satterthwaite_parameter_covariance; + Matrix variance_parameter_covariance; + Matrix kr_adjusted_covariance; +}; + +std::vector apply_projection( + const ColMajorMatrix& x, const std::vector& inverse_v, + const Matrix& covariance, const std::vector& input) { + const std::size_t n = x.rows(); + const std::size_t p = x.cols(); + std::vector rhs(p, 0.0); + for (std::size_t col = 0; col < p; ++col) { + for (std::size_t i = 0; i < n; ++i) { + rhs[col] += x(i, col) * inverse_v[i] * input[i]; + } + } + std::vector coefficients(p, 0.0); + for (std::size_t row = 0; row < p; ++row) { + for (std::size_t col = 0; col < p; ++col) { + coefficients[row] += at(covariance, p, row, col) * rhs[col]; + } + } + std::vector result(n); + for (std::size_t i = 0; i < n; ++i) { + double fitted = 0.0; + for (std::size_t col = 0; col < p; ++col) { + fitted += x(i, col) * coefficients[col]; + } + result[i] = inverse_v[i] * (input[i] - fitted); + } + return result; +} + +InferenceWorkspace make_workspace(const std::vector& y, + const ColMajorMatrix& x, + const std::vector& eigenvalues, + const RemlResult& fit, + FixedEffectTestMethod method, + double covariance_floor_relative) { + InferenceWorkspace workspace; + workspace.n = x.rows(); + workspace.p = x.cols(); + workspace.boundary = fit.status == FitStatus::converged_boundary; + workspace.parameter_count = workspace.boundary ? 1 : 2; + const std::size_t k = workspace.parameter_count; + const std::size_t p = workspace.p; + + const double ve = fit.sigma_e2; + const double vg = fit.sigma_g2; + double covariance_scale = ve; + for (double lambda : eigenvalues) { + covariance_scale = std::max(covariance_scale, std::abs(vg * lambda)); + } + const double floor = covariance_floor_relative * + std::max(covariance_scale, + std::numeric_limits::min()); + std::vector inverse_v(workspace.n); + for (std::size_t i = 0; i < workspace.n; ++i) { + const double value = ve + vg * eigenvalues[i]; + if (!std::isfinite(value) || value <= floor) { + throw std::runtime_error("non-positive covariance in fixed-effect inference"); + } + inverse_v[i] = 1.0 / value; + } + + const Matrix normal = weighted_cross_product(x, inverse_v); + std::string error; + if (!invert_positive_definite(normal, p, workspace.covariance, error)) { + throw std::runtime_error("fixed-effect covariance failed: " + error); + } + + workspace.kr_p.resize(k); + workspace.covariance_gradient.resize(k); + std::vector> derivative(k, + std::vector(workspace.n, 1.0)); + if (k == 2) { + derivative[1] = eigenvalues; + } + std::vector qq(k * k); + Matrix information_core(k * k, 0.0); + for (std::size_t a = 0; a < k; ++a) { + std::vector weights(workspace.n); + for (std::size_t i = 0; i < workspace.n; ++i) { + weights[i] = derivative[a][i] * inverse_v[i] * inverse_v[i]; + } + Matrix positive = weighted_cross_product(x, weights); + workspace.kr_p[a] = positive; + for (double& value : workspace.kr_p[a]) { + value = -value; + } + workspace.covariance_gradient[a] = multiply( + multiply(workspace.covariance, positive, p), + workspace.covariance, p); + } + if (method == FixedEffectTestMethod::satterthwaite) { + std::vector py(workspace.n); + for (std::size_t i = 0; i < workspace.n; ++i) { + double fitted = 0.0; + for (std::size_t col = 0; col < p; ++col) { + fitted += x(i, col) * fit.beta[col]; + } + py[i] = inverse_v[i] * (y[i] - fitted); + } + std::vector> derivative_py( + k, std::vector(workspace.n)); + std::vector> p_derivative_py(k); + for (std::size_t a = 0; a < k; ++a) { + for (std::size_t i = 0; i < workspace.n; ++i) { + derivative_py[a][i] = derivative[a][i] * py[i]; + } + p_derivative_py[a] = apply_projection( + x, inverse_v, workspace.covariance, derivative_py[a]); + } + Matrix average_information(k * k, 0.0); + for (std::size_t a = 0; a < k; ++a) { + for (std::size_t b = a; b < k; ++b) { + const double value = + 0.5 * dot(derivative_py[a], p_derivative_py[b]); + at(average_information, k, a, b) = value; + at(average_information, k, b, a) = value; + } + } + if (!invert_positive_definite( + average_information, k, + workspace.satterthwaite_parameter_covariance, error)) { + throw std::runtime_error( + "variance-parameter average information is singular: " + error); + } + return workspace; + } + + for (std::size_t a = 0; a < k; ++a) { + for (std::size_t b = a; b < k; ++b) { + std::vector weights(workspace.n); + double ktrace = 0.0; + for (std::size_t i = 0; i < workspace.n; ++i) { + const double derivative_product = derivative[a][i] * derivative[b][i]; + const double inverse2 = inverse_v[i] * inverse_v[i]; + weights[i] = derivative_product * inverse2 * inverse_v[i]; + ktrace += derivative_product * inverse2; + } + qq[a + b * k] = weighted_cross_product(x, weights); + qq[b + a * k] = qq[a + b * k]; + const Matrix covariance_p_a = + multiply(workspace.covariance, workspace.kr_p[a], p); + const Matrix p_b_covariance = + multiply(workspace.kr_p[b], workspace.covariance, p); + const double value = + ktrace - + 2.0 * element_sum_product(workspace.covariance, + qq[a + b * k]) + + element_sum_product(covariance_p_a, p_b_covariance); + at(information_core, k, a, b) = value; + at(information_core, k, b, a) = value; + } + } + + Matrix inverse_information_core; + if (!invert_positive_definite(information_core, k, + inverse_information_core, error)) { + throw std::runtime_error("variance-parameter information is singular: " + error); + } + workspace.variance_parameter_covariance = inverse_information_core; + for (double& value : workspace.variance_parameter_covariance) { + value *= 2.0; + } + + Matrix kr_u(p * p, 0.0); + for (std::size_t a = 0; a < k; ++a) { + for (std::size_t b = 0; b < k; ++b) { + const Matrix term = multiply( + multiply(workspace.kr_p[a], workspace.covariance, p), + workspace.kr_p[b], p); + const double weight = + at(workspace.variance_parameter_covariance, k, a, b); + for (std::size_t index = 0; index < kr_u.size(); ++index) { + kr_u[index] += weight * (qq[a + b * k][index] - term[index]); + } + } + } + const Matrix gamma = multiply( + multiply(workspace.covariance, kr_u, p), workspace.covariance, p); + workspace.kr_adjusted_covariance = workspace.covariance; + for (std::size_t index = 0; index < workspace.kr_adjusted_covariance.size(); + ++index) { + workspace.kr_adjusted_covariance[index] += 2.0 * gamma[index]; + } + return workspace; +} + +struct ReducedHypothesis { + std::size_t rank = 0; + ColMajorMatrix contrast; + std::vector rhs; + std::vector eigenvalues; +}; + +ReducedHypothesis reduce_hypothesis(const FixedEffectHypothesis& hypothesis, + const Matrix& covariance, std::size_t p, + double rank_tolerance_relative) { + if (hypothesis.contrast.rows() == 0 || hypothesis.contrast.cols() != p) { + throw std::invalid_argument("contrast has invalid dimensions"); + } + const std::size_t rows = hypothesis.contrast.rows(); + std::vector rhs = hypothesis.rhs; + if (rhs.empty()) { + rhs.assign(rows, 0.0); + } + if (rhs.size() != rows) { + throw std::invalid_argument("contrast right-hand side has invalid length"); + } + + ColMajorMatrix contrast_covariance(rows, rows); + for (std::size_t col = 0; col < rows; ++col) { + for (std::size_t row = 0; row < rows; ++row) { + double value = 0.0; + for (std::size_t a = 0; a < p; ++a) { + for (std::size_t b = 0; b < p; ++b) { + value += hypothesis.contrast(row, a) * + at(covariance, p, a, b) * + hypothesis.contrast(col, b); + } + } + contrast_covariance(row, col) = value; + } + } + const std::vector eigenvalues = + symmetric_eigen_decomposition(contrast_covariance); + if (rows == 1) { + if (!(eigenvalues[0] > 0.0) || !std::isfinite(eigenvalues[0])) { + throw std::invalid_argument("contrast has zero numerical rank"); + } + ReducedHypothesis result; + result.rank = 1; + result.contrast = hypothesis.contrast; + result.rhs = std::move(rhs); + result.eigenvalues = eigenvalues; + return result; + } + const double maximum = eigenvalues.empty() ? 0.0 : eigenvalues.back(); + const double tolerance = + std::max(std::numeric_limits::min(), + rank_tolerance_relative * maximum); + std::vector retained; + for (std::size_t i = 0; i < rows; ++i) { + if (eigenvalues[i] > tolerance) { + retained.push_back(i); + } + } + if (retained.empty()) { + throw std::invalid_argument("contrast has zero numerical rank"); + } + + ReducedHypothesis result; + result.rank = retained.size(); + result.contrast = ColMajorMatrix(result.rank, p); + result.rhs.resize(result.rank, 0.0); + result.eigenvalues.resize(result.rank); + for (std::size_t reduced_row = 0; reduced_row < retained.size(); + ++reduced_row) { + const std::size_t eigen_index = retained[reduced_row]; + result.eigenvalues[reduced_row] = eigenvalues[eigen_index]; + for (std::size_t original_row = 0; original_row < rows; ++original_row) { + const double coefficient = + contrast_covariance(original_row, eigen_index); + result.rhs[reduced_row] += coefficient * rhs[original_row]; + for (std::size_t col = 0; col < p; ++col) { + result.contrast(reduced_row, col) += + coefficient * hypothesis.contrast(original_row, col); + } + } + } + return result; +} + +double contrast_variance(const std::vector& contrast, + const Matrix& covariance, std::size_t p) { + return quadratic_form(contrast, covariance, p); +} + +double satterthwaite_df(const std::vector& contrast, + const InferenceWorkspace& workspace, + double variance) { + std::vector gradient(workspace.parameter_count, 0.0); + for (std::size_t a = 0; a < workspace.parameter_count; ++a) { + gradient[a] = contrast_variance( + contrast, workspace.covariance_gradient[a], workspace.p); + } + const double denominator = quadratic_form( + gradient, workspace.satterthwaite_parameter_covariance, + workspace.parameter_count); + if (!(variance > 0.0) || !(denominator > 0.0) || + !std::isfinite(denominator)) { + throw std::runtime_error("invalid Satterthwaite denominator"); + } + const double result = 2.0 * variance * variance / denominator; + if (!(result > 0.0) || !std::isfinite(result)) { + throw std::runtime_error("invalid Satterthwaite degrees of freedom"); + } + return result; +} + +double combine_satterthwaite_df(const std::vector& values) { + if (values.size() == 1) { + return values.front(); + } + const auto [minimum, maximum] = + std::minmax_element(values.begin(), values.end()); + if (*maximum - *minimum < 1e-8) { + double sum = 0.0; + for (double value : values) { + sum += value; + } + return sum / static_cast(values.size()); + } + for (double value : values) { + if (value <= 2.0) { + return 2.0; + } + } + double expectation = 0.0; + for (double value : values) { + expectation += value / (value - 2.0); + } + return 2.0 * expectation / + (expectation - static_cast(values.size())); +} + +FixedEffectTestResult satterthwaite_test( + const FixedEffectHypothesis& hypothesis, const std::vector& beta, + const InferenceWorkspace& workspace, double rank_tolerance_relative) { + const ReducedHypothesis reduced = reduce_hypothesis( + hypothesis, workspace.covariance, workspace.p, + rank_tolerance_relative); + FixedEffectTestResult result; + result.numerator_df = reduced.rank; + double sum_t_squared = 0.0; + std::vector component_df; + for (std::size_t row = 0; row < reduced.rank; ++row) { + std::vector contrast(workspace.p); + double estimate = -reduced.rhs[row]; + for (std::size_t col = 0; col < workspace.p; ++col) { + contrast[col] = reduced.contrast(row, col); + estimate += contrast[col] * beta[col]; + } + const double variance = reduced.eigenvalues[row]; + sum_t_squared += estimate * estimate / variance; + component_df.push_back( + satterthwaite_df(contrast, workspace, variance)); + if (reduced.rank == 1) { + result.estimate = estimate + reduced.rhs[row]; + result.standard_error = std::sqrt(variance); + result.statistic = estimate / result.standard_error; + } + } + result.denominator_df = combine_satterthwaite_df(component_df); + 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); + } + result.valid = true; + return result; +} + +FixedEffectTestResult kenward_roger_test( + const FixedEffectHypothesis& hypothesis, const std::vector& beta, + const InferenceWorkspace& workspace, double rank_tolerance_relative) { + const ReducedHypothesis reduced = reduce_hypothesis( + hypothesis, workspace.covariance, workspace.p, + rank_tolerance_relative); + const std::size_t q = reduced.rank; + const std::size_t p = workspace.p; + const std::size_t k = workspace.parameter_count; + + Matrix l_phi_lt(q * q, 0.0); + Matrix l_phia_lt(q * q, 0.0); + std::vector difference(q, 0.0); + for (std::size_t row = 0; row < q; ++row) { + difference[row] = -reduced.rhs[row]; + for (std::size_t col = 0; col < p; ++col) { + difference[row] += reduced.contrast(row, col) * beta[col]; + } + for (std::size_t other = 0; other < q; ++other) { + for (std::size_t a = 0; a < p; ++a) { + for (std::size_t b = 0; b < p; ++b) { + const double product = reduced.contrast(row, a) * + reduced.contrast(other, b); + at(l_phi_lt, q, row, other) += + product * at(workspace.covariance, p, a, b); + at(l_phia_lt, q, row, other) += + product * at(workspace.kr_adjusted_covariance, p, a, b); + } + } + } + } + std::string error; + Matrix inverse_l_phi_lt; + Matrix inverse_l_phia_lt; + if (!invert_positive_definite(l_phi_lt, q, inverse_l_phi_lt, error) || + !invert_positive_definite(l_phia_lt, q, inverse_l_phia_lt, error)) { + throw std::runtime_error("KR contrast covariance failed: " + error); + } + + // Theta = L' (L Phi L')^-1 L. + Matrix theta(p * p, 0.0); + for (std::size_t col = 0; col < p; ++col) { + for (std::size_t row = 0; row < p; ++row) { + for (std::size_t a = 0; a < q; ++a) { + for (std::size_t b = 0; b < q; ++b) { + at(theta, p, row, col) += + reduced.contrast(a, row) * + at(inverse_l_phi_lt, q, a, b) * + reduced.contrast(b, col); + } + } + } + } + const Matrix theta_phi = multiply(theta, workspace.covariance, p); + std::vector u(k); + for (std::size_t a = 0; a < k; ++a) { + u[a] = multiply( + multiply(theta_phi, workspace.kr_p[a], p), + workspace.covariance, p); + } + double a1 = 0.0; + double a2 = 0.0; + for (std::size_t a = 0; a < k; ++a) { + for (std::size_t b = 0; b < k; ++b) { + const double weight = + at(workspace.variance_parameter_covariance, k, a, b); + a1 += weight * trace(u[a], p) * trace(u[b], p); + a2 += weight * element_sum_product(u[a], transpose(u[b], p)); + } + } + const double qd = static_cast(q); + if (!(a2 > 0.0) || !std::isfinite(a1) || !std::isfinite(a2)) { + throw std::runtime_error("invalid Kenward-Roger moment adjustment"); + } + const double b = (a1 + 6.0 * a2) / (2.0 * qd); + const double g = ((qd + 1.0) * a1 - (qd + 4.0) * a2) / + ((qd + 2.0) * a2); + const double common = 3.0 * qd + 2.0 * (1.0 - g); + const double c1 = g / common; + const double c2 = (qd - g) / common; + const double c3 = (qd + 2.0 - g) / common; + const double v0 = 1.0 + c1 * b; + const double v1 = 1.0 - c2 * b; + const double v2 = 1.0 - c3 * b; + const double rho = (1.0 / qd) * + std::pow((1.0 - a2 / qd) / v1, 2.0) * v0 / v2; + const double denominator_df = + 4.0 + (qd + 2.0) / (qd * rho - 1.0); + if (!(denominator_df > 0.0) || !std::isfinite(denominator_df)) { + throw std::runtime_error("invalid Kenward-Roger degrees of freedom"); + } + const double scaling = + std::abs(denominator_df - 2.0) < 1e-2 + ? 1.0 + : denominator_df * (1.0 - a2 / qd) / + (denominator_df - 2.0); + const double unscaled_f = + quadratic_form(difference, inverse_l_phia_lt, q) / qd; + + FixedEffectTestResult result; + result.valid = true; + result.numerator_df = q; + result.denominator_df = denominator_df; + result.statistic = scaling * unscaled_f; + result.p_value = + f_upper_tail(result.statistic, qd, result.denominator_df); + if (q == 1) { + double raw_estimate = 0.0; + for (std::size_t col = 0; col < p; ++col) { + raw_estimate += reduced.contrast(0, col) * beta[col]; + } + result.estimate = raw_estimate; + result.standard_error = std::sqrt(at(l_phia_lt, 1, 0, 0)); + } + return result; +} + +FixedEffectTestResult run_test(const FixedEffectHypothesis& hypothesis, + const std::vector& beta, + const InferenceWorkspace& workspace, + FixedEffectTestMethod method, + double rank_tolerance_relative) { + if (method == FixedEffectTestMethod::satterthwaite) { + return satterthwaite_test(hypothesis, beta, workspace, + rank_tolerance_relative); + } + if (method == FixedEffectTestMethod::kenward_roger) { + return kenward_roger_test(hypothesis, beta, workspace, + rank_tolerance_relative); + } + throw std::invalid_argument("fixed-effect inference method is none"); +} + +} // namespace + +FixedEffectInferenceResult infer_fixed_effects_spectral( + const std::vector& y_star, const ColMajorMatrix& x_star, + const std::vector& eigenvalues, const RemlResult& fit, + FixedEffectTestMethod method, + const std::vector& hypotheses, + double rank_tolerance_relative, double covariance_floor_relative) { + FixedEffectInferenceResult result; + result.method = method; + if (method == FixedEffectTestMethod::none) { + result.status = FixedEffectInferenceStatus::not_requested; + return result; + } + if (fit.status != FitStatus::converged && + fit.status != FitStatus::converged_boundary) { + result.status = FixedEffectInferenceStatus::fit_not_converged; + result.error = "fixed-effect tests require a converged REML fit"; + return result; + } + if (y_star.size() != x_star.rows() || eigenvalues.size() != x_star.rows() || + fit.beta.size() != x_star.cols()) { + result.status = FixedEffectInferenceStatus::numerical_error; + result.error = "fixed-effect inference inputs have inconsistent dimensions"; + return result; + } + try { + const InferenceWorkspace workspace = make_workspace( + y_star, x_star, eigenvalues, fit, method, + covariance_floor_relative); + 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(run_test( + hypothesis, fit.beta, workspace, method, + rank_tolerance_relative)); + } + result.hypothesis_tests.reserve(hypotheses.size()); + for (const auto& hypothesis : hypotheses) { + result.hypothesis_tests.push_back(run_test( + hypothesis, fit.beta, workspace, method, + rank_tolerance_relative)); + } + result.status = workspace.boundary + ? FixedEffectInferenceStatus::boundary_conditional + : FixedEffectInferenceStatus::ok; + } catch (const std::invalid_argument& exception) { + result.status = FixedEffectInferenceStatus::invalid_contrast; + result.error = exception.what(); + } catch (const std::runtime_error& exception) { + const std::string message = exception.what(); + result.status = message.find("information is singular") != std::string::npos + ? FixedEffectInferenceStatus::information_singular + : FixedEffectInferenceStatus::numerical_error; + result.error = message; + } catch (const std::exception& exception) { + result.status = FixedEffectInferenceStatus::numerical_error; + result.error = exception.what(); + } + return result; +} + +} // namespace spectra::reml diff --git a/src/main.cpp b/src/main.cpp index 0c2328b..988dadf 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -17,6 +17,7 @@ namespace { using spectra::reml::BatchDimensions; using spectra::reml::BatchInputPaths; using spectra::reml::BatchOptions; +using spectra::reml::FixedEffectTestMethod; void print_usage(std::ostream& output) { output @@ -44,6 +45,7 @@ void print_usage(std::ostream& output) { << " --resume skip blocks with .complete marker\n" << " --overwrite replace existing block files\n\n" << "Optional AI-REML controls:\n" + << " --fixed-effect-test METHOD satterthwaite (default), kenward-roger, or none\n" << " --max-iterations N default 100\n" << " --line-search-max-evals N default 48\n" << " --line-search-max-zoom N default 48\n" @@ -179,6 +181,7 @@ void reject_unknown_options(const Arguments& arguments) { "--n-extra-covariate-rows", "--block-size", "--threads", + "--fixed-effect-test", "--max-iterations", "--line-search-max-evals", "--line-search-max-zoom", @@ -245,6 +248,22 @@ int main(int argc, char** argv) { BatchOptions options; options.resume = arguments.flags.count("--resume") != 0U; options.overwrite = arguments.flags.count("--overwrite") != 0U; + set_if_present(arguments, "--fixed-effect-test", + [&](const std::string& value) { + if (value == "satterthwaite") { + options.reml.fixed_effect_test = + FixedEffectTestMethod::satterthwaite; + } else if (value == "kenward-roger") { + options.reml.fixed_effect_test = + FixedEffectTestMethod::kenward_roger; + } else if (value == "none") { + options.reml.fixed_effect_test = + FixedEffectTestMethod::none; + } else { + throw std::invalid_argument( + "--fixed-effect-test must be satterthwaite, kenward-roger, or none"); + } + }); set_if_present(arguments, "--block-size", [&](const std::string& value) { options.block_size = parse_size(value, "--block-size"); }); diff --git a/src/reml_core.cpp b/src/reml_core.cpp index 9ef1973..7aa4ba4 100644 --- a/src/reml_core.cpp +++ b/src/reml_core.cpp @@ -734,6 +734,38 @@ const char* to_string(FitStatus status) noexcept { return "unknown"; } +const char* to_string(FixedEffectTestMethod method) noexcept { + switch (method) { + case FixedEffectTestMethod::none: + return "none"; + case FixedEffectTestMethod::satterthwaite: + return "satterthwaite"; + case FixedEffectTestMethod::kenward_roger: + return "kenward-roger"; + } + return "unknown"; +} + +const char* to_string(FixedEffectInferenceStatus status) noexcept { + switch (status) { + case FixedEffectInferenceStatus::not_requested: + return "not_requested"; + case FixedEffectInferenceStatus::ok: + return "ok"; + case FixedEffectInferenceStatus::boundary_conditional: + return "boundary_conditional"; + case FixedEffectInferenceStatus::fit_not_converged: + return "fit_not_converged"; + case FixedEffectInferenceStatus::invalid_contrast: + return "invalid_contrast"; + case FixedEffectInferenceStatus::information_singular: + return "information_singular"; + case FixedEffectInferenceStatus::numerical_error: + return "numerical_error"; + } + return "unknown"; +} + RemlEvaluation evaluate_reml_spectral( const std::vector& y_star, const ColMajorMatrix& x_star, const std::vector& eigenvalues, double sigma_e, double sigma_g, diff --git a/tests/cpp/test_reml_synthetic.cpp b/tests/cpp/test_reml_synthetic.cpp index bc57f54..4011803 100644 --- a/tests/cpp/test_reml_synthetic.cpp +++ b/tests/cpp/test_reml_synthetic.cpp @@ -1,4 +1,6 @@ #include "spectra_reml/reml.hpp" +#include "spectra_reml/fixed_effects.hpp" +#include "spectra_reml/distributions.hpp" #include #include @@ -163,6 +165,119 @@ void test_ai_reml_fit_improves_likelihood() { "reported sigma_e2 is inconsistent"); require_near(fitted.sigma_g2, fitted.sigma_g * fitted.sigma_g, 1e-14, "reported sigma_g2 is inconsistent"); + + spectra::reml::FixedEffectHypothesis joint; + joint.contrast = spectra::reml::ColMajorMatrix(2, fixture.x.cols()); + joint.contrast(0, 1) = 1.0; + joint.contrast(1, 2) = 1.0; + const auto satterthwaite = spectra::reml::infer_fixed_effects_spectral( + fixture.y, fixture.x, fixture.lambda, fitted, + spectra::reml::FixedEffectTestMethod::satterthwaite, {joint}); + require(satterthwaite.status == + spectra::reml::FixedEffectInferenceStatus::ok, + "interior Satterthwaite inference failed: " + satterthwaite.error); + require(satterthwaite.coefficient_tests.size() == fixture.x.cols(), + "interior Satterthwaite coefficient tests are missing"); + require(satterthwaite.hypothesis_tests.size() == 1 && + satterthwaite.hypothesis_tests.front().numerator_df == 2, + "interior Satterthwaite joint test is missing"); + require(satterthwaite.hypothesis_tests.front().valid && + satterthwaite.hypothesis_tests.front().p_value >= 0.0 && + satterthwaite.hypothesis_tests.front().p_value <= 1.0, + "interior Satterthwaite joint p-value is invalid"); + + const auto kr = spectra::reml::infer_fixed_effects_spectral( + fixture.y, fixture.x, fixture.lambda, fitted, + spectra::reml::FixedEffectTestMethod::kenward_roger, {joint}); + require(kr.status == spectra::reml::FixedEffectInferenceStatus::ok, + "interior KR inference failed: " + kr.error); + require(kr.hypothesis_tests.size() == 1 && + kr.hypothesis_tests.front().valid && + kr.hypothesis_tests.front().numerator_df == 2, + "interior KR joint test is invalid"); +} + +void test_probability_distributions_against_r() { + require_near(spectra::reml::student_t_two_sided_p(2.1, 7.3), + 0.072246713424853351, 2e-14, + "Student-t tail probability disagrees with R"); + require_near(spectra::reml::f_upper_tail(3.7, 2.5, 11.2), + 0.051114185450108554, 2e-14, + "F tail probability disagrees with R"); + require_near(spectra::reml::student_t_two_sided_p(12.0, 3.5), + 0.0005745341928561279, 2e-14, + "extreme Student-t tail probability disagrees with R"); + require_near(spectra::reml::f_upper_tail(100.0, 1.0, 8.25), + 6.803231179451616e-06, 2e-14, + "extreme F tail probability disagrees with R"); +} + +void test_fixed_effect_inference_against_dense_reference() { + 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, false, true); + require(evaluated.valid, "fixed-effect reference evaluation failed"); + spectra::reml::RemlResult fit; + fit.status = spectra::reml::FitStatus::converged; + fit.sigma_e = sigma_e; + fit.sigma_g = sigma_g; + fit.sigma_e2 = sigma_e * sigma_e; + fit.sigma_g2 = sigma_g * sigma_g; + fit.beta = evaluated.beta; + + spectra::reml::FixedEffectHypothesis joint; + joint.contrast = spectra::reml::ColMajorMatrix(2, fixture.x.cols()); + joint.contrast(0, 1) = 1.0; + joint.contrast(1, 2) = 1.0; + const auto satt = spectra::reml::infer_fixed_effects_spectral( + fixture.y, fixture.x, fixture.lambda, fit, + spectra::reml::FixedEffectTestMethod::satterthwaite, {joint}); + require(satt.status == spectra::reml::FixedEffectInferenceStatus::ok, + "dense-reference Satterthwaite inference failed: " + satt.error); + const std::vector expected_satt_df = { + 11.271857730561006, 11.012661088553992, 8.7818755201562855}; + const std::vector expected_satt_t = { + 1.9994269321713118, -1.6229945279594145, 0.26284041305853334}; + for (std::size_t i = 0; i < fixture.x.cols(); ++i) { + require_near(satt.coefficient_tests[i].denominator_df, + expected_satt_df[i], 3e-11, + "Satterthwaite df disagrees with dense reference"); + require_near(satt.coefficient_tests[i].statistic, + expected_satt_t[i], 3e-12, + "Satterthwaite t disagrees with dense reference"); + } + require_near(satt.hypothesis_tests[0].statistic, + 1.3172663846541173, 3e-12, + "Satterthwaite joint F disagrees with dense reference"); + require_near(satt.hypothesis_tests[0].denominator_df, + 9.7667759456518599, 3e-11, + "Satterthwaite joint df disagrees with dense reference"); + + const auto kr = spectra::reml::infer_fixed_effects_spectral( + fixture.y, fixture.x, fixture.lambda, fit, + spectra::reml::FixedEffectTestMethod::kenward_roger, {joint}); + require(kr.status == spectra::reml::FixedEffectInferenceStatus::ok, + "dense-reference KR inference failed: " + kr.error); + const std::vector expected_kr_df = { + 14.488430939129156, 14.905537974410061, 14.020384661469432}; + const std::vector expected_kr_f = { + 3.9170644186264694, 2.1360773364244796, 0.0544837858871276}; + for (std::size_t i = 0; i < fixture.x.cols(); ++i) { + require_near(kr.coefficient_tests[i].denominator_df, + expected_kr_df[i], 3e-11, + "KR df disagrees with dense reference"); + require_near(kr.coefficient_tests[i].statistic, + expected_kr_f[i], 3e-11, + "KR F disagrees with dense reference"); + } + require_near(kr.hypothesis_tests[0].statistic, + 1.0684514230854425, 3e-11, + "KR joint F disagrees with dense reference"); + require_near(kr.hypothesis_tests[0].denominator_df, + 14.615505880683108, 3e-11, + "KR joint df disagrees with dense reference"); } void test_internal_phenotype_scaling_restores_original_units() { @@ -212,15 +327,15 @@ void test_residual_only_kkt_boundary() { constexpr std::size_t n = 14; spectra::reml::ColMajorMatrix x(n, 1); std::vector lambda(n); - std::vector y(n, 0.0); + std::vector y(n, 1.5); for (std::size_t i = 0; i < n; ++i) { x(i, 0) = 1.0; lambda[i] = 0.1 * static_cast(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; + y[0] += 1.0; + y[1] -= 1.0; const double expected_sigma_e2 = 2.0 / static_cast(n - 1); const auto boundary = spectra::reml::evaluate_reml_spectral( y, x, lambda, std::sqrt(expected_sigma_e2), 0.0, false, true); @@ -250,6 +365,42 @@ void test_residual_only_kkt_boundary() { "residual-only variance is not RSS/(n-p)"); require_near(fitted.log_likelihood, boundary.log_likelihood, 2e-12, "reported boundary likelihood is inconsistent"); + + const double expected_se = std::sqrt(expected_sigma_e2 / n); + const double expected_t = 1.5 / expected_se; + const auto satterthwaite = spectra::reml::infer_fixed_effects_spectral( + y, x, lambda, fitted, + spectra::reml::FixedEffectTestMethod::satterthwaite); + require(satterthwaite.status == + spectra::reml::FixedEffectInferenceStatus::boundary_conditional, + "Satterthwaite boundary inference did not report conditional status"); + require(satterthwaite.coefficient_tests.size() == 1, + "Satterthwaite coefficient test is missing"); + const auto& satt = satterthwaite.coefficient_tests.front(); + require(satt.valid, "Satterthwaite boundary test is invalid"); + require_near(satt.denominator_df, static_cast(n - 1), 2e-11, + "Satterthwaite did not recover OLS residual df"); + require_near(satt.standard_error, expected_se, 2e-12, + "Satterthwaite did not recover OLS standard error"); + require_near(satt.statistic, expected_t, 2e-11, + "Satterthwaite did not recover OLS t statistic"); + + const auto kr = spectra::reml::infer_fixed_effects_spectral( + 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"); + require(kr.coefficient_tests.size() == 1, + "KR coefficient test is missing"); + const auto& kr_test = kr.coefficient_tests.front(); + require(kr_test.valid, "KR boundary test is invalid: " + kr.error); + require_near(kr_test.denominator_df, static_cast(n - 1), 2e-10, + "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"); } } // namespace @@ -259,6 +410,8 @@ int main() { test_score_matches_finite_difference(); test_signed_parameterization_and_row_permutation(); test_ai_reml_fit_improves_likelihood(); + test_probability_distributions_against_r(); + test_fixed_effect_inference_against_dense_reference(); test_internal_phenotype_scaling_restores_original_units(); test_residual_only_kkt_boundary(); std::cout << "test_reml_synthetic: PASS\n"; diff --git a/tests/python/test_spectra_reml_cli.py b/tests/python/test_spectra_reml_cli.py index 93e5526..e069e8f 100644 --- a/tests/python/test_spectra_reml_cli.py +++ b/tests/python/test_spectra_reml_cli.py @@ -64,14 +64,29 @@ class SpectraRemlCliTests(unittest.TestCase): root.mkdir(parents=True, exist_ok=True) (root / "block_000000.summary.tsv").write_text( "\t".join(cli.SUMMARY_HEADER) + "\n" - "0\ttrait_a\tconverged\t2\t0\t0\t0\t1\t1\t0.5\t-1\t4\t8\t1e-8\t\n" - "1\ttrait_b\tconverged_boundary\t3\t1\t2\t3\t0\t1\t0\t-2\t3\t6\t1e-9\t\n", + "0\ttrait_a\tconverged\t2\t0\t0\t0\t1\t1\t0.5\t-1\t4\t8\t1e-8" + "\tsatterthwaite\tok\t0\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\t1\t8\t2\t0.2\t\t\n", encoding="utf-8", ) np.asarray([1, 2, 3, 4, 5], dtype="