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

Add Satterthwaite and Kenward-Roger fixed-effect tests

Browse Source
This commit is contained in:
Jiarui Li
2026-08-13 07:02:10 +08:00
parent 76dd8b1379
commit 8557daecac
14 changed files with 1330 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
CMakeLists.txt
View File

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

8
README.md
View File

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

59
docs/FORMAT.md
View File

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

13
include/spectra_reml/distributions.hpp Normal file
View File

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

27
include/spectra_reml/fixed_effects.hpp Normal file
View File

@@ -0,0 +1,27 @@
#pragma once
#include "spectra_reml/types.hpp"
#include <vector>
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<double> 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<double>& y_star, const ColMajorMatrix& x_star,
const std::vector<double>& eigenvalues, const RemlResult& fit,
FixedEffectTestMethod method,
const std::vector<FixedEffectHypothesis>& hypotheses = {},
double rank_tolerance_relative = 1e-10,
double covariance_floor_relative = 1e-12);
} // namespace spectra::reml

51
include/spectra_reml/types.hpp
View File

@@ -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<double>::quiet_NaN();
double statistic = std::numeric_limits<double>::quiet_NaN();
double p_value = std::numeric_limits<double>::quiet_NaN();
double estimate = std::numeric_limits<double>::quiet_NaN();
double standard_error = std::numeric_limits<double>::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<FixedEffectTestResult> coefficient_tests;
// Optional general linear hypotheses requested by the caller.
std::vector<FixedEffectTestResult> 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<double>::quiet_NaN();
double initial_sigma_g = std::numeric_limits<double>::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<double> beta_covariance_packed_lower;
FixedEffectInferenceResult fixed_effect_inference;
FixedEffectTestResult extra_fixed_effect_joint_test;
std::string error;
[[nodiscard]] bool has_estimates() const noexcept {

64
python/spectra_reml.py
View File

@@ -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="<f8")
cov = np.fromfile(cov_path, dtype="<f8")
fixed_se = np.fromfile(fixed_se_path, dtype="<f8")
fixed_stat = np.fromfile(fixed_stat_path, dtype="<f8")
fixed_ddf = np.fromfile(fixed_ddf_path, dtype="<f8")
fixed_p = np.fromfile(fixed_p_path, dtype="<f8")
next_beta_offset = 0
next_cov_offset = 0
next_fixed_test_offset = 0
for row in rows:
index = int(row["task_index"])
if index in results or index < 0 or index >= 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)

88
src/batch_io.cpp
View File

