Add reusable SpectraREML batch AI-REML engine

include/spectra_reml/batch_io.hpp

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

include/spectra_reml/grm.hpp

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

include/spectra_reml/linalg.hpp

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

include/spectra_reml/reml.hpp

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

include/spectra_reml/types.hpp

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