Add reusable SpectraREML batch AI-REML engine

This commit is contained in:
2026-08-09 14:13:38 +08:00
parent 777eeb288c
commit da971a2ca7
22 changed files with 4736 additions and 0 deletions

10
.gitignore vendored Normal file
View File

@@ -0,0 +1,10 @@
/build*/
/*.obj
/*.lib
/*.exp
/out/
/results/
__pycache__/
*.py[cod]
*.tmp.*
/examples/synthetic_work/

118
CMakeLists.txt Normal file
View File

@@ -0,0 +1,118 @@
cmake_minimum_required(VERSION 3.20)
project(
SpectraREML
VERSION 0.1.0
DESCRIPTION "Reusable batched spectral AI-REML engine"
LANGUAGES CXX
)
include(GNUInstallDirs)
include(CTest)
list(PREPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake")
set(REML_BLAS "MKL" CACHE STRING "BLAS/LAPACKE backend: MKL or OPENBLAS")
set_property(CACHE REML_BLAS PROPERTY STRINGS MKL OPENBLAS)
string(TOUPPER "${REML_BLAS}" REML_BLAS)
option(REML_ENABLE_OPENMP "Parallelize independent REML tasks with OpenMP" ON)
option(REML_ENABLE_NATIVE_ARCH "Enable compiler tuning for the build host" OFF)
add_library(reml_blas_lapacke INTERFACE)
add_library(REML::BLAS_LAPACKE ALIAS reml_blas_lapacke)
if(REML_BLAS STREQUAL "MKL")
# These variables are consumed by Intel's MKLConfig.cmake. Sequential is
# deliberately the portable default because tasks are parallelized with
# OpenMP; a threaded BLAS can otherwise create nested thread pools.
set(MKL_INTERFACE "lp64" CACHE STRING "oneMKL integer interface")
set(MKL_LINK "dynamic" CACHE STRING "oneMKL link type")
set(MKL_THREADING "sequential" CACHE STRING "oneMKL threading layer")
set_property(CACHE MKL_INTERFACE PROPERTY STRINGS lp64 ilp64)
set_property(CACHE MKL_LINK PROPERTY STRINGS dynamic static)
set_property(CACHE MKL_THREADING PROPERTY STRINGS sequential intel_thread gnu_thread tbb_thread)
find_package(MKL CONFIG REQUIRED)
if(NOT TARGET MKL::MKL)
message(FATAL_ERROR "MKL was found, but its MKL::MKL CMake target is missing")
endif()
target_link_libraries(reml_blas_lapacke INTERFACE MKL::MKL)
target_compile_definitions(reml_blas_lapacke INTERFACE REML_BLAS_MKL=1)
elseif(REML_BLAS STREQUAL "OPENBLAS")
find_package(REML_OpenBLAS REQUIRED)
target_link_libraries(reml_blas_lapacke INTERFACE REML::OpenBLASLAPACKE)
target_compile_definitions(reml_blas_lapacke INTERFACE REML_BLAS_OPENBLAS=1)
else()
message(FATAL_ERROR "Unsupported REML_BLAS='${REML_BLAS}'. Choose MKL or OPENBLAS.")
endif()
if(REML_ENABLE_OPENMP)
find_package(OpenMP REQUIRED COMPONENTS CXX)
endif()
add_library(spectra_reml_core STATIC
src/linalg.cpp
src/grm.cpp
src/reml_core.cpp
src/batch_io.cpp
)
add_library(SpectraREML::core ALIAS spectra_reml_core)
target_compile_features(spectra_reml_core PUBLIC cxx_std_17)
target_include_directories(spectra_reml_core
PUBLIC
$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>
$<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}>
)
target_link_libraries(spectra_reml_core PUBLIC REML::BLAS_LAPACKE)
if(REML_ENABLE_OPENMP)
target_link_libraries(spectra_reml_core PUBLIC OpenMP::OpenMP_CXX)
target_compile_definitions(spectra_reml_core PUBLIC REML_USE_OPENMP=1)
endif()
if(MSVC)
target_compile_options(spectra_reml_core PRIVATE /W4 /permissive- /EHsc)
else()
target_compile_options(spectra_reml_core PRIVATE -Wall -Wextra -Wpedantic)
target_compile_definitions(spectra_reml_core PUBLIC _FILE_OFFSET_BITS=64)
if(REML_ENABLE_NATIVE_ARCH)
target_compile_options(spectra_reml_core PRIVATE -march=native)
endif()
endif()
add_executable(spectra_reml src/main.cpp)
target_link_libraries(spectra_reml PRIVATE SpectraREML::core)
target_compile_features(spectra_reml PRIVATE cxx_std_17)
if(MSVC)
target_compile_options(spectra_reml PRIVATE /W4 /permissive- /EHsc)
else()
target_compile_options(spectra_reml PRIVATE -Wall -Wextra -Wpedantic)
if(REML_ENABLE_NATIVE_ARCH)
target_compile_options(spectra_reml PRIVATE -march=native)
endif()
endif()
if(BUILD_TESTING)
add_executable(test_reml_synthetic tests/cpp/test_reml_synthetic.cpp)
target_link_libraries(test_reml_synthetic PRIVATE SpectraREML::core)
add_test(NAME reml_synthetic COMMAND test_reml_synthetic)
add_executable(test_grm_io tests/cpp/test_grm_io.cpp)
target_link_libraries(test_grm_io PRIVATE SpectraREML::core)
add_test(NAME grm_io COMMAND test_grm_io)
endif()
install(TARGETS spectra_reml_core spectra_reml
RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR}
LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
)
install(DIRECTORY include/ DESTINATION ${CMAKE_INSTALL_INCLUDEDIR})
message(STATUS "SpectraREML ${PROJECT_VERSION}")
message(STATUS " BLAS/LAPACKE backend : ${REML_BLAS}")
message(STATUS " OpenMP : ${REML_ENABLE_OPENMP}")
message(STATUS " Native architecture : ${REML_ENABLE_NATIVE_ARCH}")

153
README.md
View File

@@ -1,2 +1,155 @@
# SpectraREML
SpectraREML is a reusable batch AI-REML engine for many continuous traits that share one genomic relationship matrix (GRM).
For each task, it fits
\[
y = X\beta + Zu + e, \qquad
u \sim N(0,\sigma_g^2G), \qquad
e \sim N(0,\sigma_e^2I).
\]
The engine diagonalizes the GRM once, rotates the common design and every unique task-specific covariate once, rotates phenotypes in blocks, and fits independent REML tasks in parallel.
## Features
- C++17 numerical core using oneMKL or OpenBLAS/LAPACKE.
- GRM eigendecomposition once per task set.
- AI-REML with a strong-Wolfe line search; no EM updates.
- 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.
- Atomic block output and safe resume/force semantics.
- Generic Python CLI for manifest creation, validation, execution, provenance signatures, status, and result export.
The numerical core has no knowledge of cohorts, molecular assay types, or domain-specific variable names.
## Repository layout
```text
include/spectra_reml/ public C++ API
src/ numerical core, batch I/O, and CLI
python/spectra_reml.py generic Python CLI
examples/example.py fully synthetic input example
scripts/run_server.sh generic Linux build/run wrapper
tests/ numerical and file-contract tests
docs/FORMAT.md binary and tabular file contract
```
## Build with Intel oneMKL
Load the oneAPI environment first:
```bash
source /opt/intel/oneapi/setvars.sh
cmake -S . -B build-mkl \
-DCMAKE_BUILD_TYPE=Release \
-DREML_BLAS=MKL \
-DMKL_INTERFACE=lp64 \
-DMKL_LINK=dynamic \
-DMKL_THREADING=sequential
cmake --build build-mkl --parallel
ctest --test-dir build-mkl --output-on-failure
```
For a nonstandard installation, locate `MKLConfig.cmake` and set its directory explicitly:
```bash
cmake -S . -B build-mkl \
-DREML_BLAS=MKL \
-DMKL_DIR=/path/to/mkl/latest/lib/cmake/mkl \
-DMKL_THREADING=sequential
```
Sequential BLAS is recommended because independent tasks are already parallelized by OpenMP.
## Build with OpenBLAS
```bash
cmake -S . -B build-openblas \
-DCMAKE_BUILD_TYPE=Release \
-DREML_BLAS=OPENBLAS \
-DOpenBLAS_ROOT=/path/to/openblas \
-DLAPACKE_ROOT=/path/to/lapacke
cmake --build build-openblas --parallel
ctest --test-dir build-openblas --output-on-failure
```
The configuration performs a real CBLAS/LAPACKE link check. It supports LAPACKE either inside OpenBLAS or in a separate library.
## Quick start
Create a synthetic bundle:
```bash
python examples/example.py
```
Create a manifest for your own data:
```bash
python python/spectra_reml.py make-manifest \
--manifest work/manifest.json \
--grm-bin data/example.grm.bin \
--grm-id data/example.grm.id \
--base-x data/base_x.f64.bin \
--phenotypes data/phenotypes.f64.bin \
--extra-covariates data/extra_covariates.f32.bin \
--tasks data/tasks.tsv \
--extra-offsets data/extra_offsets.i64.bin \
--extra-indices data/extra_indices.i32.bin \
--output-dir work/blocks \
--n-samples 3523 \
--n-base-covariates 8 \
--n-phenotype-rows 10000 \
--n-extra-covariate-rows 2000
```
Validate, run, and export results:
```bash
python python/spectra_reml.py validate --manifest work/manifest.json
python python/spectra_reml.py run \
--manifest work/manifest.json \
--engine build-mkl/spectra_reml \
--threads 28 \
--blas-threads 1 \
--block-size 256 \
--resume
python python/spectra_reml.py finalize \
--manifest work/manifest.json \
--output work/results.tsv.gz
```
`results.tsv.gz` retains the complete per-task summary and stores the fixed-effect vector and row-wise packed lower covariance as JSON arrays.
## 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.
- `--resume` reuses complete blocks only when the signature matches exactly.
- `--force` invalidates the old signature before deleting old blocks and starting a new generation.
- `finalize` refuses blocks that are not bound to the current manifest.
- `--dry-run` does not mutate output state.
## Threading
Use one BLAS thread with multiple outer task threads unless benchmarking shows otherwise:
```bash
export MKL_NUM_THREADS=1
export OPENBLAS_NUM_THREADS=1
export OMP_DYNAMIC=FALSE
```
Then set `--threads` to the physical cores allocated to the process.
See [docs/FORMAT.md](docs/FORMAT.md) for the exact file contract.
## License
BSD 3-Clause. See [LICENSE](LICENSE).

View File

@@ -0,0 +1,135 @@
# Locate an LP64 OpenBLAS implementation together with the CBLAS and LAPACKE
# C interfaces. LAPACKE may be provided by libopenblas itself or by a separate
# liblapacke library.
#
# Result:
# REML_OpenBLAS_FOUND
# REML::OpenBLASLAPACKE
#
# Optional hints/overrides:
# OpenBLAS_ROOT
# LAPACKE_ROOT
# REML_OPENBLAS_LIBRARY
# REML_LAPACKE_LIBRARY
# REML_CBLAS_INCLUDE_DIR
# REML_LAPACKE_INCLUDE_DIR
include(CheckCXXSourceCompiles)
include(CMakePushCheckState)
include(FindPackageHandleStandardArgs)
set(_reml_openblas_hints)
foreach(_root IN ITEMS "${OpenBLAS_ROOT}" "$ENV{OpenBLAS_ROOT}" "$ENV{OPENBLAS_ROOT}" "$ENV{CONDA_PREFIX}")
if(_root)
list(APPEND _reml_openblas_hints "${_root}")
endif()
endforeach()
set(_reml_lapacke_hints ${_reml_openblas_hints})
foreach(_root IN ITEMS "${LAPACKE_ROOT}" "$ENV{LAPACKE_ROOT}")
if(_root)
list(PREPEND _reml_lapacke_hints "${_root}")
endif()
endforeach()
# Prefer package-provided imported targets because they carry static-library
# dependencies (Fortran runtime, pthreads, and libm) correctly.
find_package(OpenBLAS CONFIG QUIET)
if(TARGET OpenBLAS::OpenBLAS)
set(_reml_openblas_link OpenBLAS::OpenBLAS)
elseif(TARGET OpenBLAS)
set(_reml_openblas_link OpenBLAS)
endif()
if(NOT _reml_openblas_link)
find_library(REML_OPENBLAS_LIBRARY
NAMES openblas libopenblas
HINTS ${_reml_openblas_hints}
PATH_SUFFIXES lib lib64 Library/lib
)
if(REML_OPENBLAS_LIBRARY)
set(_reml_openblas_link "${REML_OPENBLAS_LIBRARY}")
endif()
endif()
find_path(REML_CBLAS_INCLUDE_DIR
NAMES cblas.h
HINTS ${_reml_openblas_hints}
PATH_SUFFIXES include include/openblas Library/include Library/include/openblas
)
find_path(REML_LAPACKE_INCLUDE_DIR
NAMES lapacke.h
HINTS ${_reml_lapacke_hints}
PATH_SUFFIXES include include/openblas include/lapacke Library/include Library/include/openblas
)
find_package(LAPACKE CONFIG QUIET)
if(TARGET LAPACKE::LAPACKE)
set(_reml_lapacke_link LAPACKE::LAPACKE)
elseif(TARGET LAPACKE)
set(_reml_lapacke_link LAPACKE)
else()
find_library(REML_LAPACKE_LIBRARY
NAMES lapacke liblapacke
HINTS ${_reml_lapacke_hints}
PATH_SUFFIXES lib lib64 Library/lib
)
if(REML_LAPACKE_LIBRARY)
set(_reml_lapacke_link "${REML_LAPACKE_LIBRARY}")
endif()
endif()
if(_reml_lapacke_link)
# Keep the static-link order dependency first, provider second:
# liblapacke calls LAPACK symbols supplied by libopenblas.
set(_reml_link_items ${_reml_lapacke_link} ${_reml_openblas_link})
else()
set(_reml_link_items ${_reml_openblas_link})
endif()
unset(REML_OPENBLAS_LINK_OK CACHE)
if(REML_CBLAS_INCLUDE_DIR AND REML_LAPACKE_INCLUDE_DIR AND _reml_openblas_link)
cmake_push_check_state(RESET)
set(CMAKE_REQUIRED_INCLUDES "${REML_CBLAS_INCLUDE_DIR};${REML_LAPACKE_INCLUDE_DIR}")
set(CMAKE_REQUIRED_LIBRARIES ${_reml_link_items})
check_cxx_source_compiles([[
#include <cblas.h>
#include <lapacke.h>
int main() {
double a[1] = {1.0};
double b[1] = {1.0};
double c[1] = {0.0};
cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,
1, 1, 1, 1.0, a, 1, b, 1, 0.0, c, 1);
return LAPACKE_dpotrf(LAPACK_ROW_MAJOR, 'L', 1, a, 1);
}
]] REML_OPENBLAS_LINK_OK)
cmake_pop_check_state()
else()
set(REML_OPENBLAS_LINK_OK FALSE)
endif()
find_package_handle_standard_args(REML_OpenBLAS
REQUIRED_VARS
REML_CBLAS_INCLUDE_DIR
REML_LAPACKE_INCLUDE_DIR
_reml_openblas_link
REML_OPENBLAS_LINK_OK
FAIL_MESSAGE
"OpenBLAS with LP64 CBLAS and LAPACKE is required. Set OpenBLAS_ROOT/LAPACKE_ROOT, or the REML_* cache variables listed in cmake/FindREML_OpenBLAS.cmake."
)
if(REML_OpenBLAS_FOUND AND NOT TARGET REML::OpenBLASLAPACKE)
add_library(REML::OpenBLASLAPACKE INTERFACE IMPORTED)
set_property(TARGET REML::OpenBLASLAPACKE PROPERTY
INTERFACE_INCLUDE_DIRECTORIES "${REML_CBLAS_INCLUDE_DIR};${REML_LAPACKE_INCLUDE_DIR}")
set_property(TARGET REML::OpenBLASLAPACKE PROPERTY
INTERFACE_LINK_LIBRARIES "${_reml_link_items}")
endif()
mark_as_advanced(
REML_CBLAS_INCLUDE_DIR
REML_LAPACKE_INCLUDE_DIR
REML_OPENBLAS_LIBRARY
REML_LAPACKE_LIBRARY
)

162
docs/FORMAT.md Normal file
View File

@@ -0,0 +1,162 @@
# SpectraREML file contract
All raw binary files are little-endian, headerless, and contiguous. Integer indices and element offsets are zero based.
## Shared sample order
The following matrices must use exactly the same sample order:
1. GRM;
2. common design matrix;
3. phenotype matrix;
4. extra-covariate matrix.
When supplied, `--grm-id` is checked for the expected number of nonempty rows. Domain adapters remain responsible for verifying the actual identifiers and order.
## GRM
`--grm-bin` uses the GCTA lower-triangle packed `float32` layout:
```text
G[0,0],
G[1,0], G[1,1],
G[2,0], G[2,1], G[2,2], ...
```
For `n` samples, the exact file size is `4 * n * (n + 1) / 2` bytes.
## Common design
`--base-x` is a row-major `float64` matrix with shape
```text
n_samples × n_base_covariates
```
It must already contain every common fixed effect, including an intercept if required. SpectraREML does not add or standardize columns.
## Phenotypes
`--phenotypes` is a row-major `float64` matrix with shape
```text
n_phenotype_rows × n_samples
```
Each task selects one row through `phenotype_row`.
## Extra fixed-effect covariates
`--extra-covariates` is an optional row-major `float32` matrix with shape
```text
n_extra_covariate_rows × n_samples
```
Only rows referenced by at least one task are read and rotated. The file may be omitted when the row count and all task-specific counts are zero.
## Task table
`--tasks` is a UTF-8 tab-separated file with exactly four columns:
```text
task_index task_id phenotype_row n_extra_covariates
0 trait_a 0 0
1 trait_b 1 2
```
Requirements:
- `task_index` is consecutive and zero based;
- `task_id` is nonempty and unique;
- `phenotype_row` is within the phenotype matrix;
- `n_extra_covariates` agrees with the CSR offsets.
## CSR task-to-covariate mapping
`--extra-offsets` is an `int64` array of length `n_tasks + 1`. It begins with zero and is nondecreasing.
`--extra-indices` is an `int32` array of length `offsets[-1]`. For task `i`, its extra-covariate row indices are
```text
indices[offsets[i]:offsets[i+1]]
```
An index must be in `[0, n_extra_covariate_rows)`, and a task cannot reference the same row twice.
## Block output
For block number `KKKKKK`:
```text
block_KKKKKK.summary.tsv
block_KKKKKK.beta.f64.bin
block_KKKKKK.cov.f64.bin
block_KKKKKK.complete
```
The summary header is:
```text
task_index
task_id
status
n_fixed
n_extra_covariates
beta_offset
cov_offset
sigma_g2
sigma_e2
h2
logL
iterations
line_search_steps
grad_inf
error
```
`beta_offset` and `cov_offset` count `float64` elements, not bytes. A negative offset indicates that no estimates were emitted for that task.
The covariance array uses the row-wise packed lower triangle:
```text
(0,0), (1,0), (1,1), (2,0), (2,1), (2,2), ...
```
The `.complete` marker is written last and contains tab-separated key/value rows:
```text
format spectra-reml-block-v1
block 0
tasks 256
beta_elements 4096
cov_elements 34816
```
Consumers must ignore blocks without `.complete`.
## Status values
```text
converged
converged_boundary
max_iterations
line_search_failed
rank_deficient
invalid_input
non_positive_covariance
numerical_error
```
`converged_boundary` is a successful residual-only solution accepted after the one-sided variance-component score and likelihood checks.
## Generic finalized output
The Python CLI exports one TSV row per task. It includes the full summary plus:
```text
beta_json
covariance_packed_lower_json
```
Project-specific software can attach coefficient names and derive contrasts without changing the numerical engine.

83
examples/example.py Normal file
View File

@@ -0,0 +1,83 @@
#!/usr/bin/env python3
"""Create a small, fully synthetic SpectraREML input bundle."""
from __future__ import annotations
import sys
from pathlib import Path
import numpy as np
PROJECT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT / "python"))
import spectra_reml # noqa: E402
def main() -> None:
output = Path(__file__).resolve().parent / "synthetic_work"
output.mkdir(parents=True, exist_ok=True)
rng = np.random.default_rng(20260809)
n = 12
z = rng.normal(size=(n, 5))
grm = z @ z.T / z.shape[1]
diagonal = np.sqrt(np.diag(grm))
grm /= np.outer(diagonal, diagonal)
packed = np.asarray(
[grm[row, col] for row in range(n) for col in range(row + 1)],
dtype="<f4",
)
packed.tofile(output / "example.grm.bin")
(output / "example.grm.id").write_text(
"".join("sample{0}\tsample{0}\n".format(i + 1) for i in range(n)),
encoding="utf-8",
)
base_x = np.column_stack((np.ones(n), np.linspace(-1.0, 1.0, n)))
extra = rng.normal(size=(2, n)).astype("<f4")
phenotypes = np.vstack(
(
0.3 + 0.5 * base_x[:, 1] + rng.normal(scale=0.2, size=n),
-0.2 + 0.8 * extra[0] + rng.normal(scale=0.2, size=n),
0.1 + 0.4 * extra[0] - 0.3 * extra[1] + rng.normal(scale=0.2, size=n),
)
)
np.asarray(base_x, dtype="<f8").tofile(output / "base_x.f64.bin")
np.asarray(phenotypes, dtype="<f8").tofile(output / "phenotypes.f64.bin")
extra.tofile(output / "extra_covariates.f32.bin")
(output / "tasks.tsv").write_text(
"task_index\ttask_id\tphenotype_row\tn_extra_covariates\n"
"0\ttrait_0\t0\t0\n"
"1\ttrait_1\t1\t1\n"
"2\ttrait_2\t2\t2\n",
encoding="utf-8",
)
np.asarray([0, 0, 1, 3], dtype="<i8").tofile(output / "extra_offsets.i64.bin")
np.asarray([0, 0, 1], dtype="<i4").tofile(output / "extra_indices.i32.bin")
manifest = output / "manifest.json"
spectra_reml.main(
[
"make-manifest",
"--manifest", str(manifest),
"--grm-bin", str(output / "example.grm.bin"),
"--grm-id", str(output / "example.grm.id"),
"--base-x", str(output / "base_x.f64.bin"),
"--phenotypes", str(output / "phenotypes.f64.bin"),
"--extra-covariates", str(output / "extra_covariates.f32.bin"),
"--tasks", str(output / "tasks.tsv"),
"--extra-offsets", str(output / "extra_offsets.i64.bin"),
"--extra-indices", str(output / "extra_indices.i32.bin"),
"--output-dir", str(output / "blocks"),
"--n-samples", str(n),
"--n-base-covariates", str(base_x.shape[1]),
"--n-phenotype-rows", str(phenotypes.shape[0]),
"--n-extra-covariate-rows", str(extra.shape[0]),
]
)
print("\nNext:")
print("python python/spectra_reml.py run --manifest {} --engine /path/to/spectra_reml --threads 4".format(manifest))
if __name__ == "__main__":
main()

View File

@@ -0,0 +1,62 @@
#pragma once
#include "spectra_reml/grm.hpp"
#include "spectra_reml/types.hpp"
#include <cstddef>
#include <cstdint>
#include <filesystem>
#include <string>
#include <vector>
namespace spectra::reml {
struct BatchInputPaths {
std::filesystem::path grm_bin;
std::filesystem::path grm_id;
std::filesystem::path base_x_f64;
std::filesystem::path phenotype_f64;
std::filesystem::path extra_covariate_f32;
std::filesystem::path tasks_tsv;
std::filesystem::path extra_offsets_i64;
std::filesystem::path extra_indices_i32;
std::filesystem::path output_directory;
};
struct BatchDimensions {
std::size_t sample_count = 0;
std::size_t base_covariate_count = 0;
std::size_t phenotype_row_count = 0;
std::size_t extra_covariate_row_count = 0;
};
struct BatchOptions {
std::size_t block_size = 256;
std::size_t outer_threads = 1;
bool resume = false;
bool overwrite = false;
RemlOptions reml;
};
[[nodiscard]] ColMajorMatrix read_row_major_f64_matrix(
const std::filesystem::path& path, std::size_t rows, std::size_t cols);
[[nodiscard]] std::vector<RemlTask> read_tasks(
const std::filesystem::path& tasks_tsv,
const std::filesystem::path& extra_offsets_i64,
const std::filesystem::path& extra_indices_i32,
const BatchDimensions& dimensions);
// Runs one complete task set. The GRM and common variables are transformed
// once; phenotypes are transformed in blocks and tasks are fitted in parallel.
// Each completed block contains:
// block_NNNNNN.summary.tsv
// block_NNNNNN.beta.f64.bin
// block_NNNNNN.cov.f64.bin
// block_NNNNNN.complete
// The .complete marker is renamed last and is the sole resume criterion.
void run_task_batch(const BatchInputPaths& paths,
const BatchDimensions& dimensions,
const BatchOptions& options);
} // namespace spectra::reml

View File

@@ -0,0 +1,36 @@
#pragma once
#include "spectra_reml/types.hpp"
#include <cstddef>
#include <filesystem>
#include <vector>
namespace spectra::reml {
struct SpectralGrm {
std::vector<double> eigenvalues;
ColMajorMatrix eigenvectors;
double minimum_eigenvalue = 0.0;
double maximum_eigenvalue = 0.0;
};
// GCTA stores the lower triangle, row by row, as IEEE-754 float32 values:
// (0,0), (1,0), (1,1), (2,0), ... . The returned matrix is full symmetric,
// double precision and column-major.
[[nodiscard]] ColMajorMatrix read_gcta_grm_lower_triangle(
const std::filesystem::path& grm_bin, std::size_t sample_count);
// Reads and diagonalizes a GRM once. Eigenvalues are deliberately not
// clipped: preserving the supplied GRM is required for numerical equivalence.
// The REML evaluator rejects parameter trials for which sigma_g^2 lambda_i +
// sigma_e^2 is not positive.
[[nodiscard]] SpectralGrm read_and_diagonalize_gcta_grm(
const std::filesystem::path& grm_bin, std::size_t sample_count);
// Optional structural check for a GCTA .grm.id file. The Python driver owns
// sample-ID alignment; this guard catches wrong dimensions at the CLI boundary.
void validate_grm_id_count(const std::filesystem::path& grm_id,
std::size_t expected_sample_count);
} // namespace spectra::reml

View File

@@ -0,0 +1,41 @@
#pragma once
#include "spectra_reml/types.hpp"
#include <cstddef>
#include <string>
#include <vector>
namespace spectra::reml {
// C = A' B. A is n x k and B is n x m; C is k x m.
[[nodiscard]] ColMajorMatrix cross_product(const ColMajorMatrix& a,
const ColMajorMatrix& b);
// C = A' B, specialized for an orthogonal n x n matrix A. This is kept as a
// distinct API because it is the dominant level-3 BLAS operation in a batch.
[[nodiscard]] ColMajorMatrix rotate_to_eigenspace(
const ColMajorMatrix& eigenvectors, const ColMajorMatrix& variables);
// Overwrites a symmetric matrix with its eigenvectors (columns) and returns
// eigenvalues in ascending order. Only the lower triangle is inspected.
[[nodiscard]] std::vector<double> symmetric_eigen_decomposition(
ColMajorMatrix& symmetric_matrix);
// Cholesky helpers. The factor is lower triangular and stored in the lower
// triangle of a full column-major matrix.
[[nodiscard]] bool cholesky_factor_in_place(std::vector<double>& matrix,
std::size_t order,
std::string* error = nullptr);
[[nodiscard]] bool cholesky_solve_in_place(
const std::vector<double>& factor, std::size_t order, double* rhs,
std::size_t rhs_columns, std::string* error = nullptr);
[[nodiscard]] bool cholesky_inverse(
const std::vector<double>& factor, std::size_t order,
std::vector<double>& inverse, std::string* error = nullptr);
[[nodiscard]] double dot(const std::vector<double>& a,
const std::vector<double>& b);
[[nodiscard]] double infinity_norm(const std::vector<double>& x);
} // namespace spectra::reml

View File

@@ -0,0 +1,49 @@
#pragma once
#include "spectra_reml/types.hpp"
#include <limits>
#include <string>
#include <vector>
namespace spectra::reml {
// Fits y = X beta + g + e in the GRM eigenspace, where
// Var(y) = sigma_g^2 diag(lambda) + sigma_e^2 I.
//
// The parameters are the signed standard deviations (sigma_e, sigma_g), as in
// thesis Equations 4.24-4.25. The covariance depends on their squares, so the
// optimization is unconstrained. It uses AI-REML and a strong-Wolfe line
// search; no EM update is performed and P is never materialized.
[[nodiscard]] RemlResult fit_ai_reml_spectral(
const std::vector<double>& y_star, const ColMajorMatrix& x_star,
const std::vector<double>& eigenvalues,
const RemlOptions& options = {});
// Exposed for finite-difference and independent-oracle tests. The ordering is
// theta=(sigma_e,sigma_g); ai is [ee,eg;eg,gg] in column-major order.
struct RemlEvaluation {
bool valid = false;
double log_likelihood = -std::numeric_limits<double>::infinity();
double gradient_e = std::numeric_limits<double>::quiet_NaN();
double gradient_g = std::numeric_limits<double>::quiet_NaN();
// Scores with respect to the variance parameters v_e=sigma_e^2 and
// v_g=sigma_g^2. These remain informative when sigma_g=0, unlike the
// signed-standard-deviation gradient gradient_g=2*sigma_g*score_v_g.
double variance_score_e = std::numeric_limits<double>::quiet_NaN();
double variance_score_g = std::numeric_limits<double>::quiet_NaN();
double ai_ee = std::numeric_limits<double>::quiet_NaN();
double ai_eg = std::numeric_limits<double>::quiet_NaN();
double ai_gg = std::numeric_limits<double>::quiet_NaN();
std::vector<double> beta;
std::vector<double> beta_covariance;
std::string error;
};
[[nodiscard]] 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,
bool compute_ai = true, bool compute_beta_covariance = false,
double covariance_floor_relative = 1e-12);
} // namespace spectra::reml

View File

@@ -0,0 +1,133 @@
#pragma once
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>
namespace spectra::reml {
// Dense matrices use Fortran/BLAS column-major storage throughout the
// numerical core. Keeping this convention at the API boundary prevents
// hidden transposes before LAPACK calls.
class ColMajorMatrix {
public:
ColMajorMatrix() = default;
ColMajorMatrix(std::size_t rows, std::size_t cols)
: rows_(rows), cols_(cols), values_(rows * cols, 0.0) {}
ColMajorMatrix(std::size_t rows, std::size_t cols,
std::vector<double> values)
: rows_(rows), cols_(cols), values_(std::move(values)) {
if (values_.size() != rows_ * cols_) {
throw std::invalid_argument("ColMajorMatrix data has the wrong size");
}
}
[[nodiscard]] std::size_t rows() const noexcept { return rows_; }
[[nodiscard]] std::size_t cols() const noexcept { return cols_; }
[[nodiscard]] std::size_t size() const noexcept { return values_.size(); }
[[nodiscard]] bool empty() const noexcept { return values_.empty(); }
[[nodiscard]] std::size_t leading_dimension() const noexcept { return rows_; }
double* data() noexcept { return values_.data(); }
const double* data() const noexcept { return values_.data(); }
std::vector<double>& values() noexcept { return values_; }
const std::vector<double>& values() const noexcept { return values_; }
double& operator()(std::size_t row, std::size_t col) noexcept {
return values_[row + col * rows_];
}
double operator()(std::size_t row, std::size_t col) const noexcept {
return values_[row + col * rows_];
}
private:
std::size_t rows_ = 0;
std::size_t cols_ = 0;
std::vector<double> values_;
};
enum class FitStatus {
converged,
converged_boundary,
max_iterations,
line_search_failed,
rank_deficient,
invalid_input,
non_positive_covariance,
numerical_error
};
[[nodiscard]] const char* to_string(FitStatus status) noexcept;
struct RemlOptions {
std::size_t max_iterations = 100;
std::size_t line_search_max_evaluations = 48;
std::size_t line_search_max_zoom_iterations = 48;
double gradient_absolute_tolerance = 1e-7;
double gradient_relative_tolerance = 1e-8;
double step_relative_tolerance = 1e-9;
double likelihood_relative_tolerance = 1e-11;
double wolfe_c1 = 1e-4;
double wolfe_c2 = 0.9;
double initial_line_search_step = 1.0;
double maximum_line_search_step = 64.0;
double covariance_floor_relative = 1e-12;
double rank_tolerance_relative = 1e-10;
double ai_ridge_relative = 1e-10;
std::size_t ai_ridge_attempts = 8;
// KKT check for the residual-only boundary sigma_g^2=0. The boundary is
// accepted only when the one-sided variance-component score is no larger
// than boundary_score_tolerance and its likelihood is not inferior to the
// current interior iterate (within the stated relative roundoff tolerance).
double boundary_score_tolerance = 1e-8;
double boundary_likelihood_relative_tolerance = 1e-12;
double boundary_h2_trigger = 1e-3;
// NaN means use the paper's initial value: both standard deviations are
// 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();
};
struct RemlResult {
FitStatus status = FitStatus::invalid_input;
double sigma_e = std::numeric_limits<double>::quiet_NaN();
double sigma_g = std::numeric_limits<double>::quiet_NaN();
double sigma_e2 = std::numeric_limits<double>::quiet_NaN();
double sigma_g2 = std::numeric_limits<double>::quiet_NaN();
double h2 = std::numeric_limits<double>::quiet_NaN();
// Equation B.5 of the thesis, omitting additive constants.
double log_likelihood = -std::numeric_limits<double>::infinity();
std::size_t iterations = 0;
std::size_t line_search_evaluations = 0;
double gradient_inf_norm = std::numeric_limits<double>::infinity();
std::vector<double> beta;
// Row-wise packed lower triangle:
// (0,0), (1,0), (1,1), (2,0), (2,1), (2,2), ...
std::vector<double> beta_covariance_packed_lower;
std::string error;
[[nodiscard]] bool has_estimates() const noexcept {
return !beta.empty() && std::isfinite(log_likelihood);
}
};
struct RemlTask {
std::uint64_t task_index = 0;
std::string task_id;
std::uint64_t phenotype_row = 0;
std::vector<std::int32_t> extra_covariate_rows;
};
struct TaskResult {
RemlTask task;
RemlResult fit;
};
} // namespace spectra::reml

766
python/spectra_reml.py Normal file
View File

@@ -0,0 +1,766 @@
#!/usr/bin/env python3
"""Generic manifest, execution, recovery, and result CLI for SpectraREML."""
from __future__ import annotations
import argparse
import csv
import gzip
import hashlib
import json
import os
import shutil
import subprocess
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Mapping, NoReturn, Sequence
for _name in (
"OMP_NUM_THREADS",
"OPENBLAS_NUM_THREADS",
"MKL_NUM_THREADS",
"NUMEXPR_NUM_THREADS",
):
os.environ.setdefault(_name, "1")
try:
import numpy as np
except ImportError as exc: # pragma: no cover
raise SystemExit("NumPy is required: {}".format(exc))
MANIFEST_FORMAT = "spectra-reml-manifest-v1"
BLOCK_FORMAT = "spectra-reml-block-v1"
RUN_SIGNATURE_FORMAT = "spectra-reml-run-signature-v1"
FINALIZE_FORMAT = "spectra-reml-finalize-v1"
TASK_HEADER = (
"task_index",
"task_id",
"phenotype_row",
"n_extra_covariates",
)
SUMMARY_HEADER = (
"task_index",
"task_id",
"status",
"n_fixed",
"n_extra_covariates",
"beta_offset",
"cov_offset",
"sigma_g2",
"sigma_e2",
"h2",
"logL",
"iterations",
"line_search_steps",
"grad_inf",
"error",
)
SUCCESS_STATUSES = frozenset(("converged", "converged_boundary"))
BLOCK_PATTERNS = (
"block_*.summary.tsv",
"block_*.beta.f64.bin",
"block_*.cov.f64.bin",
"block_*.complete",
)
def fail(message: str) -> NoReturn:
raise RuntimeError(message)
def log(message: str) -> None:
print(time.strftime("%Y-%m-%d %H:%M:%S"), "|", message, flush=True)
def utc_now() -> str:
return datetime.now(timezone.utc).isoformat(timespec="seconds")
def resolved(path: Path) -> Path:
return path.expanduser().resolve()
def file_identity(path: Path) -> dict[str, Any]:
path = resolved(path)
stat = path.stat()
return {
"path": str(path),
"size_bytes": int(stat.st_size),
"mtime_ns": int(stat.st_mtime_ns),
}
def sha256_file(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as handle:
for block in iter(lambda: handle.read(1024 * 1024), b""):
digest.update(block)
return digest.hexdigest()
def canonical_sha256(value: Any) -> str:
encoded = json.dumps(
value,
ensure_ascii=False,
allow_nan=False,
sort_keys=True,
separators=(",", ":"),
).encode("utf-8")
return hashlib.sha256(encoded).hexdigest()
def atomic_json(path: Path, value: Mapping[str, Any]) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
temporary = path.with_name(path.name + ".tmp.{}".format(os.getpid()))
try:
with temporary.open("w", encoding="utf-8", newline="") as handle:
json.dump(value, handle, ensure_ascii=False, allow_nan=False, indent=2)
handle.write("\n")
handle.flush()
os.fsync(handle.fileno())
os.replace(temporary, path)
finally:
try:
temporary.unlink()
except FileNotFoundError:
pass
def read_json(path: Path) -> dict[str, Any]:
try:
with path.open("r", encoding="utf-8-sig") as handle:
value = json.load(handle)
except (OSError, json.JSONDecodeError) as exc:
fail("Cannot read JSON {}: {}".format(path, exc))
if not isinstance(value, dict):
fail("JSON root is not an object: {}".format(path))
return value
def require_file(path: Path, label: str, allow_empty: bool = False) -> None:
if not path.is_file():
fail("{} is missing: {}".format(label, path))
if not allow_empty and path.stat().st_size == 0:
fail("{} is empty: {}".format(label, path))
def require_size(path: Path, expected: int, label: str) -> None:
require_file(path, label, allow_empty=(expected == 0))
observed = path.stat().st_size
if observed != expected:
fail("{} size mismatch: observed {}; expected {}: {}".format(
label, observed, expected, path
))
def manifest_paths(manifest: Mapping[str, Any]) -> dict[str, Path | None]:
raw = manifest.get("paths")
if not isinstance(raw, dict):
fail("Manifest has no paths object.")
required = (
"grm_bin",
"base_x_f64",
"phenotype_f64",
"tasks_tsv",
"extra_offsets_i64",
"extra_indices_i32",
"output_directory",
)
missing = [name for name in required if name not in raw]
if missing:
fail("Manifest paths are missing: {}".format(", ".join(missing)))
result: dict[str, Path | None] = {}
for name in required + ("grm_id", "extra_covariate_f32"):
value = raw.get(name)
result[name] = Path(str(value)) if value not in (None, "") else None
return result
def read_tasks(path: Path) -> list[dict[str, str]]:
with path.open("r", encoding="utf-8-sig", newline="") as handle:
reader = csv.DictReader(handle, delimiter="\t")
if tuple(reader.fieldnames or ()) != TASK_HEADER:
fail("Task header must be exactly: {}".format("\t".join(TASK_HEADER)))
return list(reader)
def validate_manifest(path: Path) -> dict[str, Any]:
path = resolved(path)
manifest = read_json(path)
if manifest.get("format") != MANIFEST_FORMAT:
fail("Unsupported manifest format: {!r}".format(manifest.get("format")))
dimensions = manifest.get("dimensions")
if not isinstance(dimensions, dict):
fail("Manifest has no dimensions object.")
names = (
"sample_count",
"base_covariate_count",
"phenotype_row_count",
"extra_covariate_row_count",
"task_count",
"extra_index_count",
)
parsed: dict[str, int] = {}
for name in names:
value = dimensions.get(name)
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
fail("Invalid dimension {}={!r}".format(name, value))
parsed[name] = value
if min(parsed["sample_count"], parsed["base_covariate_count"], parsed["task_count"]) <= 0:
fail("Sample, base-covariate, and task counts must be positive.")
paths = manifest_paths(manifest)
n = parsed["sample_count"]
require_size(paths["grm_bin"], n * (n + 1) // 2 * 4, "GRM") # type: ignore[arg-type]
require_size(paths["base_x_f64"], n * parsed["base_covariate_count"] * 8, "base design") # type: ignore[arg-type]
require_size(paths["phenotype_f64"], n * parsed["phenotype_row_count"] * 8, "phenotypes") # type: ignore[arg-type]
require_size(paths["extra_offsets_i64"], (parsed["task_count"] + 1) * 8, "extra offsets") # type: ignore[arg-type]
require_size(paths["extra_indices_i32"], parsed["extra_index_count"] * 4, "extra indices") # type: ignore[arg-type]
if parsed["extra_covariate_row_count"]:
if paths["extra_covariate_f32"] is None:
fail("extra_covariate_f32 is required when its row count is nonzero.")
require_size(paths["extra_covariate_f32"], n * parsed["extra_covariate_row_count"] * 4, "extra covariates")
if paths["grm_id"] is not None:
require_file(paths["grm_id"], "GRM IDs")
with paths["grm_id"].open("r", encoding="utf-8-sig") as handle:
if sum(bool(line.strip()) for line in handle) != n:
fail("GRM ID row count differs from sample_count.")
sources = manifest.get("source_files")
if not isinstance(sources, dict) or not sources:
fail("Manifest has no source identities.")
for label, expected in sources.items():
if not isinstance(expected, dict):
fail("Malformed source identity: {}".format(label))
source = Path(str(expected.get("path", "")))
if not source.is_file():
fail("Source is missing: {}: {}".format(label, source))
observed = source.stat()
if (
observed.st_size != int(expected.get("size_bytes", -1))
or observed.st_mtime_ns != int(expected.get("mtime_ns", -1))
):
fail("Source changed since manifest creation: {}: {}".format(label, source))
controls = {
"base_x_f64": paths["base_x_f64"],
"tasks_tsv": paths["tasks_tsv"],
"extra_offsets_i64": paths["extra_offsets_i64"],
"extra_indices_i32": paths["extra_indices_i32"],
}
hashes = manifest.get("control_sha256")
if not isinstance(hashes, dict) or set(hashes) != set(controls):
fail("control_sha256 does not cover all control files.")
for name, control in controls.items():
assert control is not None
if sha256_file(control) != hashes[name]:
fail("Control checksum mismatch: {}".format(control))
tasks = read_tasks(paths["tasks_tsv"]) # type: ignore[arg-type]
if len(tasks) != parsed["task_count"]:
fail("Task count differs from manifest.")
offsets = np.fromfile(paths["extra_offsets_i64"], dtype="<i8")
indices = np.fromfile(paths["extra_indices_i32"], dtype="<i4")
if offsets[0] != 0 or np.any(offsets[1:] < offsets[:-1]) or int(offsets[-1]) != indices.size:
fail("Extra-covariate CSR offsets are invalid.")
if indices.size and (
int(indices.min()) < 0
or int(indices.max()) >= parsed["extra_covariate_row_count"]
):
fail("An extra-covariate index is outside its matrix.")
seen_ids: set[str] = set()
for expected, row in enumerate(tasks):
try:
task_index = int(row["task_index"])
phenotype_row = int(row["phenotype_row"])
extra_count = int(row["n_extra_covariates"])
except ValueError as exc:
fail("Invalid task integer at row {}: {}".format(expected, exc))
if task_index != expected:
fail("Task indices must be consecutive and zero based.")
if not row["task_id"] or row["task_id"] in seen_ids:
fail("Task IDs must be nonempty and unique.")
seen_ids.add(row["task_id"])
if phenotype_row < 0 or phenotype_row >= parsed["phenotype_row_count"]:
fail("Task phenotype row is outside its matrix.")
if extra_count != int(offsets[expected + 1] - offsets[expected]):
fail("Task extra count disagrees with CSR offsets.")
return manifest
def make_manifest(args: argparse.Namespace) -> Path:
target = resolved(args.manifest)
paths: dict[str, Path | None] = {
"grm_bin": resolved(args.grm_bin),
"grm_id": resolved(args.grm_id) if args.grm_id else None,
"base_x_f64": resolved(args.base_x),
"phenotype_f64": resolved(args.phenotypes),
"extra_covariate_f32": resolved(args.extra_covariates) if args.extra_covariates else None,
"tasks_tsv": resolved(args.tasks),
"extra_offsets_i64": resolved(args.extra_offsets),
"extra_indices_i32": resolved(args.extra_indices),
"output_directory": resolved(args.output_dir),
}
tasks = read_tasks(paths["tasks_tsv"]) # type: ignore[arg-type]
extra_index_count = paths["extra_indices_i32"].stat().st_size // 4 # type: ignore[union-attr]
source_names = ("grm_bin", "phenotype_f64")
sources = {name: file_identity(paths[name]) for name in source_names} # type: ignore[arg-type]
for name in ("grm_id", "extra_covariate_f32"):
if paths[name] is not None:
sources[name] = file_identity(paths[name]) # type: ignore[arg-type]
controls = ("base_x_f64", "tasks_tsv", "extra_offsets_i64", "extra_indices_i32")
manifest = {
"format": MANIFEST_FORMAT,
"created_utc": utc_now(),
"dimensions": {
"sample_count": args.n_samples,
"base_covariate_count": args.n_base_covariates,
"phenotype_row_count": args.n_phenotype_rows,
"extra_covariate_row_count": args.n_extra_covariate_rows,
"task_count": len(tasks),
"extra_index_count": extra_index_count,
},
"paths": {name: str(value) if value is not None else None for name, value in paths.items()},
"source_files": sources,
"control_sha256": {name: sha256_file(paths[name]) for name in controls}, # type: ignore[arg-type]
}
atomic_json(target, manifest)
validate_manifest(target)
log("manifest created: {} tasks -> {}".format(len(tasks), target))
return target
def resolve_engine(value: str) -> Path:
candidate = Path(value).expanduser()
if candidate.is_file():
return candidate.resolve()
found = shutil.which(value)
if found:
return Path(found).resolve()
fail("SpectraREML engine was not found: {}".format(value))
def build_engine_command(engine: Path, manifest: Mapping[str, Any], args: argparse.Namespace) -> list[str]:
paths = manifest_paths(manifest)
dims = manifest["dimensions"]
command = [
str(engine),
"--grm-bin", str(paths["grm_bin"]),
"--base-x", str(paths["base_x_f64"]),
"--phenotypes", str(paths["phenotype_f64"]),
"--tasks", str(paths["tasks_tsv"]),
"--extra-offsets", str(paths["extra_offsets_i64"]),
"--extra-indices", str(paths["extra_indices_i32"]),
"--out-dir", str(paths["output_directory"]),
"--n-samples", str(dims["sample_count"]),
"--n-base-covariates", str(dims["base_covariate_count"]),
"--n-phenotype-rows", str(dims["phenotype_row_count"]),
"--n-extra-covariate-rows", str(dims["extra_covariate_row_count"]),
"--block-size", str(args.block_size),
"--threads", str(args.threads),
]
if paths["grm_id"] is not None:
command.extend(("--grm-id", str(paths["grm_id"])))
if int(dims["extra_index_count"]):
command.extend(("--extra-covariates", str(paths["extra_covariate_f32"])))
option_map = (
("max_iterations", "--max-iterations"),
("line_search_max_evals", "--line-search-max-evals"),
("line_search_max_zoom", "--line-search-max-zoom"),
("gradient_abs_tol", "--gradient-abs-tol"),
("gradient_rel_tol", "--gradient-rel-tol"),
("step_rel_tol", "--step-rel-tol"),
("likelihood_rel_tol", "--likelihood-rel-tol"),
("wolfe_c1", "--wolfe-c1"),
("wolfe_c2", "--wolfe-c2"),
("initial_step", "--initial-step"),
("maximum_step", "--maximum-step"),
("rank_tol", "--rank-tol"),
("covariance_floor", "--covariance-floor"),
("boundary_h2_trigger", "--boundary-h2-trigger"),
("boundary_score_tol", "--boundary-score-tol"),
("boundary_logl_rel_tol", "--boundary-logl-rel-tol"),
)
for attribute, option in option_map:
command.extend((option, repr(getattr(args, attribute))))
for attribute, option in (("initial_sigma_e", "--initial-sigma-e"), ("initial_sigma_g", "--initial-sigma-g")):
value = getattr(args, attribute)
if value is not None:
command.extend((option, repr(value)))
if args.resume:
command.append("--resume")
elif args.force:
command.append("--overwrite")
return command
def engine_identity(engine: Path) -> dict[str, Any]:
value = file_identity(engine)
value["sha256"] = sha256_file(engine)
return value
def signature_payload(manifest: Mapping[str, Any], engine: Path, args: argparse.Namespace) -> dict[str, Any]:
excluded = {"action", "manifest", "engine", "dry_run", "output"}
options = {
key: value for key, value in vars(args).items()
if key not in excluded and key not in {"resume", "force"}
}
return {
"manifest_sha256": canonical_sha256(manifest),
"engine": engine_identity(engine),
"options": options,
}
def make_signature(manifest: Mapping[str, Any], engine: Path, args: argparse.Namespace) -> dict[str, Any]:
payload = signature_payload(manifest, engine, args)
return {
"format": RUN_SIGNATURE_FORMAT,
"created_utc": utc_now(),
"signature_sha256": canonical_sha256(payload),
"payload": payload,
}
def read_signature(path: Path) -> dict[str, Any]:
value = read_json(path)
payload = value.get("payload")
if value.get("format") != RUN_SIGNATURE_FORMAT or not isinstance(payload, dict):
fail("Invalid run signature: {}".format(path))
if value.get("signature_sha256") != canonical_sha256(payload):
fail("Run signature integrity check failed: {}".format(path))
return value
def block_files(output: Path) -> list[Path]:
values: set[Path] = set()
for pattern in BLOCK_PATTERNS:
values.update(path for path in output.glob(pattern) if path.is_file())
return sorted(values)
def validate_resume_layout(output: Path, task_count: int, block_size: int) -> None:
for marker in output.glob("block_*.complete"):
digits = marker.name[len("block_"):-len(".complete")]
if len(digits) != 6 or not digits.isdigit():
fail("Invalid block marker: {}".format(marker))
index = int(digits)
expected = min(block_size, max(0, task_count - index * block_size))
values: dict[str, str] = {}
with marker.open("r", encoding="utf-8") as handle:
for line in handle:
fields = line.rstrip("\r\n").split("\t", 1)
if len(fields) != 2:
fail("Malformed marker: {}".format(marker))
values[fields[0]] = fields[1]
if expected <= 0 or values.get("format") != BLOCK_FORMAT or int(values.get("tasks", -1)) != expected:
fail("Block marker is incompatible with this manifest/block size: {}".format(marker))
def read_marker(path: Path) -> dict[str, str]:
values: dict[str, str] = {}
with path.open("r", encoding="utf-8") as handle:
for line in handle:
fields = line.rstrip("\r\n").split("\t", 1)
if len(fields) != 2 or not fields[0] or fields[0] in values:
fail("Malformed completion marker: {}".format(path))
values[fields[0]] = fields[1]
return values
def run_engine(manifest: Mapping[str, Any], args: argparse.Namespace) -> None:
if args.dry_run:
print(json.dumps(build_engine_command(Path(args.engine), manifest, args), ensure_ascii=False))
return
engine = resolve_engine(args.engine)
expected_signature = make_signature(manifest, engine, args)
paths = manifest_paths(manifest)
output = paths["output_directory"]
assert output is not None
output.mkdir(parents=True, exist_ok=True)
signature_path = output / "run.signature.json"
markers = list(output.glob("block_*.complete"))
if args.force:
try:
signature_path.unlink()
except FileNotFoundError:
pass
for path in block_files(output):
path.unlink()
atomic_json(signature_path, expected_signature)
elif args.resume:
if markers:
if not signature_path.is_file():
fail("Completed blocks have no run signature; use --force.")
observed = read_signature(signature_path)
if observed["signature_sha256"] != expected_signature["signature_sha256"]:
fail("Run signature differs from engine, manifest, or options; use --force.")
validate_resume_layout(output, int(manifest["dimensions"]["task_count"]), args.block_size)
else:
atomic_json(signature_path, expected_signature)
else:
if block_files(output):
fail("Block outputs already exist; use --resume or --force.")
atomic_json(signature_path, expected_signature)
command = build_engine_command(engine, manifest, args)
environment = os.environ.copy()
environment.update({
"OMP_NUM_THREADS": str(args.threads),
"MKL_NUM_THREADS": str(args.blas_threads),
"OPENBLAS_NUM_THREADS": str(args.blas_threads),
"OMP_DYNAMIC": "FALSE",
})
log("starting SpectraREML engine")
with (output / "engine.log").open("a" if args.resume else "w", encoding="utf-8") as handle:
handle.write("# {}\n# command={}\n".format(utc_now(), json.dumps(command, ensure_ascii=False)))
handle.flush()
completed = subprocess.run(command, stdout=handle, stderr=subprocess.STDOUT, env=environment, check=False)
atomic_json(output / "run.json", {
"finished_utc": utc_now(),
"return_code": completed.returncode,
"command": command,
"run_signature_sha256": expected_signature["signature_sha256"],
})
if completed.returncode:
fail("SpectraREML failed with exit code {}; see engine.log".format(completed.returncode))
def finalizable_signature(manifest: Mapping[str, Any]) -> dict[str, Any]:
output = manifest_paths(manifest)["output_directory"]
assert output is not None
path = output / "run.signature.json"
if not path.is_file():
fail("Cannot finalize without a run signature.")
value = read_signature(path)
if value["payload"].get("manifest_sha256") != canonical_sha256(manifest):
fail("Run signature does not belong to the current manifest.")
return value
def read_completed(manifest: Mapping[str, Any]) -> list[dict[str, Any]]:
output = manifest_paths(manifest)["output_directory"]
assert output is not None
task_count = int(manifest["dimensions"]["task_count"])
results: dict[int, dict[str, Any]] = {}
for marker in sorted(output.glob("block_*.complete")):
stem = marker.name[:-len(".complete")]
summary_path = output / (stem + ".summary.tsv")
beta_path = output / (stem + ".beta.f64.bin")
cov_path = output / (stem + ".cov.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"])
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
) < 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")
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:
fail("Unexpected summary header: {}".format(summary_path))
rows = list(reader)
if len(rows) != declared_tasks:
fail("Summary row count disagrees with marker: {}".format(marker))
beta = np.fromfile(beta_path, dtype="<f8")
cov = np.fromfile(cov_path, dtype="<f8")
next_beta_offset = 0
next_cov_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"])
if beta_offset < 0 or cov_offset < 0:
if beta_offset != cov_offset:
fail("Only one output offset is negative for task {}".format(index))
beta_values: list[float] = []
cov_values: list[float] = []
else:
packed = p * (p + 1) // 2
if beta_offset != next_beta_offset or cov_offset != next_cov_offset:
fail("Non-contiguous output offsets for task {}".format(index))
if beta_offset + p > beta.size or cov_offset + packed > cov.size:
fail("Output offsets exceed binary arrays for task {}".format(index))
beta_values = beta[beta_offset:beta_offset + p].tolist()
cov_values = cov[cov_offset:cov_offset + packed].tolist()
next_beta_offset += p
next_cov_offset += packed
result: dict[str, Any] = dict(row)
result["beta_json"] = json.dumps(beta_values, separators=(",", ":"))
result["covariance_packed_lower_json"] = json.dumps(cov_values, separators=(",", ":"))
results[index] = result
if next_beta_offset != beta_elements or next_cov_offset != cov_elements:
fail("Block binary arrays contain unused elements: {}".format(stem))
missing = sorted(set(range(task_count)).difference(results))
if missing:
fail("Completed blocks are missing {} tasks.".format(len(missing)))
return [results[index] for index in range(task_count)]
def finalize(manifest: Mapping[str, Any], output: Path) -> Path:
signature = finalizable_signature(manifest)
rows = read_completed(manifest)
output = resolved(output)
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"]
try:
with opener(temporary, "wt", encoding="utf-8", newline="") as handle:
writer = csv.DictWriter(handle, fieldnames=fields, delimiter="\t", lineterminator="\n")
writer.writeheader()
writer.writerows(rows)
os.replace(temporary, output)
finally:
try:
temporary.unlink()
except FileNotFoundError:
pass
counts: dict[str, int] = {}
for row in rows:
counts[row["status"]] = counts.get(row["status"], 0) + 1
atomic_json(output.with_suffix(output.suffix + ".json"), {
"format": FINALIZE_FORMAT,
"finished_utc": utc_now(),
"manifest_sha256": canonical_sha256(manifest),
"run_signature_sha256": signature["signature_sha256"],
"tasks": len(rows),
"status_counts": counts,
})
log("finalized {} tasks -> {}".format(len(rows), output))
return output
def status(manifest: Mapping[str, Any]) -> None:
output = manifest_paths(manifest)["output_directory"]
assert output is not None
counts: dict[str, int] = {}
completed = 0
for marker in output.glob("block_*.complete"):
summary = output / (marker.name[:-len(".complete")] + ".summary.tsv")
with summary.open("r", encoding="utf-8-sig", newline="") as handle:
for row in csv.DictReader(handle, delimiter="\t"):
completed += 1
counts[row.get("status", "unknown")] = counts.get(row.get("status", "unknown"), 0) + 1
print(json.dumps({
"tasks": manifest["dimensions"]["task_count"],
"completed": completed,
"status_counts": counts,
"output_directory": str(output),
}, ensure_ascii=False, indent=2))
def positive_integer(value: str) -> int:
parsed = int(value)
if parsed <= 0:
raise argparse.ArgumentTypeError("must be positive")
return parsed
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("--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)
parser.add_argument("--line-search-max-zoom", type=positive_integer, default=48)
parser.add_argument("--gradient-abs-tol", type=float, default=1e-7)
parser.add_argument("--gradient-rel-tol", type=float, default=1e-8)
parser.add_argument("--step-rel-tol", type=float, default=1e-9)
parser.add_argument("--likelihood-rel-tol", type=float, default=1e-11)
parser.add_argument("--wolfe-c1", type=float, default=1e-4)
parser.add_argument("--wolfe-c2", type=float, default=0.9)
parser.add_argument("--initial-step", type=float, default=1.0)
parser.add_argument("--maximum-step", type=float, default=64.0)
parser.add_argument("--rank-tol", type=float, default=1e-10)
parser.add_argument("--covariance-floor", type=float, default=1e-12)
parser.add_argument("--boundary-h2-trigger", type=float, default=1e-3)
parser.add_argument("--boundary-score-tol", type=float, default=1e-8)
parser.add_argument("--boundary-logl-rel-tol", type=float, default=1e-12)
parser.add_argument("--initial-sigma-e", type=float)
parser.add_argument("--initial-sigma-g", type=float)
mode = parser.add_mutually_exclusive_group()
mode.add_argument("--resume", action="store_true")
mode.add_argument("--force", action="store_true")
parser.add_argument("--dry-run", action="store_true")
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description=__doc__)
subparsers = parser.add_subparsers(dest="action", required=True)
make = subparsers.add_parser("make-manifest")
make.add_argument("--manifest", type=Path, required=True)
make.add_argument("--grm-bin", type=Path, required=True)
make.add_argument("--grm-id", type=Path)
make.add_argument("--base-x", type=Path, required=True)
make.add_argument("--phenotypes", type=Path, required=True)
make.add_argument("--extra-covariates", type=Path)
make.add_argument("--tasks", type=Path, required=True)
make.add_argument("--extra-offsets", type=Path, required=True)
make.add_argument("--extra-indices", type=Path, required=True)
make.add_argument("--output-dir", type=Path, required=True)
make.add_argument("--n-samples", type=positive_integer, required=True)
make.add_argument("--n-base-covariates", type=positive_integer, required=True)
make.add_argument("--n-phenotype-rows", type=positive_integer, required=True)
make.add_argument("--n-extra-covariate-rows", type=int, default=0)
for action in ("validate", "status"):
command = subparsers.add_parser(action)
command.add_argument("--manifest", type=Path, required=True)
finalize_parser = subparsers.add_parser("finalize")
finalize_parser.add_argument("--manifest", type=Path, required=True)
finalize_parser.add_argument("--output", type=Path, required=True)
for action in ("run", "all"):
command = subparsers.add_parser(action)
command.add_argument("--manifest", type=Path, required=True)
if action == "all":
command.add_argument("--output", type=Path, required=True)
add_run_options(command)
return parser
def main(argv: Sequence[str] | None = None) -> int:
args = build_parser().parse_args(argv)
if args.action == "make-manifest":
make_manifest(args)
return 0
manifest = validate_manifest(args.manifest)
if args.action == "validate":
print(json.dumps({
"manifest": str(resolved(args.manifest)),
"manifest_sha256": canonical_sha256(manifest),
"dimensions": manifest["dimensions"],
}, ensure_ascii=False, indent=2))
elif args.action == "status":
status(manifest)
elif args.action == "finalize":
finalize(manifest, args.output)
elif args.action == "run":
run_engine(manifest, args)
elif args.action == "all":
run_engine(manifest, args)
if not args.dry_run:
finalize(manifest, args.output)
return 0
if __name__ == "__main__":
try:
raise SystemExit(main())
except KeyboardInterrupt:
raise SystemExit(130)
except Exception as exc:
print("ERROR: {}".format(exc), file=sys.stderr)
raise SystemExit(1)

112
scripts/run_server.sh Normal file
View File

@@ -0,0 +1,112 @@
#!/usr/bin/env bash
set -euo pipefail
ACTION="${1:-}"
if [[ ! "$ACTION" =~ ^(build|test|validate|run|dry-run|finalize|status|all)$ ]]; then
echo "Usage: bash $0 {build|test|validate|run|dry-run|finalize|status|all}" >&2
exit 2
fi
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_DIR="${PROJECT_DIR:-$(cd -- "$SCRIPT_DIR/.." && pwd)}"
PYTHON="${PYTHON:-python3}"
BACKEND="${BACKEND:-MKL}"
BACKEND="${BACKEND^^}"
BUILD_DIR="${BUILD_DIR:-$PROJECT_DIR/build-${BACKEND,,}}"
ENGINE="${ENGINE:-$BUILD_DIR/spectra_reml}"
MANIFEST="${MANIFEST:-}"
OUTPUT="${OUTPUT:-}"
THREADS="${THREADS:-1}"
BLAS_THREADS="${BLAS_THREADS:-1}"
BLOCK_SIZE="${BLOCK_SIZE:-256}"
RESUME="${RESUME:-1}"
FORCE="${FORCE:-0}"
CLI="$PROJECT_DIR/python/spectra_reml.py"
die() {
echo "ERROR: $*" >&2
exit 1
}
[[ "$BACKEND" == "MKL" || "$BACKEND" == "OPENBLAS" ]] || die "BACKEND must be MKL or OPENBLAS"
[[ "$THREADS" =~ ^[1-9][0-9]*$ ]] || die "THREADS must be positive"
[[ "$BLAS_THREADS" =~ ^[1-9][0-9]*$ ]] || die "BLAS_THREADS must be positive"
[[ "$BLOCK_SIZE" =~ ^[1-9][0-9]*$ ]] || die "BLOCK_SIZE must be positive"
[[ "$RESUME" == "0" || "$RESUME" == "1" ]] || die "RESUME must be 0 or 1"
[[ "$FORCE" == "0" || "$FORCE" == "1" ]] || die "FORCE must be 0 or 1"
[[ "$RESUME" != "1" || "$FORCE" != "1" ]] || die "RESUME and FORCE are mutually exclusive"
build_project() {
local args=(
-S "$PROJECT_DIR"
-B "$BUILD_DIR"
-DCMAKE_BUILD_TYPE=Release
-DREML_BLAS="$BACKEND"
-DREML_ENABLE_OPENMP=ON
)
if [[ "$BACKEND" == "MKL" ]]; then
args+=(
-DMKL_INTERFACE=lp64
-DMKL_LINK=dynamic
-DMKL_THREADING=sequential
)
fi
cmake "${args[@]}"
cmake --build "$BUILD_DIR" --parallel
}
if [[ "$ACTION" == "build" ]]; then
build_project
exit 0
fi
command -v "$PYTHON" >/dev/null 2>&1 || die "Python was not found: $PYTHON"
[[ -s "$CLI" ]] || die "Generic CLI is missing: $CLI"
if [[ "$ACTION" == "test" ]]; then
[[ -d "$BUILD_DIR" ]] || die "Build first: bash $0 build"
ctest --test-dir "$BUILD_DIR" --output-on-failure
"$PYTHON" "$PROJECT_DIR/tests/test_math_reference.py"
"$PYTHON" "$PROJECT_DIR/tests/python/test_spectra_reml_cli.py"
exit 0
fi
[[ -n "$MANIFEST" ]] || die "Set MANIFEST=/absolute/path/to/manifest.json"
case "$ACTION" in
validate|status)
"$PYTHON" "$CLI" "$ACTION" --manifest "$MANIFEST"
;;
finalize)
[[ -n "$OUTPUT" ]] || die "Set OUTPUT=/absolute/path/to/results.tsv.gz"
"$PYTHON" "$CLI" finalize --manifest "$MANIFEST" --output "$OUTPUT"
;;
run|dry-run|all)
if [[ "$ACTION" != "dry-run" ]]; then
[[ -x "$ENGINE" ]] || die "Engine is missing/not executable: $ENGINE"
fi
mode=()
if [[ "$FORCE" == "1" ]]; then
mode+=(--force)
elif [[ "$RESUME" == "1" ]]; then
mode+=(--resume)
fi
common=(
--manifest "$MANIFEST"
--engine "$ENGINE"
--threads "$THREADS"
--blas-threads "$BLAS_THREADS"
--block-size "$BLOCK_SIZE"
"${mode[@]}"
)
if [[ "$ACTION" == "dry-run" ]]; then
"$PYTHON" "$CLI" run "${common[@]}" --dry-run
elif [[ "$ACTION" == "all" ]]; then
[[ -n "$OUTPUT" ]] || die "Set OUTPUT=/absolute/path/to/results.tsv.gz"
"$PYTHON" "$CLI" all "${common[@]}" --output "$OUTPUT"
else
"$PYTHON" "$CLI" run "${common[@]}"
fi
;;
esac

711
src/batch_io.cpp Normal file
View File

@@ -0,0 +1,711 @@
#include "spectra_reml/batch_io.hpp"
#include "spectra_reml/linalg.hpp"
#include "spectra_reml/reml.hpp"
#include <algorithm>
#include <array>
#include <chrono>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <limits>
#include <set>
#include <sstream>
#include <stdexcept>
#include <string>
#include <thread>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>
#ifdef _OPENMP
#include <omp.h>
#endif
namespace spectra::reml {
namespace {
constexpr const char* kOutputFormatVersion = "spectra-reml-block-v1";
std::size_t checked_product(std::size_t left, std::size_t right,
const char* description) {
if (left != 0 && right > std::numeric_limits<std::size_t>::max() / left) {
throw std::overflow_error(std::string(description) + " size overflow");
}
return left * right;
}
void require_little_endian_ieee754() {
static_assert(sizeof(float) == 4, "batch format requires 32-bit float");
static_assert(sizeof(double) == 8, "batch format requires 64-bit double");
const std::uint32_t marker = 0x01020304U;
const auto* bytes = reinterpret_cast<const unsigned char*>(&marker);
if (bytes[0] != 0x04U || !std::numeric_limits<float>::is_iec559 ||
!std::numeric_limits<double>::is_iec559) {
throw std::runtime_error(
"binary batch format requires little-endian IEEE-754 hardware");
}
}
void validate_file_size(const std::filesystem::path& path,
std::uintmax_t expected_bytes,
const char* description) {
std::error_code error;
const auto actual_bytes = std::filesystem::file_size(path, error);
if (error) {
throw std::runtime_error("cannot stat " + std::string(description) +
" file " + path.string() + ": " +
error.message());
}
if (actual_bytes != expected_bytes) {
std::ostringstream message;
message << description << " file size mismatch for " << path.string()
<< ": expected " << expected_bytes << " bytes, found "
<< actual_bytes;
throw std::runtime_error(message.str());
}
}
template <typename T>
std::vector<T> read_exact_binary_vector(const std::filesystem::path& path,
std::size_t count,
const char* description) {
const std::size_t bytes = checked_product(count, sizeof(T), description);
validate_file_size(path, static_cast<std::uintmax_t>(bytes), description);
std::vector<T> values(count);
std::ifstream input(path, std::ios::binary);
if (!input) {
throw std::runtime_error("cannot open " + std::string(description) +
" file " + path.string());
}
if (bytes != 0) {
input.read(reinterpret_cast<char*>(values.data()),
static_cast<std::streamsize>(bytes));
}
if (!input || input.gcount() != static_cast<std::streamsize>(bytes)) {
throw std::runtime_error("short read from " + std::string(description) +
" file " + path.string());
}
return values;
}
std::vector<std::string> split_tab(const std::string& line) {
std::vector<std::string> fields;
std::size_t begin = 0;
while (true) {
const std::size_t end = line.find('\t', begin);
fields.push_back(line.substr(begin, end - begin));
if (end == std::string::npos) {
break;
}
begin = end + 1;
}
return fields;
}
std::size_t find_required_column(const std::vector<std::string>& header,
const std::string& name) {
const auto found = std::find(header.begin(), header.end(), name);
if (found == header.end()) {
throw std::runtime_error("tasks TSV is missing required column " + name);
}
return static_cast<std::size_t>(std::distance(header.begin(), found));
}
std::uint64_t parse_u64(const std::string& text, const char* field,
std::size_t line_number) {
if (text.empty() || text.front() == '-') {
std::ostringstream message;
message << "invalid " << field << " at tasks TSV line " << line_number;
throw std::runtime_error(message.str());
}
std::size_t consumed = 0;
unsigned long long value = 0;
try {
value = std::stoull(text, &consumed, 10);
} catch (const std::exception&) {
std::ostringstream message;
message << "invalid " << field << " at tasks TSV line " << line_number;
throw std::runtime_error(message.str());
}
if (consumed != text.size()) {
std::ostringstream message;
message << "invalid " << field << " at tasks TSV line " << line_number;
throw std::runtime_error(message.str());
}
return static_cast<std::uint64_t>(value);
}
std::string sanitize_tsv(std::string text) {
for (char& character : text) {
if (character == '\t' || character == '\n' || character == '\r') {
character = ' ';
}
}
return text;
}
ColMajorMatrix read_selected_f32_rows(
const std::filesystem::path& path, std::size_t total_rows,
std::size_t sample_count, const std::vector<std::int32_t>& selected_rows) {
const std::size_t element_count =
checked_product(total_rows, sample_count, "extra_covariates matrix");
const std::size_t byte_count =
checked_product(element_count, sizeof(float), "extra_covariates matrix");
validate_file_size(path, byte_count, "extra_covariates matrix");
ColMajorMatrix selected(sample_count, selected_rows.size());
if (selected_rows.empty()) {
return selected;
}
std::ifstream input(path, std::ios::binary);
if (!input) {
throw std::runtime_error("cannot open extra_covariates matrix " + path.string());
}
std::vector<float> buffer(sample_count);
const std::size_t row_bytes =
checked_product(sample_count, sizeof(float), "extra_covariates row");
for (std::size_t column = 0; column < selected_rows.size(); ++column) {
const auto row = selected_rows[column];
if (row < 0 || static_cast<std::size_t>(row) >= total_rows) {
throw std::runtime_error("extra_covariates row index is outside the matrix");
}
const std::uintmax_t offset =
static_cast<std::uintmax_t>(row) * row_bytes;
input.seekg(static_cast<std::streamoff>(offset), std::ios::beg);
input.read(reinterpret_cast<char*>(buffer.data()),
static_cast<std::streamsize>(row_bytes));
if (!input || input.gcount() != static_cast<std::streamsize>(row_bytes)) {
throw std::runtime_error("short read from extra_covariates row " +
std::to_string(row));
}
for (std::size_t sample = 0; sample < sample_count; ++sample) {
const double value = static_cast<double>(buffer[sample]);
if (!std::isfinite(value)) {
throw std::runtime_error("non-finite extra_covariates at row " +
std::to_string(row) + ", sample " +
std::to_string(sample));
}
selected(sample, column) = value;
}
input.clear();
}
return selected;
}
ColMajorMatrix read_selected_f64_rows(
std::ifstream& input, const std::filesystem::path& path,
std::size_t total_rows, std::size_t sample_count,
const std::vector<RemlTask>& tasks, std::size_t begin, std::size_t end) {
if (end < begin || end > tasks.size()) {
throw std::invalid_argument("invalid phenotypes block bounds");
}
ColMajorMatrix selected(sample_count, end - begin);
const std::size_t row_bytes =
checked_product(sample_count, sizeof(double), "phenotypes row");
for (std::size_t column = 0; column < end - begin; ++column) {
const auto row = tasks[begin + column].phenotype_row;
if (row >= total_rows) {
throw std::runtime_error("phenotypes row index is outside the matrix");
}
const std::uintmax_t offset = row * row_bytes;
input.seekg(static_cast<std::streamoff>(offset), std::ios::beg);
input.read(reinterpret_cast<char*>(selected.data() + column * sample_count),
static_cast<std::streamsize>(row_bytes));
if (!input || input.gcount() != static_cast<std::streamsize>(row_bytes)) {
throw std::runtime_error("short read from phenotypes file " +
path.string() + " at row " +
std::to_string(row));
}
input.clear();
}
return selected;
}
std::string block_stem(std::size_t block_index) {
std::ostringstream name;
name << "block_" << std::setw(6) << std::setfill('0') << block_index;
return name.str();
}
std::filesystem::path temporary_path(const std::filesystem::path& final_path) {
const auto clock_value =
std::chrono::high_resolution_clock::now().time_since_epoch().count();
const auto thread_hash =
std::hash<std::thread::id>{}(std::this_thread::get_id());
return final_path.string() + ".tmp." + std::to_string(clock_value) + "." +
std::to_string(thread_hash);
}
void rename_checked(const std::filesystem::path& source,
const std::filesystem::path& destination) {
std::error_code error;
std::filesystem::rename(source, destination, error);
if (error) {
throw std::runtime_error("cannot atomically rename " + source.string() +
" to " + destination.string() + ": " +
error.message());
}
}
void remove_if_exists(const std::filesystem::path& path) {
std::error_code error;
std::filesystem::remove(path, error);
if (error) {
throw std::runtime_error("cannot remove existing output " + path.string() +
": " + error.message());
}
}
void write_binary_doubles(const std::filesystem::path& path,
const std::vector<double>& values) {
std::ofstream output(path, std::ios::binary | std::ios::trunc);
if (!output) {
throw std::runtime_error("cannot create output " + path.string());
}
if (!values.empty()) {
output.write(reinterpret_cast<const char*>(values.data()),
static_cast<std::streamsize>(values.size() * sizeof(double)));
}
output.flush();
if (!output) {
throw std::runtime_error("failed writing output " + path.string());
}
}
void write_block_atomic(const std::filesystem::path& output_directory,
std::size_t block_index,
const std::vector<TaskResult>& results,
std::size_t base_covariate_count,
bool overwrite) {
const std::string stem = block_stem(block_index);
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 complete_path = output_directory / (stem + ".complete");
const std::array<std::filesystem::path, 4> final_paths = {
summary_path, beta_path, covariance_path, complete_path};
for (const auto& path : final_paths) {
if (std::filesystem::exists(path)) {
if (!overwrite) {
throw std::runtime_error("output already exists: " + path.string());
}
remove_if_exists(path);
}
}
std::vector<double> beta_values;
std::vector<double> covariance_values;
std::vector<std::int64_t> beta_offsets(results.size(), -1);
std::vector<std::int64_t> covariance_offsets(results.size(), -1);
for (std::size_t index = 0; index < results.size(); ++index) {
const auto& fit = results[index].fit;
if (!fit.has_estimates()) {
continue;
}
beta_offsets[index] = static_cast<std::int64_t>(beta_values.size());
covariance_offsets[index] =
static_cast<std::int64_t>(covariance_values.size());
beta_values.insert(beta_values.end(), fit.beta.begin(), fit.beta.end());
covariance_values.insert(covariance_values.end(),
fit.beta_covariance_packed_lower.begin(),
fit.beta_covariance_packed_lower.end());
}
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 complete_temp = temporary_path(complete_path);
const std::array<std::filesystem::path, 4> temporary_paths = {
summary_temp, beta_temp, covariance_temp, complete_temp};
try {
write_binary_doubles(beta_temp, beta_values);
write_binary_doubles(covariance_temp, covariance_values);
{
std::ofstream summary(summary_temp, std::ios::trunc);
if (!summary) {
throw std::runtime_error("cannot create output " +
summary_temp.string());
}
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";
summary << std::setprecision(17);
for (std::size_t index = 0; index < results.size(); ++index) {
const auto& item = results[index];
const auto& fit = item.fit;
const std::size_t expected_fixed =
fit.beta.empty() ? base_covariate_count +
item.task.extra_covariate_rows.size()
: fit.beta.size();
summary << item.task.task_index << '\t'
<< sanitize_tsv(item.task.task_id) << '\t'
<< to_string(fit.status) << '\t' << expected_fixed << '\t'
<< item.task.extra_covariate_rows.size() << '\t'
<< beta_offsets[index] << '\t' << covariance_offsets[index]
<< '\t' << fit.sigma_g2 << '\t' << fit.sigma_e2 << '\t'
<< fit.h2 << '\t' << fit.log_likelihood << '\t'
<< fit.iterations << '\t' << fit.line_search_evaluations
<< '\t' << fit.gradient_inf_norm << '\t'
<< sanitize_tsv(fit.error) << '\n';
}
summary.flush();
if (!summary) {
throw std::runtime_error("failed writing output " +
summary_temp.string());
}
}
{
std::ofstream complete(complete_temp, std::ios::trunc);
if (!complete) {
throw std::runtime_error("cannot create completion marker " +
complete_temp.string());
}
complete << "format\t" << kOutputFormatVersion << '\n'
<< "block\t" << block_index << '\n'
<< "tasks\t" << results.size() << '\n'
<< "beta_elements\t" << beta_values.size() << '\n'
<< "cov_elements\t" << covariance_values.size() << '\n';
complete.flush();
if (!complete) {
throw std::runtime_error("failed writing completion marker " +
complete_temp.string());
}
}
rename_checked(beta_temp, beta_path);
rename_checked(covariance_temp, covariance_path);
rename_checked(summary_temp, summary_path);
rename_checked(complete_temp, complete_path);
} catch (...) {
for (const auto& path : temporary_paths) {
std::error_code ignored;
std::filesystem::remove(path, ignored);
}
throw;
}
}
bool block_is_complete(const std::filesystem::path& output_directory,
std::size_t block_index) {
return std::filesystem::exists(
output_directory / (block_stem(block_index) + ".complete"));
}
} // namespace
ColMajorMatrix read_row_major_f64_matrix(const std::filesystem::path& path,
std::size_t rows,
std::size_t cols) {
require_little_endian_ieee754();
const std::size_t element_count =
checked_product(rows, cols, "float64 matrix");
const auto row_major =
read_exact_binary_vector<double>(path, element_count, "float64 matrix");
ColMajorMatrix result(rows, cols);
for (std::size_t row = 0; row < rows; ++row) {
for (std::size_t col = 0; col < cols; ++col) {
const double value = row_major[row * cols + col];
if (!std::isfinite(value)) {
throw std::runtime_error("non-finite value in float64 matrix " +
path.string());
}
result(row, col) = value;
}
}
return result;
}
std::vector<RemlTask> read_tasks(
const std::filesystem::path& tasks_tsv,
const std::filesystem::path& extra_offsets_i64,
const std::filesystem::path& extra_indices_i32,
const BatchDimensions& dimensions) {
require_little_endian_ieee754();
std::ifstream input(tasks_tsv);
if (!input) {
throw std::runtime_error("cannot open tasks TSV " + tasks_tsv.string());
}
std::string line;
if (!std::getline(input, line)) {
throw std::runtime_error("tasks TSV is empty");
}
if (!line.empty() && line.back() == '\r') {
line.pop_back();
}
const auto header = split_tab(line);
const auto task_column = find_required_column(header, "task_index");
const auto task_id_column = find_required_column(header, "task_id");
const auto phenotype_column =
find_required_column(header, "phenotype_row");
const auto extra_count_column = find_required_column(header, "n_extra_covariates");
const std::size_t required_field_count =
std::max({task_column, task_id_column, phenotype_column, extra_count_column}) +
1;
struct TaskRow {
RemlTask task;
std::size_t declared_extra_count = 0;
};
std::vector<TaskRow> task_rows;
std::unordered_set<std::uint64_t> task_indices;
std::size_t line_number = 1;
while (std::getline(input, line)) {
++line_number;
if (!line.empty() && line.back() == '\r') {
line.pop_back();
}
if (line.empty()) {
continue;
}
const auto fields = split_tab(line);
if (fields.size() < required_field_count) {
throw std::runtime_error("too few fields at tasks TSV line " +
std::to_string(line_number));
}
TaskRow row;
row.task.task_index =
parse_u64(fields[task_column], "task_index", line_number);
row.task.task_id = fields[task_id_column];
row.task.phenotype_row =
parse_u64(fields[phenotype_column], "phenotype_row", line_number);
row.declared_extra_count = static_cast<std::size_t>(
parse_u64(fields[extra_count_column], "n_extra_covariates", line_number));
if (row.task.task_id.empty()) {
throw std::runtime_error("empty task_id at tasks TSV line " +
std::to_string(line_number));
}
if (row.task.phenotype_row >= dimensions.phenotype_row_count) {
throw std::runtime_error("phenotype_row outside matrix at tasks TSV line " +
std::to_string(line_number));
}
if (!task_indices.insert(row.task.task_index).second) {
throw std::runtime_error("duplicate task_index in tasks TSV: " +
std::to_string(row.task.task_index));
}
task_rows.push_back(std::move(row));
}
if (!input.eof()) {
throw std::runtime_error("failed while reading tasks TSV " +
tasks_tsv.string());
}
const auto offsets = read_exact_binary_vector<std::int64_t>(
extra_offsets_i64, task_rows.size() + 1, "extra-covariate offsets");
if (offsets.empty() || offsets.front() != 0) {
throw std::runtime_error("extra-covariate offsets must start at zero");
}
for (std::size_t i = 1; i < offsets.size(); ++i) {
if (offsets[i] < offsets[i - 1]) {
throw std::runtime_error("extra-covariate offsets must be non-decreasing");
}
}
if (offsets.back() < 0) {
throw std::runtime_error("final extra-covariate offset is negative");
}
const auto extra_indices = read_exact_binary_vector<std::int32_t>(
extra_indices_i32, static_cast<std::size_t>(offsets.back()), "extra-covariate indices");
std::vector<RemlTask> tasks;
tasks.reserve(task_rows.size());
for (std::size_t i = 0; i < task_rows.size(); ++i) {
const auto begin = static_cast<std::size_t>(offsets[i]);
const auto end = static_cast<std::size_t>(offsets[i + 1]);
if (end - begin != task_rows[i].declared_extra_count) {
throw std::runtime_error("n_extra_covariates disagrees with offsets for task " +
std::to_string(task_rows[i].task.task_index));
}
std::unordered_set<std::int32_t> within_task;
for (std::size_t offset = begin; offset < end; ++offset) {
const auto row = extra_indices[offset];
if (row < 0 ||
static_cast<std::size_t>(row) >= dimensions.extra_covariate_row_count) {
throw std::runtime_error("extra-covariate row outside extra_covariates matrix for task " +
std::to_string(
task_rows[i].task.task_index));
}
if (!within_task.insert(row).second) {
throw std::runtime_error("duplicate extra covariate within task " +
std::to_string(
task_rows[i].task.task_index));
}
task_rows[i].task.extra_covariate_rows.push_back(row);
}
tasks.push_back(std::move(task_rows[i].task));
}
return tasks;
}
void run_task_batch(const BatchInputPaths& paths,
const BatchDimensions& dimensions,
const BatchOptions& options) {
require_little_endian_ieee754();
if (dimensions.sample_count == 0 ||
dimensions.base_covariate_count == 0 || options.block_size == 0 ||
options.outer_threads == 0) {
throw std::invalid_argument(
"sample, covariate, block and thread counts must be positive");
}
if (options.resume && options.overwrite) {
throw std::invalid_argument("--resume and --overwrite are mutually exclusive");
}
std::error_code directory_error;
std::filesystem::create_directories(paths.output_directory, directory_error);
if (directory_error) {
throw std::runtime_error("cannot create output directory " +
paths.output_directory.string() + ": " +
directory_error.message());
}
if (!paths.grm_id.empty()) {
validate_grm_id_count(paths.grm_id, dimensions.sample_count);
}
const auto tasks = read_tasks(paths.tasks_tsv, paths.extra_offsets_i64,
paths.extra_indices_i32, dimensions);
if (tasks.empty()) {
throw std::runtime_error("tasks TSV contains no tasks");
}
std::cerr << "Reading and diagonalizing the GRM for "
<< dimensions.sample_count << " samples...\n";
const SpectralGrm spectral = read_and_diagonalize_gcta_grm(
paths.grm_bin, dimensions.sample_count);
std::cerr << "GRM eigenvalue range: [" << std::setprecision(17)
<< spectral.minimum_eigenvalue << ", "
<< spectral.maximum_eigenvalue << "]\n";
const ColMajorMatrix base_x = read_row_major_f64_matrix(
paths.base_x_f64, dimensions.sample_count,
dimensions.base_covariate_count);
const ColMajorMatrix base_x_star =
rotate_to_eigenspace(spectral.eigenvectors, base_x);
std::set<std::int32_t> unique_extra_set;
for (const auto& task : tasks) {
unique_extra_set.insert(task.extra_covariate_rows.begin(), task.extra_covariate_rows.end());
}
const std::vector<std::int32_t> unique_extra_rows(unique_extra_set.begin(),
unique_extra_set.end());
std::unordered_map<std::int32_t, std::size_t> extra_to_column;
extra_to_column.reserve(unique_extra_rows.size());
for (std::size_t column = 0; column < unique_extra_rows.size(); ++column) {
extra_to_column.emplace(unique_extra_rows[column], column);
}
ColMajorMatrix extra_covariate_star(dimensions.sample_count, 0);
if (!unique_extra_rows.empty()) {
if (paths.extra_covariate_f32.empty()) {
throw std::runtime_error(
"extra_covariates file is required because tasks contain extra covariates");
}
std::cerr << "Reading and rotating " << unique_extra_rows.size()
<< " unique extra covariate rows...\n";
const ColMajorMatrix extra_covariates = read_selected_f32_rows(
paths.extra_covariate_f32, dimensions.extra_covariate_row_count,
dimensions.sample_count, unique_extra_rows);
extra_covariate_star =
rotate_to_eigenspace(spectral.eigenvectors, extra_covariates);
}
const std::size_t phenotype_elements = checked_product(
dimensions.phenotype_row_count, dimensions.sample_count,
"phenotypes matrix");
const std::size_t phenotype_bytes = checked_product(
phenotype_elements, sizeof(double), "phenotypes matrix");
validate_file_size(paths.phenotype_f64, phenotype_bytes,
"phenotypes matrix");
std::ifstream phenotype_input(paths.phenotype_f64, std::ios::binary);
if (!phenotype_input) {
throw std::runtime_error("cannot open phenotypes matrix " +
paths.phenotype_f64.string());
}
#ifdef _OPENMP
if (options.outer_threads >
static_cast<std::size_t>(std::numeric_limits<int>::max())) {
throw std::invalid_argument("thread count exceeds OpenMP integer range");
}
omp_set_dynamic(0);
omp_set_num_threads(static_cast<int>(options.outer_threads));
#else
if (options.outer_threads != 1) {
std::cerr << "Warning: binary was built without OpenMP; using one outer "
"thread.\n";
}
#endif
const std::size_t block_count =
(tasks.size() + options.block_size - 1) / options.block_size;
for (std::size_t block = 0; block < block_count; ++block) {
if (options.resume && block_is_complete(paths.output_directory, block)) {
std::cerr << "Skipping completed " << block_stem(block) << "\n";
continue;
}
const std::size_t begin = block * options.block_size;
const std::size_t end =
std::min(tasks.size(), begin + options.block_size);
std::cerr << "Processing " << block_stem(block) << " (tasks " << begin
<< ".." << (end - 1) << ")...\n";
const ColMajorMatrix phenotypes = read_selected_f64_rows(
phenotype_input, paths.phenotype_f64,
dimensions.phenotype_row_count, dimensions.sample_count, tasks,
begin, end);
const ColMajorMatrix phenotype_star =
rotate_to_eigenspace(spectral.eigenvectors, phenotypes);
std::vector<TaskResult> results(end - begin);
#ifdef _OPENMP
#pragma omp parallel for schedule(dynamic)
#endif
for (std::int64_t local_signed = 0;
local_signed < static_cast<std::int64_t>(end - begin);
++local_signed) {
const auto local = static_cast<std::size_t>(local_signed);
const RemlTask& task = tasks[begin + local];
TaskResult item;
item.task = task;
try {
const std::size_t fixed_count =
dimensions.base_covariate_count + task.extra_covariate_rows.size();
ColMajorMatrix design(dimensions.sample_count, fixed_count);
std::copy(base_x_star.values().begin(),
base_x_star.values().end(), design.values().begin());
for (std::size_t extra = 0; extra < task.extra_covariate_rows.size(); ++extra) {
const auto found = extra_to_column.find(task.extra_covariate_rows[extra]);
if (found == extra_to_column.end()) {
throw std::runtime_error(
"internal error: rotated extra covariate is missing");
}
const double* source =
extra_covariate_star.data() +
found->second * dimensions.sample_count;
double* destination =
design.data() +
(dimensions.base_covariate_count + extra) *
dimensions.sample_count;
std::copy(source, source + dimensions.sample_count,
destination);
}
std::vector<double> phenotype(
phenotype_star.data() + local * dimensions.sample_count,
phenotype_star.data() +
(local + 1) * dimensions.sample_count);
item.fit = fit_ai_reml_spectral(
phenotype, design, spectral.eigenvalues, options.reml);
} catch (const std::exception& exception) {
item.fit.status = FitStatus::numerical_error;
item.fit.error = exception.what();
}
results[local] = std::move(item);
}
write_block_atomic(paths.output_directory, block, results,
dimensions.base_covariate_count,
options.overwrite || options.resume);
}
}
} // namespace spectra::reml

121
src/grm.cpp Normal file
View File

@@ -0,0 +1,121 @@
#include "spectra_reml/grm.hpp"
#include "spectra_reml/linalg.hpp"
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <fstream>
#include <limits>
#include <sstream>
#include <stdexcept>
#include <string>
#include <vector>
namespace spectra::reml {
namespace {
std::uintmax_t expected_grm_bytes(std::size_t n) {
if (n == 0) {
throw std::invalid_argument("sample count must be positive");
}
if (n > (std::numeric_limits<std::size_t>::max() / (n + 1))) {
throw std::overflow_error("sample count overflows GRM element count");
}
const std::size_t elements = n * (n + 1) / 2;
if (elements > std::numeric_limits<std::uintmax_t>::max() / sizeof(float)) {
throw std::overflow_error("GRM byte count overflow");
}
return static_cast<std::uintmax_t>(elements) * sizeof(float);
}
} // namespace
ColMajorMatrix read_gcta_grm_lower_triangle(
const std::filesystem::path& grm_bin, std::size_t sample_count) {
const auto expected_bytes = expected_grm_bytes(sample_count);
std::error_code file_error;
const auto actual_bytes = std::filesystem::file_size(grm_bin, file_error);
if (file_error) {
throw std::runtime_error("cannot stat GRM file " + grm_bin.string() +
": " + file_error.message());
}
if (actual_bytes != expected_bytes) {
std::ostringstream message;
message << "GRM file size mismatch for " << grm_bin.string()
<< ": expected " << expected_bytes << " bytes for "
<< sample_count << " samples, found " << actual_bytes;
throw std::runtime_error(message.str());
}
const std::size_t element_count = sample_count * (sample_count + 1) / 2;
std::vector<float> lower_triangle(element_count);
std::ifstream input(grm_bin, std::ios::binary);
if (!input) {
throw std::runtime_error("cannot open GRM file " + grm_bin.string());
}
input.read(reinterpret_cast<char*>(lower_triangle.data()),
static_cast<std::streamsize>(expected_bytes));
if (!input || input.gcount() != static_cast<std::streamsize>(expected_bytes)) {
throw std::runtime_error("short read from GRM file " + grm_bin.string());
}
ColMajorMatrix grm(sample_count, sample_count);
std::size_t source = 0;
for (std::size_t row = 0; row < sample_count; ++row) {
for (std::size_t col = 0; col <= row; ++col, ++source) {
const double value = static_cast<double>(lower_triangle[source]);
if (!std::isfinite(value)) {
std::ostringstream message;
message << "non-finite GRM value at (" << row << ',' << col
<< ')';
throw std::runtime_error(message.str());
}
grm(row, col) = value;
grm(col, row) = value;
}
}
return grm;
}
SpectralGrm read_and_diagonalize_gcta_grm(
const std::filesystem::path& grm_bin, std::size_t sample_count) {
SpectralGrm result;
result.eigenvectors =
read_gcta_grm_lower_triangle(grm_bin, sample_count);
result.eigenvalues =
symmetric_eigen_decomposition(result.eigenvectors);
result.minimum_eigenvalue = result.eigenvalues.front();
result.maximum_eigenvalue = result.eigenvalues.back();
return result;
}
void validate_grm_id_count(const std::filesystem::path& grm_id,
std::size_t expected_sample_count) {
std::ifstream input(grm_id);
if (!input) {
throw std::runtime_error("cannot open GRM ID file " + grm_id.string());
}
std::size_t count = 0;
std::string line;
while (std::getline(input, line)) {
if (!line.empty() && line.back() == '\r') {
line.pop_back();
}
if (line.find_first_not_of(" \t") != std::string::npos) {
++count;
}
}
if (!input.eof()) {
throw std::runtime_error("failed while reading GRM ID file " +
grm_id.string());
}
if (count != expected_sample_count) {
std::ostringstream message;
message << "GRM ID count mismatch: expected " << expected_sample_count
<< ", found " << count;
throw std::runtime_error(message.str());
}
}
} // namespace spectra::reml

201
src/linalg.cpp Normal file
View File

@@ -0,0 +1,201 @@
#include "spectra_reml/linalg.hpp"
#if defined(REML_BLAS_MKL)
#include <mkl_cblas.h>
#include <mkl_lapacke.h>
#else
#include <cblas.h>
#include <lapacke.h>
#endif
#include <algorithm>
#include <cmath>
#include <limits>
#include <sstream>
#include <stdexcept>
namespace spectra::reml {
namespace {
lapack_int checked_lapack_int(std::size_t value, const char* name) {
if (value > static_cast<std::size_t>(std::numeric_limits<lapack_int>::max())) {
throw std::overflow_error(std::string(name) + " exceeds LAPACK integer range");
}
return static_cast<lapack_int>(value);
}
int checked_blas_int(std::size_t value, const char* name) {
if (value > static_cast<std::size_t>(std::numeric_limits<int>::max())) {
throw std::overflow_error(std::string(name) + " exceeds CBLAS integer range");
}
return static_cast<int>(value);
}
} // namespace
ColMajorMatrix cross_product(const ColMajorMatrix& a,
const ColMajorMatrix& b) {
if (a.rows() != b.rows()) {
throw std::invalid_argument("cross_product requires equal row counts");
}
ColMajorMatrix result(a.cols(), b.cols());
cblas_dgemm(CblasColMajor, CblasTrans, CblasNoTrans,
checked_blas_int(a.cols(), "A columns"),
checked_blas_int(b.cols(), "B columns"),
checked_blas_int(a.rows(), "row count"), 1.0, a.data(),
checked_blas_int(a.leading_dimension(), "A leading dimension"),
b.data(),
checked_blas_int(b.leading_dimension(), "B leading dimension"),
0.0, result.data(),
checked_blas_int(result.leading_dimension(),
"result leading dimension"));
return result;
}
ColMajorMatrix rotate_to_eigenspace(const ColMajorMatrix& eigenvectors,
const ColMajorMatrix& variables) {
if (eigenvectors.rows() != eigenvectors.cols()) {
throw std::invalid_argument("eigenvector matrix must be square");
}
if (eigenvectors.rows() != variables.rows()) {
throw std::invalid_argument(
"eigenvectors and variables have inconsistent sample counts");
}
ColMajorMatrix rotated(variables.rows(), variables.cols());
cblas_dgemm(
CblasColMajor, CblasTrans, CblasNoTrans,
checked_blas_int(eigenvectors.cols(), "eigenvector column count"),
checked_blas_int(variables.cols(), "variable column count"),
checked_blas_int(eigenvectors.rows(), "sample count"), 1.0,
eigenvectors.data(),
checked_blas_int(eigenvectors.leading_dimension(),
"eigenvector leading dimension"),
variables.data(),
checked_blas_int(variables.leading_dimension(),
"variable leading dimension"),
0.0, rotated.data(),
checked_blas_int(rotated.leading_dimension(),
"rotated leading dimension"));
return rotated;
}
std::vector<double> symmetric_eigen_decomposition(
ColMajorMatrix& symmetric_matrix) {
if (symmetric_matrix.rows() != symmetric_matrix.cols()) {
throw std::invalid_argument("symmetric eigen decomposition requires a square matrix");
}
const auto n = checked_lapack_int(symmetric_matrix.rows(), "matrix order");
std::vector<double> eigenvalues(symmetric_matrix.rows());
const lapack_int info = LAPACKE_dsyevd(
LAPACK_COL_MAJOR, 'V', 'L', n, symmetric_matrix.data(), n,
eigenvalues.data());
if (info < 0) {
std::ostringstream message;
message << "LAPACKE_dsyevd rejected argument " << -info;
throw std::runtime_error(message.str());
}
if (info > 0) {
std::ostringstream message;
message << "LAPACKE_dsyevd failed to converge (info=" << info << ')';
throw std::runtime_error(message.str());
}
return eigenvalues;
}
bool cholesky_factor_in_place(std::vector<double>& matrix, std::size_t order,
std::string* error) {
if (matrix.size() != order * order) {
if (error != nullptr) {
*error = "Cholesky matrix has the wrong size";
}
return false;
}
const auto n = checked_lapack_int(order, "Cholesky order");
const lapack_int info =
LAPACKE_dpotrf(LAPACK_COL_MAJOR, 'L', n, matrix.data(), n);
if (info == 0) {
return true;
}
if (error != nullptr) {
std::ostringstream message;
if (info < 0) {
message << "LAPACKE_dpotrf rejected argument " << -info;
} else {
message << "matrix is not positive definite at leading minor "
<< info;
}
*error = message.str();
}
return false;
}
bool cholesky_solve_in_place(const std::vector<double>& factor,
std::size_t order, double* rhs,
std::size_t rhs_columns, std::string* error) {
if (factor.size() != order * order || rhs == nullptr) {
if (error != nullptr) {
*error = "invalid Cholesky solve dimensions";
}
return false;
}
const auto n = checked_lapack_int(order, "Cholesky order");
const auto nrhs = checked_lapack_int(rhs_columns, "right-hand-side count");
const lapack_int info = LAPACKE_dpotrs(
LAPACK_COL_MAJOR, 'L', n, nrhs, factor.data(), n, rhs, n);
if (info == 0) {
return true;
}
if (error != nullptr) {
std::ostringstream message;
message << "LAPACKE_dpotrs failed (info=" << info << ')';
*error = message.str();
}
return false;
}
bool cholesky_inverse(const std::vector<double>& factor, std::size_t order,
std::vector<double>& inverse, std::string* error) {
if (factor.size() != order * order) {
if (error != nullptr) {
*error = "invalid Cholesky inverse dimensions";
}
return false;
}
inverse = factor;
const auto n = checked_lapack_int(order, "Cholesky order");
const lapack_int info =
LAPACKE_dpotri(LAPACK_COL_MAJOR, 'L', n, inverse.data(), n);
if (info != 0) {
if (error != nullptr) {
std::ostringstream message;
message << "LAPACKE_dpotri failed (info=" << info << ')';
*error = message.str();
}
inverse.clear();
return false;
}
for (std::size_t col = 0; col < order; ++col) {
for (std::size_t row = 0; row < col; ++row) {
inverse[row + col * order] = inverse[col + row * order];
}
}
return true;
}
double dot(const std::vector<double>& a, const std::vector<double>& b) {
if (a.size() != b.size()) {
throw std::invalid_argument("dot product requires equal vector lengths");
}
return cblas_ddot(checked_blas_int(a.size(), "vector length"), a.data(), 1,
b.data(), 1);
}
double infinity_norm(const std::vector<double>& x) {
double result = 0.0;
for (const double value : x) {
result = std::max(result, std::abs(value));
}
return result;
}
} // namespace spectra::reml

354
src/main.cpp Normal file
View File

@@ -0,0 +1,354 @@
#include "spectra_reml/batch_io.hpp"
#include <cmath>
#include <cstddef>
#include <cstdlib>
#include <exception>
#include <filesystem>
#include <iostream>
#include <limits>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <unordered_set>
namespace {
using spectra::reml::BatchDimensions;
using spectra::reml::BatchInputPaths;
using spectra::reml::BatchOptions;
void print_usage(std::ostream& output) {
output
<< "spectra_reml - reusable spectral AI-REML batch engine\n\n"
<< "Required input paths:\n"
<< " --grm-bin PATH GCTA lower-triangle .grm.bin\n"
<< " --base-x PATH row-major n x p0 float64 matrix\n"
<< " --phenotypes PATH row-major rows x n float64 matrix\n"
<< " --tasks PATH TSV with task_index,task_id,phenotype_row,n_extra_covariates\n"
<< " --extra-offsets PATH int64 CSR offsets, m+1 entries\n"
<< " --extra-indices PATH int32 zero-based extra-covariate row indices\n"
<< " --out-dir PATH task-set output directory\n\n"
<< "Required dimensions:\n"
<< " --n-samples N\n"
<< " --n-base-covariates P\n"
<< " --n-phenotype-rows R\n"
<< " --n-extra-covariate-rows R\n\n"
<< "Conditionally required:\n"
<< " --extra-covariates PATH row-major rows x n float32 matrix; required\n"
<< " when any task uses extra covariates\n\n"
<< "Optional validation and scheduling:\n"
<< " --grm-id PATH verify .grm.id line count\n"
<< " --block-size N tasks per atomic block (default 256)\n"
<< " --threads N OpenMP task threads (default 1)\n"
<< " --resume skip blocks with .complete marker\n"
<< " --overwrite replace existing block files\n\n"
<< "Optional AI-REML controls:\n"
<< " --max-iterations N default 100\n"
<< " --line-search-max-evals N default 48\n"
<< " --line-search-max-zoom N default 48\n"
<< " --gradient-abs-tol X default 1e-7\n"
<< " --gradient-rel-tol X default 1e-8\n"
<< " --step-rel-tol X default 1e-9\n"
<< " --likelihood-rel-tol X default 1e-11\n"
<< " --wolfe-c1 X default 1e-4\n"
<< " --wolfe-c2 X default 0.9\n"
<< " --initial-step X default 1\n"
<< " --maximum-step X default 64\n"
<< " --initial-sigma-e X override OLS-based initial value\n"
<< " --initial-sigma-g X override OLS-based initial value\n"
<< " --rank-tol X default 1e-10\n"
<< " --covariance-floor X relative floor, default 1e-12\n"
<< " --boundary-h2-trigger X test v_g=0 once h2<=X (default 1e-3)\n"
<< " --boundary-score-tol X one-sided score tolerance (default 1e-8)\n"
<< " --boundary-logl-rel-tol X likelihood comparison tolerance (default 1e-12)\n"
<< " --help show this message\n\n"
<< "All binary inputs and outputs are little-endian. Offsets are counted in\n"
<< "elements, not bytes. BLAS thread counts are configured outside this CLI.\n";
}
struct Arguments {
std::unordered_map<std::string, std::string> values;
std::unordered_set<std::string> flags;
};
Arguments parse_arguments(int argc, char** argv) {
const std::unordered_set<std::string> flag_names = {
"--help", "--resume", "--overwrite"};
Arguments result;
for (int i = 1; i < argc; ++i) {
const std::string name = argv[i];
if (name.rfind("--", 0) != 0) {
throw std::invalid_argument("unexpected positional argument: " + name);
}
if (flag_names.count(name) != 0U) {
if (!result.flags.insert(name).second) {
throw std::invalid_argument("duplicate flag: " + name);
}
continue;
}
if (i + 1 >= argc) {
throw std::invalid_argument("missing value for " + name);
}
if (!result.values.emplace(name, argv[++i]).second) {
throw std::invalid_argument("duplicate option: " + name);
}
}
return result;
}
const std::string& require_value(const Arguments& arguments,
const std::string& name) {
const auto found = arguments.values.find(name);
if (found == arguments.values.end() || found->second.empty()) {
throw std::invalid_argument("missing required option " + name);
}
return found->second;
}
std::string optional_value(const Arguments& arguments, const std::string& name,
const std::string& fallback = {}) {
const auto found = arguments.values.find(name);
return found == arguments.values.end() ? fallback : found->second;
}
std::size_t parse_size(const std::string& text, const std::string& name,
bool allow_zero = false) {
if (text.empty() || text.front() == '-') {
throw std::invalid_argument("invalid integer for " + name + ": " + text);
}
std::size_t consumed = 0;
unsigned long long parsed = 0;
try {
parsed = std::stoull(text, &consumed, 10);
} catch (const std::exception&) {
throw std::invalid_argument("invalid integer for " + name + ": " + text);
}
if (consumed != text.size() ||
parsed > static_cast<unsigned long long>(
std::numeric_limits<std::size_t>::max()) ||
(!allow_zero && parsed == 0)) {
throw std::invalid_argument("invalid integer for " + name + ": " + text);
}
return static_cast<std::size_t>(parsed);
}
double parse_double(const std::string& text, const std::string& name,
bool require_positive = false,
bool allow_zero = false) {
std::size_t consumed = 0;
double parsed = 0.0;
try {
parsed = std::stod(text, &consumed);
} catch (const std::exception&) {
throw std::invalid_argument("invalid number for " + name + ": " + text);
}
if (consumed != text.size() || !std::isfinite(parsed) ||
(require_positive &&
(allow_zero ? parsed < 0.0 : parsed <= 0.0))) {
throw std::invalid_argument("invalid number for " + name + ": " + text);
}
return parsed;
}
template <typename Setter>
void set_if_present(const Arguments& arguments, const std::string& name,
Setter setter) {
const auto found = arguments.values.find(name);
if (found != arguments.values.end()) {
setter(found->second);
}
}
void reject_unknown_options(const Arguments& arguments) {
const std::unordered_set<std::string> known = {
"--grm-bin",
"--grm-id",
"--base-x",
"--phenotypes",
"--extra-covariates",
"--tasks",
"--extra-offsets",
"--extra-indices",
"--out-dir",
"--n-samples",
"--n-base-covariates",
"--n-phenotype-rows",
"--n-extra-covariate-rows",
"--block-size",
"--threads",
"--max-iterations",
"--line-search-max-evals",
"--line-search-max-zoom",
"--gradient-abs-tol",
"--gradient-rel-tol",
"--step-rel-tol",
"--likelihood-rel-tol",
"--wolfe-c1",
"--wolfe-c2",
"--initial-step",
"--maximum-step",
"--initial-sigma-e",
"--initial-sigma-g",
"--rank-tol",
"--covariance-floor",
"--boundary-h2-trigger",
"--boundary-score-tol",
"--boundary-logl-rel-tol"};
for (const auto& [name, value] : arguments.values) {
(void)value;
if (known.count(name) == 0U) {
throw std::invalid_argument("unknown option: " + name);
}
}
}
} // namespace
int main(int argc, char** argv) {
try {
const Arguments arguments = parse_arguments(argc, argv);
if (arguments.flags.count("--help") != 0U) {
print_usage(std::cout);
return EXIT_SUCCESS;
}
reject_unknown_options(arguments);
BatchInputPaths paths;
paths.grm_bin = require_value(arguments, "--grm-bin");
paths.grm_id = optional_value(arguments, "--grm-id");
paths.base_x_f64 = require_value(arguments, "--base-x");
paths.phenotype_f64 = require_value(arguments, "--phenotypes");
paths.extra_covariate_f32 = optional_value(arguments, "--extra-covariates");
paths.tasks_tsv = require_value(arguments, "--tasks");
paths.extra_offsets_i64 = require_value(arguments, "--extra-offsets");
paths.extra_indices_i32 = require_value(arguments, "--extra-indices");
paths.output_directory = require_value(arguments, "--out-dir");
BatchDimensions dimensions;
dimensions.sample_count = parse_size(
require_value(arguments, "--n-samples"), "--n-samples");
dimensions.base_covariate_count = parse_size(
require_value(arguments, "--n-base-covariates"),
"--n-base-covariates");
dimensions.phenotype_row_count = parse_size(
require_value(arguments, "--n-phenotype-rows"),
"--n-phenotype-rows");
dimensions.extra_covariate_row_count = parse_size(
require_value(arguments, "--n-extra-covariate-rows"),
"--n-extra-covariate-rows", true);
BatchOptions options;
options.resume = arguments.flags.count("--resume") != 0U;
options.overwrite = arguments.flags.count("--overwrite") != 0U;
set_if_present(arguments, "--block-size", [&](const std::string& value) {
options.block_size = parse_size(value, "--block-size");
});
set_if_present(arguments, "--threads", [&](const std::string& value) {
options.outer_threads = parse_size(value, "--threads");
});
set_if_present(arguments, "--max-iterations",
[&](const std::string& value) {
options.reml.max_iterations =
parse_size(value, "--max-iterations");
});
set_if_present(arguments, "--line-search-max-evals",
[&](const std::string& value) {
options.reml.line_search_max_evaluations =
parse_size(value, "--line-search-max-evals");
});
set_if_present(arguments, "--line-search-max-zoom",
[&](const std::string& value) {
options.reml.line_search_max_zoom_iterations =
parse_size(value, "--line-search-max-zoom");
});
set_if_present(arguments, "--gradient-abs-tol",
[&](const std::string& value) {
options.reml.gradient_absolute_tolerance =
parse_double(value, "--gradient-abs-tol", true,
true);
});
set_if_present(arguments, "--gradient-rel-tol",
[&](const std::string& value) {
options.reml.gradient_relative_tolerance =
parse_double(value, "--gradient-rel-tol", true,
true);
});
set_if_present(arguments, "--step-rel-tol",
[&](const std::string& value) {
options.reml.step_relative_tolerance =
parse_double(value, "--step-rel-tol", true, true);
});
set_if_present(arguments, "--likelihood-rel-tol",
[&](const std::string& value) {
options.reml.likelihood_relative_tolerance =
parse_double(value, "--likelihood-rel-tol", true,
true);
});
set_if_present(arguments, "--wolfe-c1", [&](const std::string& value) {
options.reml.wolfe_c1 = parse_double(value, "--wolfe-c1", true);
});
set_if_present(arguments, "--wolfe-c2", [&](const std::string& value) {
options.reml.wolfe_c2 = parse_double(value, "--wolfe-c2", true);
});
set_if_present(arguments, "--initial-step",
[&](const std::string& value) {
options.reml.initial_line_search_step =
parse_double(value, "--initial-step", true);
});
set_if_present(arguments, "--maximum-step",
[&](const std::string& value) {
options.reml.maximum_line_search_step =
parse_double(value, "--maximum-step", true);
});
set_if_present(arguments, "--initial-sigma-e",
[&](const std::string& value) {
options.reml.initial_sigma_e =
parse_double(value, "--initial-sigma-e");
});
set_if_present(arguments, "--initial-sigma-g",
[&](const std::string& value) {
options.reml.initial_sigma_g =
parse_double(value, "--initial-sigma-g");
});
set_if_present(arguments, "--rank-tol", [&](const std::string& value) {
options.reml.rank_tolerance_relative =
parse_double(value, "--rank-tol", true, true);
});
set_if_present(arguments, "--covariance-floor",
[&](const std::string& value) {
options.reml.covariance_floor_relative = parse_double(
value, "--covariance-floor", true, true);
});
set_if_present(arguments, "--boundary-h2-trigger",
[&](const std::string& value) {
options.reml.boundary_h2_trigger = parse_double(
value, "--boundary-h2-trigger", true, true);
if (options.reml.boundary_h2_trigger > 1.0) {
throw std::invalid_argument(
"--boundary-h2-trigger must be <= 1");
}
});
set_if_present(arguments, "--boundary-score-tol",
[&](const std::string& value) {
options.reml.boundary_score_tolerance = parse_double(
value, "--boundary-score-tol", true, true);
});
set_if_present(arguments, "--boundary-logl-rel-tol",
[&](const std::string& value) {
options.reml.boundary_likelihood_relative_tolerance =
parse_double(value, "--boundary-logl-rel-tol",
true, true);
});
spectra::reml::run_task_batch(paths, dimensions, options);
return EXIT_SUCCESS;
} catch (const std::invalid_argument& exception) {
std::cerr << "Argument error: " << exception.what() << "\n\n";
print_usage(std::cerr);
return 2;
} catch (const std::exception& exception) {
std::cerr << "Fatal error: " << exception.what() << '\n';
return EXIT_FAILURE;
}
}

842
src/reml_core.cpp Normal file
View File

@@ -0,0 +1,842 @@
#include "spectra_reml/reml.hpp"
#include "spectra_reml/linalg.hpp"
#include <algorithm>
#include <array>
#include <cmath>
#include <cstddef>
#include <limits>
#include <sstream>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>
namespace spectra::reml {
namespace {
constexpr double kTiny = 64.0 * std::numeric_limits<double>::epsilon();
struct ValidatedProblem {
const std::vector<double>& y;
const ColMajorMatrix& x;
const std::vector<double>& lambda;
std::size_t n;
std::size_t p;
};
ValidatedProblem validate_problem(const std::vector<double>& y,
const ColMajorMatrix& x,
const std::vector<double>& lambda) {
if (y.empty()) {
throw std::invalid_argument("phenotype is empty");
}
if (x.rows() != y.size()) {
throw std::invalid_argument("phenotype and design row counts differ");
}
if (lambda.size() != y.size()) {
throw std::invalid_argument("phenotype and eigenvalue counts differ");
}
if (x.cols() == 0) {
throw std::invalid_argument("design matrix has no columns");
}
if (x.cols() >= x.rows()) {
throw std::invalid_argument("REML requires more observations than fixed effects");
}
for (std::size_t i = 0; i < y.size(); ++i) {
if (!std::isfinite(y[i])) {
throw std::invalid_argument("phenotype contains a non-finite value");
}
if (!std::isfinite(lambda[i])) {
throw std::invalid_argument("GRM eigenvalues contain a non-finite value");
}
}
for (const double value : x.values()) {
if (!std::isfinite(value)) {
throw std::invalid_argument("design matrix contains a non-finite value");
}
}
return {y, x, lambda, y.size(), x.cols()};
}
void fill_symmetric_upper(std::vector<double>& matrix, std::size_t p) {
for (std::size_t col = 0; col < p; ++col) {
for (std::size_t row = 0; row < col; ++row) {
matrix[row + col * p] = matrix[col + row * p];
}
}
}
bool factor_is_numerically_full_rank(const std::vector<double>& factor,
std::size_t p,
double relative_tolerance) {
double min_diagonal = std::numeric_limits<double>::infinity();
double max_diagonal = 0.0;
for (std::size_t i = 0; i < p; ++i) {
const double diagonal = std::abs(factor[i + i * p]);
min_diagonal = std::min(min_diagonal, diagonal);
max_diagonal = std::max(max_diagonal, diagonal);
}
if (!(max_diagonal > 0.0) || !std::isfinite(min_diagonal)) {
return false;
}
const double ratio = min_diagonal / max_diagonal;
return ratio * ratio > relative_tolerance;
}
bool apply_p(const ValidatedProblem& problem,
const std::vector<double>& inverse_h,
const std::vector<double>& cholesky,
const std::vector<double>& input,
std::vector<double>& output, std::string& error) {
const auto n = problem.n;
const auto p = problem.p;
std::vector<double> rhs(p, 0.0);
for (std::size_t col = 0; col < p; ++col) {
const double* x_col = problem.x.data() + col * n;
double value = 0.0;
for (std::size_t i = 0; i < n; ++i) {
value += x_col[i] * inverse_h[i] * input[i];
}
rhs[col] = value;
}
if (!cholesky_solve_in_place(cholesky, p, rhs.data(), 1, &error)) {
return false;
}
output.resize(n);
for (std::size_t i = 0; i < n; ++i) {
double fitted = 0.0;
for (std::size_t col = 0; col < p; ++col) {
fitted += problem.x(i, col) * rhs[col];
}
output[i] = inverse_h[i] * (input[i] - fitted);
}
return true;
}
RemlEvaluation evaluate_impl(const ValidatedProblem& problem, double sigma_e,
double sigma_g, bool compute_ai,
bool compute_beta_covariance,
double covariance_floor_relative,
double rank_tolerance_relative) {
RemlEvaluation result;
if (!std::isfinite(sigma_e) || !std::isfinite(sigma_g)) {
result.error = "variance parameters are not finite";
return result;
}
const double sigma_e2 = sigma_e * sigma_e;
const double sigma_g2 = sigma_g * sigma_g;
if (!std::isfinite(sigma_e2) || !std::isfinite(sigma_g2) ||
!(sigma_e2 > 0.0 || sigma_g2 > 0.0)) {
result.error = "both variance components are zero or non-finite";
return result;
}
const auto n = problem.n;
const auto p = problem.p;
double covariance_scale = sigma_e2;
for (const double lambda : problem.lambda) {
covariance_scale =
std::max(covariance_scale, std::abs(sigma_g2 * lambda));
}
const double covariance_floor = std::max(
std::numeric_limits<double>::min(),
covariance_floor_relative * std::max(covariance_scale,
std::numeric_limits<double>::min()));
std::vector<double> inverse_h(n);
double log_determinant_h = 0.0;
double sum_inverse_h = 0.0;
double sum_lambda_inverse_h = 0.0;
for (std::size_t i = 0; i < n; ++i) {
const double h = sigma_e2 + sigma_g2 * problem.lambda[i];
if (!std::isfinite(h) || h <= covariance_floor) {
std::ostringstream message;
message << "non-positive spectral covariance at index " << i
<< " (value=" << h << ')';
result.error = message.str();
return result;
}
inverse_h[i] = 1.0 / h;
log_determinant_h += std::log(h);
sum_inverse_h += inverse_h[i];
sum_lambda_inverse_h += problem.lambda[i] * inverse_h[i];
}
std::vector<double> normal_matrix(p * p, 0.0);
std::vector<double> normal_rhs(p, 0.0);
for (std::size_t i = 0; i < n; ++i) {
const double weight = inverse_h[i];
for (std::size_t row = 0; row < p; ++row) {
const double x_row = problem.x(i, row);
normal_rhs[row] += weight * x_row * problem.y[i];
for (std::size_t col = 0; col <= row; ++col) {
normal_matrix[row + col * p] +=
weight * x_row * problem.x(i, col);
}
}
}
fill_symmetric_upper(normal_matrix, p);
std::vector<double> cholesky = normal_matrix;
std::string linear_algebra_error;
if (!cholesky_factor_in_place(cholesky, p, &linear_algebra_error)) {
result.error = "fixed-effect normal matrix is singular: " +
linear_algebra_error;
return result;
}
if (!factor_is_numerically_full_rank(cholesky, p,
rank_tolerance_relative)) {
result.error = "fixed-effect normal matrix is numerically rank deficient";
return result;
}
double log_determinant_normal = 0.0;
for (std::size_t i = 0; i < p; ++i) {
log_determinant_normal += 2.0 * std::log(cholesky[i + i * p]);
}
result.beta = normal_rhs;
if (!cholesky_solve_in_place(cholesky, p, result.beta.data(), 1,
&linear_algebra_error)) {
result.error = linear_algebra_error;
result.beta.clear();
return result;
}
std::vector<double> py(n);
double quadratic = 0.0;
for (std::size_t i = 0; i < n; ++i) {
double fitted = 0.0;
for (std::size_t col = 0; col < p; ++col) {
fitted += problem.x(i, col) * result.beta[col];
}
py[i] = inverse_h[i] * (problem.y[i] - fitted);
quadratic += problem.y[i] * py[i];
}
result.log_likelihood =
-0.5 * (log_determinant_normal + log_determinant_h + quadratic);
if (!std::isfinite(result.log_likelihood)) {
result.error = "restricted log-likelihood is not finite";
result.beta.clear();
return result;
}
std::vector<double> correction_e(p * p, 0.0);
std::vector<double> correction_g(p * p, 0.0);
for (std::size_t i = 0; i < n; ++i) {
const double weight2 = inverse_h[i] * inverse_h[i];
const double lambda_weight2 = problem.lambda[i] * weight2;
for (std::size_t row = 0; row < p; ++row) {
const double x_row = problem.x(i, row);
for (std::size_t col = 0; col <= row; ++col) {
const double product = x_row * problem.x(i, col);
correction_e[row + col * p] += weight2 * product;
correction_g[row + col * p] += lambda_weight2 * product;
}
}
}
fill_symmetric_upper(correction_e, p);
fill_symmetric_upper(correction_g, p);
if (!cholesky_solve_in_place(cholesky, p, correction_e.data(), p,
&linear_algebra_error) ||
!cholesky_solve_in_place(cholesky, p, correction_g.data(), p,
&linear_algebra_error)) {
result.error = linear_algebra_error;
result.beta.clear();
return result;
}
double trace_correction_e = 0.0;
double trace_correction_g = 0.0;
for (std::size_t i = 0; i < p; ++i) {
trace_correction_e += correction_e[i + i * p];
trace_correction_g += correction_g[i + i * p];
}
const double trace_p = sum_inverse_h - trace_correction_e;
const double trace_p_lambda =
sum_lambda_inverse_h - trace_correction_g;
double py_squared = 0.0;
double py_lambda_py = 0.0;
for (std::size_t i = 0; i < n; ++i) {
const double square = py[i] * py[i];
py_squared += square;
py_lambda_py += problem.lambda[i] * square;
}
result.gradient_e = -sigma_e * (trace_p - py_squared);
result.gradient_g = -sigma_g * (trace_p_lambda - py_lambda_py);
result.variance_score_e = 0.5 * (py_squared - trace_p);
result.variance_score_g = 0.5 * (py_lambda_py - trace_p_lambda);
if (compute_ai) {
std::vector<double> p_py;
if (!apply_p(problem, inverse_h, cholesky, py, p_py,
linear_algebra_error)) {
result.error = linear_algebra_error;
result.beta.clear();
return result;
}
std::vector<double> lambda_py(n);
for (std::size_t i = 0; i < n; ++i) {
lambda_py[i] = problem.lambda[i] * py[i];
}
std::vector<double> p_lambda_py;
if (!apply_p(problem, inverse_h, cholesky, lambda_py, p_lambda_py,
linear_algebra_error)) {
result.error = linear_algebra_error;
result.beta.clear();
return result;
}
const double y_p3_y = dot(py, p_py);
const double y_pp_lambda_p_y = dot(py, p_lambda_py);
const double y_p_lambda_p_lambda_p_y =
dot(lambda_py, p_lambda_py);
result.ai_ee = 2.0 * sigma_e2 * y_p3_y;
result.ai_eg =
2.0 * sigma_e * sigma_g * y_pp_lambda_p_y;
result.ai_gg = 2.0 * sigma_g2 * y_p_lambda_p_lambda_p_y;
}
if (compute_beta_covariance) {
if (!cholesky_inverse(cholesky, p, result.beta_covariance,
&linear_algebra_error)) {
result.error = linear_algebra_error;
result.beta.clear();
return result;
}
}
if (!std::isfinite(result.gradient_e) ||
!std::isfinite(result.gradient_g) ||
!std::isfinite(result.variance_score_e) ||
!std::isfinite(result.variance_score_g) ||
(compute_ai &&
(!std::isfinite(result.ai_ee) || !std::isfinite(result.ai_eg) ||
!std::isfinite(result.ai_gg)))) {
result.error = "gradient or average-information matrix is not finite";
result.beta.clear();
result.beta_covariance.clear();
return result;
}
result.valid = true;
return result;
}
struct InitialGuess {
bool valid = false;
double sigma_e = 0.0;
double sigma_g = 0.0;
double residual_mean_square = 0.0;
std::string error;
};
InitialGuess initial_guess(const ValidatedProblem& problem,
const RemlOptions& options) {
InitialGuess result;
const auto n = problem.n;
const auto p = problem.p;
std::vector<double> normal(p * p, 0.0);
std::vector<double> rhs(p, 0.0);
for (std::size_t i = 0; i < n; ++i) {
for (std::size_t row = 0; row < p; ++row) {
const double x_row = problem.x(i, row);
rhs[row] += x_row * problem.y[i];
for (std::size_t col = 0; col <= row; ++col) {
normal[row + col * p] += x_row * problem.x(i, col);
}
}
}
fill_symmetric_upper(normal, p);
std::vector<double> factor = normal;
if (!cholesky_factor_in_place(factor, p, &result.error)) {
result.error = "ordinary fixed-effect design is rank deficient: " +
result.error;
return result;
}
if (!factor_is_numerically_full_rank(factor, p,
options.rank_tolerance_relative)) {
result.error = "ordinary fixed-effect design is numerically rank deficient";
return result;
}
if (!cholesky_solve_in_place(factor, p, rhs.data(), 1, &result.error)) {
return result;
}
double residual_sum_squares = 0.0;
double phenotype_sum_squares = 0.0;
for (std::size_t i = 0; i < n; ++i) {
double fitted = 0.0;
for (std::size_t col = 0; col < p; ++col) {
fitted += problem.x(i, col) * rhs[col];
}
const double residual = problem.y[i] - fitted;
residual_sum_squares += residual * residual;
phenotype_sum_squares += problem.y[i] * problem.y[i];
}
const double residual_mean_square =
residual_sum_squares / static_cast<double>(n - p);
result.residual_mean_square = residual_mean_square;
const double phenotype_rms =
std::sqrt(phenotype_sum_squares / static_cast<double>(n));
const double sigma_floor =
std::max(1e-12, std::sqrt(std::numeric_limits<double>::epsilon()) *
std::max(1.0, phenotype_rms));
const double paper_initial =
std::max(sigma_floor, std::sqrt(0.5 * residual_mean_square));
result.sigma_e = std::isfinite(options.initial_sigma_e)
? options.initial_sigma_e
: paper_initial;
result.sigma_g = std::isfinite(options.initial_sigma_g)
? options.initial_sigma_g
: paper_initial;
if (!std::isfinite(result.sigma_e) || !std::isfinite(result.sigma_g) ||
!(result.sigma_e != 0.0 || result.sigma_g != 0.0)) {
result.error = "initial variance parameters are invalid";
return result;
}
result.valid = true;
return result;
}
bool ai_direction(const RemlEvaluation& evaluation,
const RemlOptions& options, std::array<double, 2>& direction,
std::string& error) {
const double scale = std::max(
{1.0, std::abs(evaluation.ai_ee), std::abs(evaluation.ai_eg),
std::abs(evaluation.ai_gg)});
double ridge = 0.0;
for (std::size_t attempt = 0; attempt <= options.ai_ridge_attempts;
++attempt) {
const double a = evaluation.ai_ee + ridge;
const double b = evaluation.ai_eg;
const double c = evaluation.ai_gg + ridge;
const double determinant = a * c - b * b;
if (a > 0.0 && c > 0.0 && determinant > kTiny * scale * scale) {
direction[0] =
(c * evaluation.gradient_e - b * evaluation.gradient_g) /
determinant;
direction[1] =
(a * evaluation.gradient_g - b * evaluation.gradient_e) /
determinant;
const double directional_derivative =
evaluation.gradient_e * direction[0] +
evaluation.gradient_g * direction[1];
if (std::isfinite(direction[0]) && std::isfinite(direction[1]) &&
directional_derivative > 0.0) {
return true;
}
}
ridge = attempt == 0 ? options.ai_ridge_relative * scale : ridge * 10.0;
}
error = "average-information matrix cannot produce an ascent direction";
return false;
}
struct LineSearchResult {
bool success = false;
double alpha = 0.0;
RemlEvaluation evaluation;
std::size_t evaluations = 0;
std::string error;
};
double directional_derivative(const RemlEvaluation& evaluation,
const std::array<double, 2>& direction) {
return evaluation.gradient_e * direction[0] +
evaluation.gradient_g * direction[1];
}
LineSearchResult strong_wolfe_line_search(
const ValidatedProblem& problem, double sigma_e, double sigma_g,
const RemlEvaluation& initial, const std::array<double, 2>& direction,
const RemlOptions& options) {
LineSearchResult result;
const double derivative_zero = directional_derivative(initial, direction);
if (!(derivative_zero > 0.0) || !std::isfinite(derivative_zero)) {
result.error = "line-search direction is not an ascent direction";
return result;
}
auto evaluate_alpha = [&](double alpha) {
++result.evaluations;
return evaluate_impl(problem, sigma_e + alpha * direction[0],
sigma_g + alpha * direction[1], false, false,
options.covariance_floor_relative,
options.rank_tolerance_relative);
};
auto armijo_bound = [&](double alpha) {
return initial.log_likelihood +
options.wolfe_c1 * alpha * derivative_zero;
};
auto zoom = [&](double alpha_lo, double alpha_hi,
RemlEvaluation evaluation_lo) -> LineSearchResult {
LineSearchResult zoom_result;
// Keep the global evaluation counter in result; copy it at return.
for (std::size_t iteration = 0;
iteration < options.line_search_max_zoom_iterations &&
result.evaluations < options.line_search_max_evaluations;
++iteration) {
const double alpha = 0.5 * (alpha_lo + alpha_hi);
const double interval_scale =
std::max({1.0, std::abs(alpha_lo), std::abs(alpha_hi)});
if (std::abs(alpha_hi - alpha_lo) <=
8.0 * std::numeric_limits<double>::epsilon() * interval_scale) {
zoom_result.error = "strong-Wolfe zoom interval collapsed";
break;
}
RemlEvaluation trial = evaluate_alpha(alpha);
if (!trial.valid || trial.log_likelihood < armijo_bound(alpha) ||
trial.log_likelihood <= evaluation_lo.log_likelihood) {
alpha_hi = alpha;
continue;
}
const double derivative =
directional_derivative(trial, direction);
if (std::abs(derivative) <=
options.wolfe_c2 * std::abs(derivative_zero)) {
zoom_result.success = true;
zoom_result.alpha = alpha;
zoom_result.evaluation = std::move(trial);
zoom_result.evaluations = result.evaluations;
return zoom_result;
}
if (derivative * (alpha_hi - alpha_lo) <= 0.0) {
alpha_hi = alpha_lo;
}
alpha_lo = alpha;
evaluation_lo = std::move(trial);
}
zoom_result.evaluations = result.evaluations;
if (zoom_result.error.empty()) {
zoom_result.error = "strong-Wolfe zoom exceeded its evaluation limit";
}
return zoom_result;
};
double alpha_previous = 0.0;
RemlEvaluation evaluation_previous = initial;
double alpha = std::min(options.initial_line_search_step,
options.maximum_line_search_step);
for (std::size_t iteration = 0;
result.evaluations < options.line_search_max_evaluations;
++iteration) {
RemlEvaluation trial = evaluate_alpha(alpha);
if (!trial.valid || trial.log_likelihood < armijo_bound(alpha) ||
(iteration > 0 &&
trial.log_likelihood <= evaluation_previous.log_likelihood)) {
return zoom(alpha_previous, alpha, std::move(evaluation_previous));
}
const double derivative = directional_derivative(trial, direction);
if (std::abs(derivative) <=
options.wolfe_c2 * std::abs(derivative_zero)) {
result.success = true;
result.alpha = alpha;
result.evaluation = std::move(trial);
return result;
}
if (derivative <= 0.0) {
return zoom(alpha, alpha_previous, std::move(trial));
}
if (alpha >= options.maximum_line_search_step) {
result.error =
"strong-Wolfe curvature condition not reached at maximum step";
return result;
}
alpha_previous = alpha;
evaluation_previous = std::move(trial);
alpha = std::min(2.0 * alpha, options.maximum_line_search_step);
}
result.error = "strong-Wolfe line search exceeded its evaluation limit";
return result;
}
void finalize_result(RemlResult& result, const ValidatedProblem& problem,
double sigma_e, double sigma_g,
const RemlOptions& options) {
RemlEvaluation final =
evaluate_impl(problem, sigma_e, sigma_g, false, true,
options.covariance_floor_relative,
options.rank_tolerance_relative);
if (!final.valid) {
result.status = FitStatus::numerical_error;
result.error = "final REML evaluation failed: " + final.error;
return;
}
result.sigma_e = sigma_e;
result.sigma_g = sigma_g;
result.sigma_e2 = sigma_e * sigma_e;
result.sigma_g2 = sigma_g * sigma_g;
const double total_variance = result.sigma_e2 + result.sigma_g2;
result.h2 = total_variance > 0.0 ? result.sigma_g2 / total_variance
: std::numeric_limits<double>::quiet_NaN();
result.log_likelihood = final.log_likelihood;
result.gradient_inf_norm =
std::max(std::abs(final.gradient_e), std::abs(final.gradient_g));
result.beta = std::move(final.beta);
result.beta_covariance_packed_lower.clear();
result.beta_covariance_packed_lower.reserve(
problem.p * (problem.p + 1) / 2);
for (std::size_t row = 0; row < problem.p; ++row) {
for (std::size_t col = 0; col <= row; ++col) {
result.beta_covariance_packed_lower.push_back(
final.beta_covariance[row + col * problem.p]);
}
}
}
void validate_options(const RemlOptions& options) {
if (options.max_iterations == 0) {
throw std::invalid_argument("max_iterations must be positive");
}
if (!(options.wolfe_c1 > 0.0 && options.wolfe_c1 < options.wolfe_c2 &&
options.wolfe_c2 < 1.0)) {
throw std::invalid_argument("strong-Wolfe constants must satisfy 0<c1<c2<1");
}
if (!(options.initial_line_search_step > 0.0) ||
!(options.maximum_line_search_step >=
options.initial_line_search_step)) {
throw std::invalid_argument("invalid line-search step bounds");
}
if (!(options.covariance_floor_relative >= 0.0) ||
!(options.rank_tolerance_relative >= 0.0) ||
!(options.boundary_score_tolerance >= 0.0) ||
!(options.boundary_likelihood_relative_tolerance >= 0.0) ||
!(options.boundary_h2_trigger >= 0.0 &&
options.boundary_h2_trigger <= 1.0)) {
throw std::invalid_argument("numerical tolerances must be non-negative");
}
}
} // namespace
const char* to_string(FitStatus status) noexcept {
switch (status) {
case FitStatus::converged:
return "converged";
case FitStatus::converged_boundary:
return "converged_boundary";
case FitStatus::max_iterations:
return "max_iterations";
case FitStatus::line_search_failed:
return "line_search_failed";
case FitStatus::rank_deficient:
return "rank_deficient";
case FitStatus::invalid_input:
return "invalid_input";
case FitStatus::non_positive_covariance:
return "non_positive_covariance";
case FitStatus::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,
bool compute_ai, bool compute_beta_covariance,
double covariance_floor_relative) {
try {
const auto problem = validate_problem(y_star, x_star, eigenvalues);
return evaluate_impl(problem, sigma_e, sigma_g, compute_ai,
compute_beta_covariance,
covariance_floor_relative, 1e-12);
} catch (const std::exception& exception) {
RemlEvaluation result;
result.error = exception.what();
return result;
}
}
RemlResult fit_ai_reml_spectral(const std::vector<double>& y_star,
const ColMajorMatrix& x_star,
const std::vector<double>& eigenvalues,
const RemlOptions& options) {
RemlResult result;
try {
validate_options(options);
const auto problem = validate_problem(y_star, x_star, eigenvalues);
const InitialGuess guess = initial_guess(problem, options);
if (!guess.valid) {
result.status = FitStatus::rank_deficient;
result.error = guess.error;
return result;
}
double sigma_e = guess.sigma_e;
double sigma_g = guess.sigma_g;
RemlEvaluation current =
evaluate_impl(problem, sigma_e, sigma_g, true, false,
options.covariance_floor_relative,
options.rank_tolerance_relative);
if (!current.valid) {
result.status = current.error.find("rank") != std::string::npos
? FitStatus::rank_deficient
: FitStatus::non_positive_covariance;
result.error = "initial REML evaluation failed: " + current.error;
return result;
}
// At v_g=sigma_g^2=0 the signed-sigma gradient is identically zero,
// even when the variance-component KKT condition is not satisfied.
// Construct the analytically optimized residual-only candidate once:
// v_e = RSS_OLS/(n-p). It is accepted below only if the one-sided
// score dL/dv_g is non-positive and its likelihood is not inferior to
// the current interior iterate. This is a boundary comparison, not an
// EM update and not an unconditional truncation of small estimates.
const double boundary_sigma_e =
guess.residual_mean_square > 0.0
? std::sqrt(guess.residual_mean_square)
: 0.0;
RemlEvaluation genetic_boundary;
if (boundary_sigma_e > 0.0 && std::isfinite(boundary_sigma_e)) {
genetic_boundary =
evaluate_impl(problem, boundary_sigma_e, 0.0, false, false,
options.covariance_floor_relative,
options.rank_tolerance_relative);
} else {
genetic_boundary.error =
"residual-only boundary has zero residual variance";
}
auto near_genetic_boundary = [&]() {
const double sigma_e2 = sigma_e * sigma_e;
const double sigma_g2 = sigma_g * sigma_g;
const double total = sigma_e2 + sigma_g2;
return total > 0.0 &&
sigma_g2 <= options.boundary_h2_trigger * total;
};
auto accept_genetic_boundary =
[&](const RemlEvaluation& reference) -> bool {
if (!genetic_boundary.valid || !reference.valid ||
genetic_boundary.variance_score_g >
options.boundary_score_tolerance) {
return false;
}
const double likelihood_scale =
std::max({1.0, std::abs(genetic_boundary.log_likelihood),
std::abs(reference.log_likelihood)});
const double likelihood_tolerance =
options.boundary_likelihood_relative_tolerance *
likelihood_scale;
if (genetic_boundary.log_likelihood + likelihood_tolerance <
reference.log_likelihood) {
return false;
}
sigma_e = boundary_sigma_e;
sigma_g = 0.0;
result.status = FitStatus::converged_boundary;
result.error.clear();
return true;
};
result.status = FitStatus::max_iterations;
for (std::size_t iteration = 0; iteration < options.max_iterations;
++iteration) {
// Once the interior path is sufficiently close to v_g=0, test the
// already-computed KKT boundary immediately. Acceptance still
// requires both the one-sided score and likelihood conditions, so
// this is an early exact boundary decision rather than truncation.
if (near_genetic_boundary() &&
accept_genetic_boundary(current)) {
break;
}
const double gradient_norm =
std::max(std::abs(current.gradient_e),
std::abs(current.gradient_g));
const double gradient_threshold =
options.gradient_absolute_tolerance +
options.gradient_relative_tolerance *
std::max(1.0, std::abs(current.log_likelihood));
if (gradient_norm <= gradient_threshold) {
// Do not mistake the automatic factor 2*sigma_g in the
// signed-parameter gradient for convergence when the
// one-sided variance score still points into the interior.
if (near_genetic_boundary() &&
current.variance_score_g >
options.boundary_score_tolerance) {
// Continue with the AI direction below.
} else {
result.status = FitStatus::converged;
break;
}
}
std::array<double, 2> direction{};
std::string direction_error;
if (!ai_direction(current, options, direction, direction_error)) {
if (!accept_genetic_boundary(current)) {
result.status = FitStatus::numerical_error;
result.error = direction_error;
}
break;
}
LineSearchResult line_search = strong_wolfe_line_search(
problem, sigma_e, sigma_g, current, direction, options);
result.line_search_evaluations += line_search.evaluations;
if (!line_search.success) {
if (!accept_genetic_boundary(current)) {
result.status = FitStatus::line_search_failed;
result.error = line_search.error;
}
break;
}
const double previous_likelihood = current.log_likelihood;
const double previous_sigma_e = sigma_e;
const double previous_sigma_g = sigma_g;
sigma_e += line_search.alpha * direction[0];
sigma_g += line_search.alpha * direction[1];
current = std::move(line_search.evaluation);
++result.iterations;
// The line search only needs score, not AI. Re-evaluate once at
// the accepted point to construct the next AI direction.
current = evaluate_impl(problem, sigma_e, sigma_g, true, false,
options.covariance_floor_relative,
options.rank_tolerance_relative);
if (!current.valid) {
result.status = FitStatus::numerical_error;
result.error = "accepted REML point could not be re-evaluated: " +
current.error;
break;
}
const double step_scale =
std::max({1.0, std::abs(previous_sigma_e),
std::abs(previous_sigma_g)});
const double step_norm =
std::max(std::abs(sigma_e - previous_sigma_e),
std::abs(sigma_g - previous_sigma_g));
const double likelihood_change =
std::abs(current.log_likelihood - previous_likelihood);
if (step_norm <= options.step_relative_tolerance * step_scale &&
likelihood_change <=
options.likelihood_relative_tolerance *
std::max(1.0, std::abs(previous_likelihood))) {
if (near_genetic_boundary() &&
accept_genetic_boundary(current)) {
break;
}
if (!(near_genetic_boundary() &&
current.variance_score_g >
options.boundary_score_tolerance)) {
result.status = FitStatus::converged;
break;
}
}
}
if (result.status == FitStatus::max_iterations) {
(void)accept_genetic_boundary(current);
}
finalize_result(result, problem, sigma_e, sigma_g, options);
return result;
} catch (const std::invalid_argument& exception) {
result.status = FitStatus::invalid_input;
result.error = exception.what();
return result;
} catch (const std::exception& exception) {
result.status = FitStatus::numerical_error;
result.error = exception.what();
return result;
}
}
} // namespace spectra::reml

130
tests/cpp/test_grm_io.cpp Normal file
View File

@@ -0,0 +1,130 @@
#include "spectra_reml/grm.hpp"
#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstdlib>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <stdexcept>
#include <string>
#include <vector>
namespace {
void require(bool condition, const std::string& message) {
if (!condition) {
throw std::runtime_error(message);
}
}
void require_near(double actual, double expected, double tolerance,
const std::string& message) {
if (std::abs(actual - expected) > tolerance) {
throw std::runtime_error(message + ": actual=" +
std::to_string(actual) +
", expected=" + std::to_string(expected));
}
}
class TemporaryDirectory {
public:
TemporaryDirectory() {
const auto seed =
std::chrono::high_resolution_clock::now().time_since_epoch().count();
path_ = std::filesystem::temp_directory_path() /
("spectra_reml_grm_test_" + std::to_string(seed));
std::filesystem::create_directories(path_);
}
~TemporaryDirectory() {
std::error_code ignored;
std::filesystem::remove_all(path_, ignored);
}
const std::filesystem::path& path() const noexcept { return path_; }
private:
std::filesystem::path path_;
};
void write_floats(const std::filesystem::path& path,
const std::vector<float>& values) {
std::ofstream output(path, std::ios::binary | std::ios::trunc);
output.write(reinterpret_cast<const char*>(values.data()),
static_cast<std::streamsize>(values.size() * sizeof(float)));
if (!output) {
throw std::runtime_error("failed to create GRM fixture");
}
}
void test_read_and_decompose() {
TemporaryDirectory temporary;
const auto grm_path = temporary.path() / "fixture.grm.bin";
const auto id_path = temporary.path() / "fixture.grm.id";
// Symmetric positive-definite matrix, packed as GCTA lower triangle.
const std::vector<float> packed = {1.2F, 0.2F, 1.1F,
-0.1F, 0.15F, 0.9F};
write_floats(grm_path, packed);
{
std::ofstream ids(id_path);
ids << "F1 I1\nF2 I2\nF3 I3\n";
}
spectra::reml::validate_grm_id_count(id_path, 3);
const auto matrix = spectra::reml::read_gcta_grm_lower_triangle(grm_path, 3);
require_near(matrix(0, 0), 1.2, 1e-7, "wrong GRM diagonal");
require_near(matrix(1, 0), 0.2, 1e-7, "wrong GRM lower triangle");
require_near(matrix(0, 1), 0.2, 1e-7, "GRM was not symmetrized");
require_near(matrix(2, 0), -0.1, 1e-7, "wrong GRM packed order");
require_near(matrix(1, 2), 0.15, 1e-7, "wrong GRM upper triangle");
const auto spectral =
spectra::reml::read_and_diagonalize_gcta_grm(grm_path, 3);
require(spectral.eigenvalues.size() == 3, "wrong eigenvalue count");
require(std::is_sorted(spectral.eigenvalues.begin(),
spectral.eigenvalues.end()),
"LAPACK eigenvalues are not ascending");
require(spectral.minimum_eigenvalue > 0.0,
"positive-definite fixture has non-positive eigenvalue");
// Reconstruct U diag(lambda) U' and compare to the original matrix.
for (std::size_t row = 0; row < 3; ++row) {
for (std::size_t col = 0; col < 3; ++col) {
double reconstructed = 0.0;
for (std::size_t axis = 0; axis < 3; ++axis) {
reconstructed += spectral.eigenvectors(row, axis) *
spectral.eigenvalues[axis] *
spectral.eigenvectors(col, axis);
}
require_near(reconstructed, matrix(row, col), 2e-12,
"eigendecomposition does not reconstruct GRM");
}
}
}
void test_wrong_file_size_is_rejected() {
TemporaryDirectory temporary;
const auto grm_path = temporary.path() / "short.grm.bin";
write_floats(grm_path, {1.0F, 0.0F});
bool rejected = false;
try {
(void)spectra::reml::read_gcta_grm_lower_triangle(grm_path, 2);
} catch (const std::runtime_error&) {
rejected = true;
}
require(rejected, "short GRM file was not rejected");
}
} // namespace
int main() {
try {
test_read_and_decompose();
test_wrong_file_size_is_rejected();
std::cout << "test_grm_io: PASS\n";
return EXIT_SUCCESS;
} catch (const std::exception& exception) {
std::cerr << "test_grm_io: FAIL: " << exception.what() << '\n';
return EXIT_FAILURE;
}
}

View File

@@ -0,0 +1,226 @@
#include "spectra_reml/reml.hpp"
#include <algorithm>
#include <cmath>
#include <cstdlib>
#include <iostream>
#include <stdexcept>
#include <string>
#include <vector>
namespace {
void require(bool condition, const std::string& message) {
if (!condition) {
throw std::runtime_error(message);
}
}
void require_near(double actual, double expected, double relative_tolerance,
const std::string& message) {
const double scale = std::max({1.0, std::abs(actual), std::abs(expected)});
if (!std::isfinite(actual) || !std::isfinite(expected) ||
std::abs(actual - expected) > relative_tolerance * scale) {
throw std::runtime_error(message + ": actual=" +
std::to_string(actual) +
", expected=" + std::to_string(expected));
}
}
struct Fixture {
std::vector<double> y;
spectra::reml::ColMajorMatrix x;
std::vector<double> lambda;
Fixture() : x(18, 3) {
lambda = {0.12, 0.18, 0.25, 0.31, 0.39, 0.48, 0.58, 0.69, 0.81,
0.95, 1.10, 1.28, 1.47, 1.69, 1.94, 2.22, 2.55, 2.91};
y.resize(lambda.size());
for (std::size_t i = 0; i < lambda.size(); ++i) {
const double z = (static_cast<double>(i) - 8.5) / 5.0;
x(i, 0) = 1.0;
x(i, 1) = z;
x(i, 2) = (i % 3 == 0 ? -0.8 : (i % 3 == 1 ? 0.1 : 0.7));
const double deterministic_noise =
std::sin(1.7 * static_cast<double>(i) + 0.3) *
std::sqrt(0.35 + 0.55 * lambda[i]) +
0.22 * std::cos(0.37 * static_cast<double>(i));
y[i] = 0.4 - 0.25 * x(i, 1) + 0.15 * x(i, 2) +
deterministic_noise;
}
}
};
void test_score_matches_finite_difference() {
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, true, true);
require(evaluated.valid, "reference REML evaluation failed: " +
evaluated.error);
require(evaluated.beta.size() == fixture.x.cols(),
"beta has wrong length");
require(evaluated.beta_covariance.size() ==
fixture.x.cols() * fixture.x.cols(),
"beta covariance has wrong size");
const double epsilon_e = 2e-6 * std::max(1.0, std::abs(sigma_e));
const double epsilon_g = 2e-6 * std::max(1.0, std::abs(sigma_g));
const auto e_plus = spectra::reml::evaluate_reml_spectral(
fixture.y, fixture.x, fixture.lambda, sigma_e + epsilon_e, sigma_g,
false, false);
const auto e_minus = spectra::reml::evaluate_reml_spectral(
fixture.y, fixture.x, fixture.lambda, sigma_e - epsilon_e, sigma_g,
false, false);
const auto g_plus = spectra::reml::evaluate_reml_spectral(
fixture.y, fixture.x, fixture.lambda, sigma_e, sigma_g + epsilon_g,
false, false);
const auto g_minus = spectra::reml::evaluate_reml_spectral(
fixture.y, fixture.x, fixture.lambda, sigma_e, sigma_g - epsilon_g,
false, false);
require(e_plus.valid && e_minus.valid && g_plus.valid && g_minus.valid,
"finite-difference REML evaluation failed");
const double finite_difference_e =
(e_plus.log_likelihood - e_minus.log_likelihood) / (2.0 * epsilon_e);
const double finite_difference_g =
(g_plus.log_likelihood - g_minus.log_likelihood) / (2.0 * epsilon_g);
require_near(evaluated.gradient_e, finite_difference_e, 2e-6,
"sigma_e score disagrees with finite difference");
require_near(evaluated.gradient_g, finite_difference_g, 2e-6,
"sigma_g score disagrees with finite difference");
require(evaluated.ai_ee > 0.0 && evaluated.ai_gg > 0.0,
"AI diagonal must be positive");
require(evaluated.ai_ee * evaluated.ai_gg -
evaluated.ai_eg * evaluated.ai_eg >=
-1e-10,
"AI matrix must be positive semidefinite");
}
void test_signed_parameterization_and_row_permutation() {
Fixture fixture;
const auto positive = spectra::reml::evaluate_reml_spectral(
fixture.y, fixture.x, fixture.lambda, 0.61, 0.52, true, false);
const auto negative = spectra::reml::evaluate_reml_spectral(
fixture.y, fixture.x, fixture.lambda, -0.61, -0.52, true, false);
require(positive.valid && negative.valid, "signed evaluations failed");
require_near(positive.log_likelihood, negative.log_likelihood, 1e-13,
"likelihood must depend on squared standard deviations");
require_near(positive.gradient_e, -negative.gradient_e, 1e-12,
"sigma_e gradient must change sign");
require_near(positive.gradient_g, -negative.gradient_g, 1e-12,
"sigma_g gradient must change sign");
std::vector<double> permuted_y(fixture.y.rbegin(), fixture.y.rend());
std::vector<double> permuted_lambda(fixture.lambda.rbegin(),
fixture.lambda.rend());
spectra::reml::ColMajorMatrix permuted_x(fixture.x.rows(), fixture.x.cols());
for (std::size_t i = 0; i < fixture.x.rows(); ++i) {
for (std::size_t j = 0; j < fixture.x.cols(); ++j) {
permuted_x(i, j) = fixture.x(fixture.x.rows() - 1 - i, j);
}
}
const auto permuted = spectra::reml::evaluate_reml_spectral(
permuted_y, permuted_x, permuted_lambda, 0.61, 0.52, true, false);
require(permuted.valid, "permuted evaluation failed");
require_near(positive.log_likelihood, permuted.log_likelihood, 2e-13,
"orthogonal-axis permutation changed REML likelihood");
require_near(positive.gradient_e, permuted.gradient_e, 2e-12,
"orthogonal-axis permutation changed sigma_e score");
require_near(positive.gradient_g, permuted.gradient_g, 2e-12,
"orthogonal-axis permutation changed sigma_g score");
}
void test_ai_reml_fit_improves_likelihood() {
Fixture fixture;
spectra::reml::RemlOptions options;
options.initial_sigma_e = 0.9;
options.initial_sigma_g = 0.9;
options.max_iterations = 150;
options.gradient_absolute_tolerance = 1e-8;
options.gradient_relative_tolerance = 1e-9;
const auto initial = spectra::reml::evaluate_reml_spectral(
fixture.y, fixture.x, fixture.lambda, options.initial_sigma_e,
options.initial_sigma_g, true, false);
require(initial.valid, "initial fit evaluation failed");
const auto fitted = spectra::reml::fit_ai_reml_spectral(
fixture.y, fixture.x, fixture.lambda, options);
require(fitted.has_estimates(),
"AI-REML returned no estimates: " + fitted.error);
require(fitted.status == spectra::reml::FitStatus::converged,
std::string("AI-REML did not converge: ") +
spectra::reml::to_string(fitted.status) + " " + fitted.error);
require(fitted.log_likelihood >= initial.log_likelihood - 1e-10,
"AI-REML decreased the likelihood");
require(fitted.gradient_inf_norm < 2e-5,
"AI-REML final score is too large");
require(fitted.beta.size() == fixture.x.cols(),
"fitted beta has wrong length");
require(fitted.beta_covariance_packed_lower.size() ==
fixture.x.cols() * (fixture.x.cols() + 1) / 2,
"packed beta covariance has wrong length");
require_near(fitted.sigma_e2, fitted.sigma_e * fitted.sigma_e, 1e-14,
"reported sigma_e2 is inconsistent");
require_near(fitted.sigma_g2, fitted.sigma_g * fitted.sigma_g, 1e-14,
"reported sigma_g2 is inconsistent");
}
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);
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;
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);
require(boundary.valid, "residual-only boundary evaluation failed: " +
boundary.error);
require(boundary.variance_score_g < -1e-3,
"fixture does not have a negative one-sided genetic score");
spectra::reml::RemlOptions options;
options.max_iterations = 100;
options.gradient_absolute_tolerance = 1e-10;
options.gradient_relative_tolerance = 1e-11;
const auto fitted =
spectra::reml::fit_ai_reml_spectral(y, x, lambda, options);
require(fitted.has_estimates(),
"boundary AI-REML returned no estimates: " + fitted.error);
require(fitted.status == spectra::reml::FitStatus::converged_boundary,
std::string("residual-only optimum was not accepted as a KKT boundary: ") +
spectra::reml::to_string(fitted.status) + " " + fitted.error);
require(fitted.iterations <= 10,
"near-boundary KKT check did not terminate promptly");
require_near(fitted.sigma_g2, 0.0, 1e-15,
"genetic boundary must report exactly zero variance");
require_near(fitted.h2, 0.0, 1e-15,
"genetic boundary must report exactly zero h2");
require_near(fitted.sigma_e2, expected_sigma_e2, 2e-12,
"residual-only variance is not RSS/(n-p)");
require_near(fitted.log_likelihood, boundary.log_likelihood, 2e-12,
"reported boundary likelihood is inconsistent");
}
} // namespace
int main() {
try {
test_score_matches_finite_difference();
test_signed_parameterization_and_row_permutation();
test_ai_reml_fit_improves_likelihood();
test_residual_only_kkt_boundary();
std::cout << "test_reml_synthetic: PASS\n";
return EXIT_SUCCESS;
} catch (const std::exception& exception) {
std::cerr << "test_reml_synthetic: FAIL: " << exception.what() << '\n';
return EXIT_FAILURE;
}
}

View File

@@ -0,0 +1,163 @@
from __future__ import annotations
import csv
import gzip
import importlib.util
import json
import tempfile
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest import mock
import numpy as np
SCRIPT = Path(__file__).resolve().parents[2] / "python" / "spectra_reml.py"
SPEC = importlib.util.spec_from_file_location("spectra_reml", SCRIPT)
assert SPEC is not None and SPEC.loader is not None
cli = importlib.util.module_from_spec(SPEC)
SPEC.loader.exec_module(cli)
class SpectraRemlCliTests(unittest.TestCase):
def fixture(self, root: Path) -> tuple[Path, dict]:
n, p0 = 3, 2
np.zeros(n * (n + 1) // 2, dtype="<f4").tofile(root / "grm.bin")
(root / "grm.id").write_text("F1 I1\nF2 I2\nF3 I3\n", encoding="utf-8")
np.asarray([[1, -1], [1, 0], [1, 1]], dtype="<f8").tofile(root / "base.bin")
np.zeros((2, n), dtype="<f8").tofile(root / "phenotypes.bin")
np.asarray([[0, 1, 2]], dtype="<f4").tofile(root / "extra.bin")
(root / "tasks.tsv").write_text(
"task_index\ttask_id\tphenotype_row\tn_extra_covariates\n"
"0\ttrait_a\t0\t0\n"
"1\ttrait_b\t1\t1\n",
encoding="utf-8",
)
np.asarray([0, 0, 1], dtype="<i8").tofile(root / "offsets.bin")
np.asarray([0], dtype="<i4").tofile(root / "indices.bin")
manifest_path = root / "manifest.json"
args = cli.build_parser().parse_args(
[
"make-manifest",
"--manifest", str(manifest_path),
"--grm-bin", str(root / "grm.bin"),
"--grm-id", str(root / "grm.id"),
"--base-x", str(root / "base.bin"),
"--phenotypes", str(root / "phenotypes.bin"),
"--extra-covariates", str(root / "extra.bin"),
"--tasks", str(root / "tasks.tsv"),
"--extra-offsets", str(root / "offsets.bin"),
"--extra-indices", str(root / "indices.bin"),
"--output-dir", str(root / "blocks"),
"--n-samples", "3",
"--n-base-covariates", "2",
"--n-phenotype-rows", "2",
"--n-extra-covariate-rows", "1",
]
)
cli.make_manifest(args)
return manifest_path, cli.validate_manifest(manifest_path)
@staticmethod
def write_blocks(root: Path) -> None:
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",
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")
(root / "block_000000.complete").write_text(
"format\t{}\nblock\t0\ntasks\t2\nbeta_elements\t5\ncov_elements\t9\n".format(cli.BLOCK_FORMAT),
encoding="utf-8",
)
def test_manifest_and_generic_command(self) -> None:
with tempfile.TemporaryDirectory() as temporary:
root = Path(temporary)
_, manifest = self.fixture(root)
args = cli.build_parser().parse_args(
["run", "--manifest", str(root / "manifest.json"), "--engine", "spectra_reml"]
)
command = cli.build_engine_command(Path("spectra_reml"), manifest, args)
for option in (
"--phenotypes",
"--extra-covariates",
"--extra-offsets",
"--extra-indices",
"--n-phenotype-rows",
"--n-extra-covariate-rows",
):
self.assertEqual(command.count(option), 1)
def test_signature_recovery_and_finalize(self) -> None:
with tempfile.TemporaryDirectory() as temporary:
root = Path(temporary)
manifest_path, manifest = self.fixture(root)
engine = root / "spectra_reml"
engine.write_bytes(b"synthetic engine")
args = cli.build_parser().parse_args(
["run", "--manifest", str(manifest_path), "--engine", str(engine), "--resume"]
)
with mock.patch.object(cli.subprocess, "run", return_value=SimpleNamespace(returncode=0)):
cli.run_engine(manifest, args)
output = Path(manifest["paths"]["output_directory"])
signature = cli.read_signature(output / "run.signature.json")
self.assertEqual(signature["format"], cli.RUN_SIGNATURE_FORMAT)
self.write_blocks(output)
changed = cli.build_parser().parse_args(
[
"run", "--manifest", str(manifest_path), "--engine", str(engine),
"--resume", "--max-iterations", "101",
]
)
with self.assertRaisesRegex(RuntimeError, "signature differs"):
cli.run_engine(manifest, changed)
result = root / "results.tsv.gz"
cli.finalize(manifest, result)
with gzip.open(result, "rt", encoding="utf-8", newline="") as handle:
rows = list(csv.DictReader(handle, delimiter="\t"))
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)
altered = dict(manifest)
altered["created_utc"] = "changed"
with self.assertRaisesRegex(RuntimeError, "does not belong"):
cli.finalize(altered, root / "unsafe.tsv")
def test_force_invalidates_before_cleanup_and_dry_run_is_read_only(self) -> None:
with tempfile.TemporaryDirectory() as temporary:
root = Path(temporary)
manifest_path, manifest = self.fixture(root)
engine = root / "spectra_reml"
engine.write_bytes(b"engine")
output = Path(manifest["paths"]["output_directory"])
self.write_blocks(output)
resume_args = cli.build_parser().parse_args(
["run", "--manifest", str(manifest_path), "--engine", str(engine), "--resume"]
)
cli.atomic_json(output / "run.signature.json", cli.make_signature(manifest, engine, resume_args))
before = sorted(path.name for path in output.iterdir())
dry = cli.build_parser().parse_args(
["run", "--manifest", str(manifest_path), "--engine", "missing", "--force", "--dry-run"]
)
cli.run_engine(manifest, dry)
self.assertEqual(before, sorted(path.name for path in output.iterdir()))
force = cli.build_parser().parse_args(
["run", "--manifest", str(manifest_path), "--engine", str(engine), "--force"]
)
with mock.patch.object(cli, "block_files", side_effect=RuntimeError("cleanup interrupted")):
with self.assertRaisesRegex(RuntimeError, "cleanup interrupted"):
cli.run_engine(manifest, force)
self.assertFalse((output / "run.signature.json").exists())
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,128 @@
#!/usr/bin/env python3
"""Independent NumPy checks for the spectral-space AI-REML equations.
This file deliberately does not import the production C++ implementation. It
serves as a small, readable oracle for gradient and transformation tests.
"""
from __future__ import annotations
import math
import numpy as np
def restricted_loglik_and_derivatives(
theta: np.ndarray,
eigenvalues: np.ndarray,
x_star: np.ndarray,
y_star: np.ndarray,
) -> tuple[float, np.ndarray, np.ndarray]:
"""Return paper-scale REML log likelihood, score, and AI matrix.
Parameters are the standard-deviation coordinates ``(sigma_e, sigma_g)``.
The function forms P explicitly because it is only used on tiny test data.
"""
sigma_e, sigma_g = (float(theta[0]), float(theta[1]))
diagonal = sigma_e**2 + sigma_g**2 * eigenvalues
if np.any(diagonal <= 0.0):
raise ValueError("H is not positive definite")
w = 1.0 / diagonal
h_inv = np.diag(w)
c = x_star.T @ h_inv @ x_star
c_inv = np.linalg.inv(c)
p = h_inv - h_inv @ x_star @ c_inv @ x_star.T @ h_inv
py = p @ y_star
loglik = -0.5 * (
np.linalg.slogdet(c)[1]
+ np.log(diagonal).sum()
+ float(y_star @ py)
)
score_e = -sigma_e * (np.trace(p) - float(py @ py))
score_g = -sigma_g * (
np.trace(p @ np.diag(eigenvalues))
- float(py @ (eigenvalues * py))
)
p_py = p @ py
lambda_py = eigenvalues * py
p_lambda_py = p @ lambda_py
ai_ee = 2.0 * sigma_e**2 * float(py @ p_py)
ai_eg = 2.0 * sigma_e * sigma_g * float(py @ p_lambda_py)
ai_gg = 2.0 * sigma_g**2 * float(lambda_py @ p_lambda_py)
ai = np.asarray([[ai_ee, ai_eg], [ai_eg, ai_gg]], dtype=np.float64)
return loglik, np.asarray([score_e, score_g]), ai
def finite_difference_gradient(
theta: np.ndarray,
eigenvalues: np.ndarray,
x_star: np.ndarray,
y_star: np.ndarray,
step: float = 1e-6,
) -> np.ndarray:
result = np.empty(2, dtype=np.float64)
for index in range(2):
delta = np.zeros(2, dtype=np.float64)
delta[index] = step
high = restricted_loglik_and_derivatives(
theta + delta, eigenvalues, x_star, y_star
)[0]
low = restricted_loglik_and_derivatives(
theta - delta, eigenvalues, x_star, y_star
)[0]
result[index] = (high - low) / (2.0 * step)
return result
def test_score_matches_finite_difference() -> None:
rng = np.random.default_rng(70123)
n = 17
p = 4
eigenvalues = np.linspace(0.05, 2.1, n)
x_star = np.column_stack((np.ones(n), rng.normal(size=(n, p - 1))))
y_star = rng.normal(size=n)
theta = np.asarray([0.83, 0.57])
_, score, ai = restricted_loglik_and_derivatives(
theta, eigenvalues, x_star, y_star
)
numerical = finite_difference_gradient(theta, eigenvalues, x_star, y_star)
np.testing.assert_allclose(score, numerical, rtol=2e-6, atol=2e-6)
np.testing.assert_allclose(ai, ai.T, rtol=0.0, atol=1e-12)
assert np.linalg.eigvalsh(ai).min() > 0.0
def test_orthogonal_transformation_preserves_reml() -> None:
rng = np.random.default_rng(70124)
n = 15
q, _ = np.linalg.qr(rng.normal(size=(n, n)))
eigenvalues = np.linspace(0.1, 1.8, n)
grm = q @ np.diag(eigenvalues) @ q.T
x = np.column_stack((np.ones(n), rng.normal(size=(n, 3))))
y = rng.normal(size=n)
theta = np.asarray([0.76, 0.62])
sigma_e, sigma_g = theta
h = sigma_g**2 * grm + sigma_e**2 * np.eye(n)
h_inv = np.linalg.inv(h)
c = x.T @ h_inv @ x
p = h_inv - h_inv @ x @ np.linalg.inv(c) @ x.T @ h_inv
original = -0.5 * (
np.linalg.slogdet(c)[1]
+ np.linalg.slogdet(h)[1]
+ float(y @ p @ y)
)
transformed = restricted_loglik_and_derivatives(
theta, eigenvalues, q.T @ x, q.T @ y
)[0]
assert math.isclose(original, transformed, rel_tol=2e-12, abs_tol=2e-12)
if __name__ == "__main__":
test_score_matches_finite_difference()
test_orthogonal_transformation_preserves_reml()
print("PASS: spectral AI-REML reference checks")