@@ -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<std::filesystem::path, 4> final_paths = {
summary_path, beta_path, covariance_path, complete_path};
const std::array<std::filesystem::path, 8> 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<double> beta_values;
std::vector<double> covariance_values;
std::vector<double> fixed_se_values;
std::vector<double> fixed_stat_values;
std::vector<double> fixed_ddf_values;
std::vector<double> fixed_p_values;
std::vector<std::int64_t> beta_offsets(results.size(), -1);
std::vector<std::int64_t> covariance_offsets(results.size(), -1);
std::vector<std::int64_t> 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<std::int64_t>(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<std::filesystem::path, 4> temporary_paths = {
summary_temp, beta_temp, covariance_temp, complete_temp};
const std::array<std::filesystem::path, 8> 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<FixedEffectHypothesis> 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();

121
src/distributions.cpp Normal file
View File

@@ -0,0 +1,121 @@
#include "spectra_reml/distributions.hpp"
#include <algorithm>
#include <cmath>
#include <limits>
#include <stdexcept>
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<double>::epsilon();
constexpr double kFloor = std::numeric_limits<double>::min() /
std::numeric_limits<double>::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<double>(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

685
src/fixed_effects.cpp Normal file
View File

@@ -0,0 +1,685 @@
#include "spectra_reml/fixed_effects.hpp"
#include "spectra_reml/distributions.hpp"
#include "spectra_reml/linalg.hpp"
#include <algorithm>
#include <array>
#include <cmath>
#include <limits>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>
namespace spectra::reml {
namespace {
using Matrix = std::vector<double>;
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<double>& 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<double>& 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<Matrix> covariance_gradient;
std::vector<Matrix> kr_p;
Matrix satterthwaite_parameter_covariance;
Matrix variance_parameter_covariance;
Matrix kr_adjusted_covariance;
};
std::vector<double> apply_projection(
const ColMajorMatrix& x, const std::vector<double>& inverse_v,
const Matrix& covariance, const std::vector<double>& input) {
const std::size_t n = x.rows();
const std::size_t p = x.cols();
std::vector<double> 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<double> 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<double> 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<double>& y,
const ColMajorMatrix& x,
const std::vector<double>& 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<double>::min());
std::vector<double> 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<std::vector<double>> derivative(k,
std::vector<double>(workspace.n, 1.0));
if (k == 2) {
derivative[1] = eigenvalues;
}
std::vector<Matrix> qq(k * k);
Matrix information_core(k * k, 0.0);
for (std::size_t a = 0; a < k; ++a) {
std::vector<double> 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<double> 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<std::vector<double>> derivative_py(
k, std::vector<double>(workspace.n));
std::vector<std::vector<double>> 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<double> 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<double> rhs;
std::vector<double> 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<double> 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<double> 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<double>::min(),
rank_tolerance_relative * maximum);
std::vector<std::size_t> 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<double>& contrast,
const Matrix& covariance, std::size_t p) {
return quadratic_form(contrast, covariance, p);
}
double satterthwaite_df(const std::vector<double>& contrast,
const InferenceWorkspace& workspace,
double variance) {
std::vector<double> 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<double>& 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<double>(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<double>(values.size()));
}
FixedEffectTestResult satterthwaite_test(
const FixedEffectHypothesis& hypothesis, const std::vector<double>& 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<double> component_df;
for (std::size_t row = 0; row < reduced.rank; ++row) {
std::vector<double> 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<double>(reduced.rank);
result.p_value = f_upper_tail(
result.statistic, static_cast<double>(reduced.rank),
result.denominator_df);
}
result.valid = true;
return result;
}
FixedEffectTestResult kenward_roger_test(
const FixedEffectHypothesis& hypothesis, const std::vector<double>& 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<double> 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<Matrix> 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<double>(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<double>& 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<double>& y_star, const ColMajorMatrix& x_star,
const std::vector<double>& eigenvalues, const RemlResult& fit,
FixedEffectTestMethod method,
const std::vector<FixedEffectHypothesis>& 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

19
src/main.cpp
View File

@@ -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");
});

32
src/reml_core.cpp
View File

@@ -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<double>& y_star, const ColMajorMatrix& x_star,
const std::vector<double>& eigenvalues, double sigma_e, double sigma_g,

159
tests/cpp/test_reml_synthetic.cpp
View File

@@ -1,4 +1,6 @@
#include "spectra_reml/reml.hpp"
#include "spectra_reml/fixed_effects.hpp"
#include "spectra_reml/distributions.hpp"
#include <algorithm>
#include <cmath>
@@ -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<double> expected_satt_df = {
11.271857730561006, 11.012661088553992, 8.7818755201562855};
const std::vector<double> 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<double> expected_kr_df = {
14.488430939129156, 14.905537974410061, 14.020384661469432};
const std::vector<double> 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<double> lambda(n);
std::vector<double> y(n, 0.0);
std::vector<double> y(n, 1.5);
for (std::size_t i = 0; i < n; ++i) {
x(i, 0) = 1.0;
lambda[i] = 0.1 * static_cast<double>(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<double>(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<double>(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<double>(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";

22
tests/python/test_spectra_reml_cli.py
View File

@@ -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="<f8").tofile(root / "block_000000.beta.f64.bin")
np.arange(1, 10, dtype="<f8").tofile(root / "block_000000.cov.f64.bin")
np.asarray([0.1, 0.2, 0.3, 0.4, 0.5], dtype="<f8").tofile(
root / "block_000000.fixed_se.f64.bin"
)
np.asarray([1, 2, 3, 4, 5], dtype="<f8").tofile(
root / "block_000000.fixed_stat.f64.bin"
)
np.asarray([10, 10, 8, 8, 8], dtype="<f8").tofile(
root / "block_000000.fixed_ddf.f64.bin"
)
np.asarray([0.5, 0.2, 0.1, 0.05, 0.01], dtype="<f8").tofile(
root / "block_000000.fixed_p.f64.bin"
)
(root / "block_000000.complete").write_text(
"format\t{}\nblock\t0\ntasks\t2\nbeta_elements\t5\ncov_elements\t9\n".format(cli.BLOCK_FORMAT),
"format\t{}\nblock\t0\ntasks\t2\nbeta_elements\t5\ncov_elements\t9\n"
"fixed_test_elements\t5\n".format(cli.BLOCK_FORMAT),
encoding="utf-8",
)
@@ -127,6 +142,7 @@ class SpectraRemlCliTests(unittest.TestCase):
self.assertEqual([row["task_id"] for row in rows], ["trait_a", "trait_b"])
self.assertEqual(json.loads(rows[0]["beta_json"]), [1.0, 2.0])
self.assertEqual(len(json.loads(rows[1]["covariance_packed_lower_json"])), 6)
self.assertEqual(len(json.loads(rows[1]["fixed_effect_p_value_json"])), 3)
altered = dict(manifest)
altered["created_utc"] = "changed"