From da971a2ca71ac2ec32a2c681b3462fb63696bf01 Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Sun, 9 Aug 2026 14:13:38 +0800 Subject: [PATCH] Add reusable SpectraREML batch AI-REML engine --- .gitignore | 10 + CMakeLists.txt | 118 ++++ README.md | 153 +++++ cmake/FindREML_OpenBLAS.cmake | 135 +++++ docs/FORMAT.md | 162 +++++ examples/example.py | 83 +++ include/spectra_reml/batch_io.hpp | 62 ++ include/spectra_reml/grm.hpp | 36 ++ include/spectra_reml/linalg.hpp | 41 ++ include/spectra_reml/reml.hpp | 49 ++ include/spectra_reml/types.hpp | 133 ++++ python/spectra_reml.py | 766 +++++++++++++++++++++++ scripts/run_server.sh | 112 ++++ src/batch_io.cpp | 711 ++++++++++++++++++++++ src/grm.cpp | 121 ++++ src/linalg.cpp | 201 ++++++ src/main.cpp | 354 +++++++++++ src/reml_core.cpp | 842 ++++++++++++++++++++++++++ tests/cpp/test_grm_io.cpp | 130 ++++ tests/cpp/test_reml_synthetic.cpp | 226 +++++++ tests/python/test_spectra_reml_cli.py | 163 +++++ tests/test_math_reference.py | 128 ++++ 22 files changed, 4736 insertions(+) create mode 100644 .gitignore create mode 100644 CMakeLists.txt create mode 100644 cmake/FindREML_OpenBLAS.cmake create mode 100644 docs/FORMAT.md create mode 100644 examples/example.py create mode 100644 include/spectra_reml/batch_io.hpp create mode 100644 include/spectra_reml/grm.hpp create mode 100644 include/spectra_reml/linalg.hpp create mode 100644 include/spectra_reml/reml.hpp create mode 100644 include/spectra_reml/types.hpp create mode 100644 python/spectra_reml.py create mode 100644 scripts/run_server.sh create mode 100644 src/batch_io.cpp create mode 100644 src/grm.cpp create mode 100644 src/linalg.cpp create mode 100644 src/main.cpp create mode 100644 src/reml_core.cpp create mode 100644 tests/cpp/test_grm_io.cpp create mode 100644 tests/cpp/test_reml_synthetic.cpp create mode 100644 tests/python/test_spectra_reml_cli.py create mode 100644 tests/test_math_reference.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..5e68301 --- /dev/null +++ b/.gitignore @@ -0,0 +1,10 @@ +/build*/ +/*.obj +/*.lib +/*.exp +/out/ +/results/ +__pycache__/ +*.py[cod] +*.tmp.* +/examples/synthetic_work/ diff --git a/CMakeLists.txt b/CMakeLists.txt new file mode 100644 index 0000000..f795b43 --- /dev/null +++ b/CMakeLists.txt @@ -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 + $ + $ +) +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}") diff --git a/README.md b/README.md index fa25ac4..a33303f 100644 --- a/README.md +++ b/README.md @@ -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). diff --git a/cmake/FindREML_OpenBLAS.cmake b/cmake/FindREML_OpenBLAS.cmake new file mode 100644 index 0000000..e514bb0 --- /dev/null +++ b/cmake/FindREML_OpenBLAS.cmake @@ -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 + #include + 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 +) diff --git a/docs/FORMAT.md b/docs/FORMAT.md new file mode 100644 index 0000000..79c6fa8 --- /dev/null +++ b/docs/FORMAT.md @@ -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. diff --git a/examples/example.py b/examples/example.py new file mode 100644 index 0000000..a41f23d --- /dev/null +++ b/examples/example.py @@ -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=" +#include +#include +#include +#include + +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 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 diff --git a/include/spectra_reml/grm.hpp b/include/spectra_reml/grm.hpp new file mode 100644 index 0000000..233b483 --- /dev/null +++ b/include/spectra_reml/grm.hpp @@ -0,0 +1,36 @@ +#pragma once + +#include "spectra_reml/types.hpp" + +#include +#include +#include + +namespace spectra::reml { + +struct SpectralGrm { + std::vector 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 diff --git a/include/spectra_reml/linalg.hpp b/include/spectra_reml/linalg.hpp new file mode 100644 index 0000000..1feb0b0 --- /dev/null +++ b/include/spectra_reml/linalg.hpp @@ -0,0 +1,41 @@ +#pragma once + +#include "spectra_reml/types.hpp" + +#include +#include +#include + +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 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& matrix, + std::size_t order, + std::string* error = nullptr); +[[nodiscard]] bool cholesky_solve_in_place( + const std::vector& factor, std::size_t order, double* rhs, + std::size_t rhs_columns, std::string* error = nullptr); +[[nodiscard]] bool cholesky_inverse( + const std::vector& factor, std::size_t order, + std::vector& inverse, std::string* error = nullptr); + +[[nodiscard]] double dot(const std::vector& a, + const std::vector& b); +[[nodiscard]] double infinity_norm(const std::vector& x); + +} // namespace spectra::reml diff --git a/include/spectra_reml/reml.hpp b/include/spectra_reml/reml.hpp new file mode 100644 index 0000000..435de6b --- /dev/null +++ b/include/spectra_reml/reml.hpp @@ -0,0 +1,49 @@ +#pragma once + +#include "spectra_reml/types.hpp" + +#include +#include +#include + +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& y_star, const ColMajorMatrix& x_star, + const std::vector& 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::infinity(); + double gradient_e = std::numeric_limits::quiet_NaN(); + double gradient_g = std::numeric_limits::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::quiet_NaN(); + double variance_score_g = std::numeric_limits::quiet_NaN(); + double ai_ee = std::numeric_limits::quiet_NaN(); + double ai_eg = std::numeric_limits::quiet_NaN(); + double ai_gg = std::numeric_limits::quiet_NaN(); + std::vector beta; + std::vector beta_covariance; + std::string error; +}; + +[[nodiscard]] RemlEvaluation evaluate_reml_spectral( + const std::vector& y_star, const ColMajorMatrix& x_star, + const std::vector& 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 diff --git a/include/spectra_reml/types.hpp b/include/spectra_reml/types.hpp new file mode 100644 index 0000000..a6a3505 --- /dev/null +++ b/include/spectra_reml/types.hpp @@ -0,0 +1,133 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +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 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& values() noexcept { return values_; } + const std::vector& 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 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::quiet_NaN(); + double initial_sigma_g = std::numeric_limits::quiet_NaN(); +}; + +struct RemlResult { + FitStatus status = FitStatus::invalid_input; + double sigma_e = std::numeric_limits::quiet_NaN(); + double sigma_g = std::numeric_limits::quiet_NaN(); + double sigma_e2 = std::numeric_limits::quiet_NaN(); + double sigma_g2 = std::numeric_limits::quiet_NaN(); + double h2 = std::numeric_limits::quiet_NaN(); + // Equation B.5 of the thesis, omitting additive constants. + double log_likelihood = -std::numeric_limits::infinity(); + std::size_t iterations = 0; + std::size_t line_search_evaluations = 0; + double gradient_inf_norm = std::numeric_limits::infinity(); + std::vector beta; + // Row-wise packed lower triangle: + // (0,0), (1,0), (1,1), (2,0), (2,1), (2,2), ... + std::vector 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 extra_covariate_rows; +}; + +struct TaskResult { + RemlTask task; + RemlResult fit; +}; + +} // namespace spectra::reml diff --git a/python/spectra_reml.py b/python/spectra_reml.py new file mode 100644 index 0000000..69420aa --- /dev/null +++ b/python/spectra_reml.py @@ -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="= 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="= 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) diff --git a/scripts/run_server.sh b/scripts/run_server.sh new file mode 100644 index 0000000..0873075 --- /dev/null +++ b/scripts/run_server.sh @@ -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 diff --git a/src/batch_io.cpp b/src/batch_io.cpp new file mode 100644 index 0000000..f4fae28 --- /dev/null +++ b/src/batch_io.cpp @@ -0,0 +1,711 @@ +#include "spectra_reml/batch_io.hpp" + +#include "spectra_reml/linalg.hpp" +#include "spectra_reml/reml.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#ifdef _OPENMP +#include +#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::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(&marker); + if (bytes[0] != 0x04U || !std::numeric_limits::is_iec559 || + !std::numeric_limits::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 +std::vector 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(bytes), description); + std::vector 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(values.data()), + static_cast(bytes)); + } + if (!input || input.gcount() != static_cast(bytes)) { + throw std::runtime_error("short read from " + std::string(description) + + " file " + path.string()); + } + return values; +} + +std::vector split_tab(const std::string& line) { + std::vector 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& 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::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(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& 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 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(row) >= total_rows) { + throw std::runtime_error("extra_covariates row index is outside the matrix"); + } + const std::uintmax_t offset = + static_cast(row) * row_bytes; + input.seekg(static_cast(offset), std::ios::beg); + input.read(reinterpret_cast(buffer.data()), + static_cast(row_bytes)); + if (!input || input.gcount() != static_cast(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(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& 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(offset), std::ios::beg); + input.read(reinterpret_cast(selected.data() + column * sample_count), + static_cast(row_bytes)); + if (!input || input.gcount() != static_cast(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::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& 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(values.data()), + static_cast(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& 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 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 beta_values; + std::vector covariance_values; + std::vector beta_offsets(results.size(), -1); + std::vector 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(beta_values.size()); + covariance_offsets[index] = + static_cast(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 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(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 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 task_rows; + std::unordered_set 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( + 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( + 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( + extra_indices_i32, static_cast(offsets.back()), "extra-covariate indices"); + + std::vector tasks; + tasks.reserve(task_rows.size()); + for (std::size_t i = 0; i < task_rows.size(); ++i) { + const auto begin = static_cast(offsets[i]); + const auto end = static_cast(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 within_task; + for (std::size_t offset = begin; offset < end; ++offset) { + const auto row = extra_indices[offset]; + if (row < 0 || + static_cast(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 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 unique_extra_rows(unique_extra_set.begin(), + unique_extra_set.end()); + std::unordered_map 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::numeric_limits::max())) { + throw std::invalid_argument("thread count exceeds OpenMP integer range"); + } + omp_set_dynamic(0); + omp_set_num_threads(static_cast(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 results(end - begin); + +#ifdef _OPENMP +#pragma omp parallel for schedule(dynamic) +#endif + for (std::int64_t local_signed = 0; + local_signed < static_cast(end - begin); + ++local_signed) { + const auto local = static_cast(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 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 diff --git a/src/grm.cpp b/src/grm.cpp new file mode 100644 index 0000000..54be65f --- /dev/null +++ b/src/grm.cpp @@ -0,0 +1,121 @@ +#include "spectra_reml/grm.hpp" + +#include "spectra_reml/linalg.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +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::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::max() / sizeof(float)) { + throw std::overflow_error("GRM byte count overflow"); + } + return static_cast(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 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(lower_triangle.data()), + static_cast(expected_bytes)); + if (!input || input.gcount() != static_cast(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(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 diff --git a/src/linalg.cpp b/src/linalg.cpp new file mode 100644 index 0000000..1843164 --- /dev/null +++ b/src/linalg.cpp @@ -0,0 +1,201 @@ +#include "spectra_reml/linalg.hpp" + +#if defined(REML_BLAS_MKL) +#include +#include +#else +#include +#include +#endif + +#include +#include +#include +#include +#include + +namespace spectra::reml { +namespace { + +lapack_int checked_lapack_int(std::size_t value, const char* name) { + if (value > static_cast(std::numeric_limits::max())) { + throw std::overflow_error(std::string(name) + " exceeds LAPACK integer range"); + } + return static_cast(value); +} + +int checked_blas_int(std::size_t value, const char* name) { + if (value > static_cast(std::numeric_limits::max())) { + throw std::overflow_error(std::string(name) + " exceeds CBLAS integer range"); + } + return static_cast(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 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 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& 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& 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& factor, std::size_t order, + std::vector& 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& a, const std::vector& 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& x) { + double result = 0.0; + for (const double value : x) { + result = std::max(result, std::abs(value)); + } + return result; +} + +} // namespace spectra::reml diff --git a/src/main.cpp b/src/main.cpp new file mode 100644 index 0000000..0c8b4e3 --- /dev/null +++ b/src/main.cpp @@ -0,0 +1,354 @@ +#include "spectra_reml/batch_io.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +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 values; + std::unordered_set flags; +}; + +Arguments parse_arguments(int argc, char** argv) { + const std::unordered_set 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( + std::numeric_limits::max()) || + (!allow_zero && parsed == 0)) { + throw std::invalid_argument("invalid integer for " + name + ": " + text); + } + return static_cast(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 +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 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; + } +} diff --git a/src/reml_core.cpp b/src/reml_core.cpp new file mode 100644 index 0000000..8b23d55 --- /dev/null +++ b/src/reml_core.cpp @@ -0,0 +1,842 @@ +#include "spectra_reml/reml.hpp" + +#include "spectra_reml/linalg.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace spectra::reml { +namespace { + +constexpr double kTiny = 64.0 * std::numeric_limits::epsilon(); + +struct ValidatedProblem { + const std::vector& y; + const ColMajorMatrix& x; + const std::vector& lambda; + std::size_t n; + std::size_t p; +}; + +ValidatedProblem validate_problem(const std::vector& y, + const ColMajorMatrix& x, + const std::vector& 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& 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& factor, + std::size_t p, + double relative_tolerance) { + double min_diagonal = std::numeric_limits::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& inverse_h, + const std::vector& cholesky, + const std::vector& input, + std::vector& output, std::string& error) { + const auto n = problem.n; + const auto p = problem.p; + std::vector 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::min(), + covariance_floor_relative * std::max(covariance_scale, + std::numeric_limits::min())); + + std::vector 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 normal_matrix(p * p, 0.0); + std::vector 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 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 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 correction_e(p * p, 0.0); + std::vector 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 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 lambda_py(n); + for (std::size_t i = 0; i < n; ++i) { + lambda_py[i] = problem.lambda[i] * py[i]; + } + std::vector 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 normal(p * p, 0.0); + std::vector 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 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(n - p); + result.residual_mean_square = residual_mean_square; + const double phenotype_rms = + std::sqrt(phenotype_sum_squares / static_cast(n)); + const double sigma_floor = + std::max(1e-12, std::sqrt(std::numeric_limits::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& 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& 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& 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::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::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 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& y_star, const ColMajorMatrix& x_star, + const std::vector& 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& y_star, + const ColMajorMatrix& x_star, + const std::vector& 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 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 diff --git a/tests/cpp/test_grm_io.cpp b/tests/cpp/test_grm_io.cpp new file mode 100644 index 0000000..bbe8452 --- /dev/null +++ b/tests/cpp/test_grm_io.cpp @@ -0,0 +1,130 @@ +#include "spectra_reml/grm.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +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& values) { + std::ofstream output(path, std::ios::binary | std::ios::trunc); + output.write(reinterpret_cast(values.data()), + static_cast(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 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; + } +} diff --git a/tests/cpp/test_reml_synthetic.cpp b/tests/cpp/test_reml_synthetic.cpp new file mode 100644 index 0000000..a8a265a --- /dev/null +++ b/tests/cpp/test_reml_synthetic.cpp @@ -0,0 +1,226 @@ +#include "spectra_reml/reml.hpp" + +#include +#include +#include +#include +#include +#include +#include + +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 y; + spectra::reml::ColMajorMatrix x; + std::vector 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(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(i) + 0.3) * + std::sqrt(0.35 + 0.55 * lambda[i]) + + 0.22 * std::cos(0.37 * static_cast(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 permuted_y(fixture.y.rbegin(), fixture.y.rend()); + std::vector 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 lambda(n); + std::vector y(n, 0.0); + for (std::size_t i = 0; i < n; ++i) { + x(i, 0) = 1.0; + lambda[i] = 0.1 * static_cast(i + 1); + } + // The residual energy is confined to the two smallest GRM eigenvalues, + // making the one-sided score for v_g at zero strictly negative. + y[0] = 1.0; + y[1] = -1.0; + const double expected_sigma_e2 = 2.0 / static_cast(n - 1); + const auto boundary = spectra::reml::evaluate_reml_spectral( + y, x, lambda, std::sqrt(expected_sigma_e2), 0.0, false, true); + 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; + } +} diff --git a/tests/python/test_spectra_reml_cli.py b/tests/python/test_spectra_reml_cli.py new file mode 100644 index 0000000..3c36488 --- /dev/null +++ b/tests/python/test_spectra_reml_cli.py @@ -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=" 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=" 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() diff --git a/tests/test_math_reference.py b/tests/test_math_reference.py new file mode 100644 index 0000000..c12f3e7 --- /dev/null +++ b/tests/test_math_reference.py @@ -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")