masreml Rust modules
The Rust kernel powering masreml is organised into four top-level groups (matrix/, reml/, solver/, gwas/) plus a thin lib.rs binding layer. Each section below is the module-level documentation extracted directly from the corresponding .rs file by _scripts/extract-rust-docs.R, so this page stays in sync with the source.
For per-function signatures and private-item details, build the full cargo doc HTML locally:
cd masreml/src/rust && cargo doc --no-deps --document-private-items --openPackage overview
lib.rs
Computational backend for the masreml R package: universal REML-BLUP genomic prediction supporting biallelic SNP, multi-allelic microhaplotype, and pedigree-based relationship matrices.
Module map
- [
matrix] — relationship-matrix construction:snp_additive/snp_dominance— VanRaden (2008) additive and Da et al. (2014) / Wang & Da (2014) dominance G matrices.mh_additive— Da (2015) \(W_{\alpha h}\) additive \(G\) for multi-allelic microhaplotypes, with the per-locus frequency-weighted row shrinkage.pedigree— Henderson (1976) recursive numerator relationship matrix \(A\).
- [
reml] — variance-component estimation:he_regression— Haseman–Elston regression (closed-form, starting values).ai_reml— Average-Information REML (Johnson & Thompson, 1995).em_reml— Expectation-Maximization REML (Dempster et al., 1977; Meyer, 1989). Slow but always non-negative.adaptive— auto-selects between HE / AI / EM based on conditioning and convergence diagnostics.
- [
solver] — fixed/random effects solvers given fitted variance components:cholesky— direct factorization; default for \(n < 10{,}000\).pcg— preconditioned conjugate gradient; default for larger \(n\).factorized— reusable Cholesky cache used by both.
- [
gwas] — single-marker association:emmax— EMMAX (Kang et al., 2010) with a pre-factorized \(V\).smoother— local moving-average smoother for likelihood-ratio statistics across markers/blocks (used by GWABLUP).
- [
utils] — smallndarray_linalghelpers used across modules.
R-facing API
Public entry points are wrapped with #[extendr] in this file and exposed to R as r_build_g_snp_add, r_build_g_snp_dom, r_build_g_mh_add, r_build_a_ped, r_run_reml, r_run_emmax, etc. The extendr_module! macro at the bottom of this file registers them all in one block.
References
- VanRaden, P. M. (2008). Efficient methods to compute genomic predictions. J. Dairy Sci., 91:4414–4423.
- Da, Y. (2015). Multi-allelic haplotype model based on genetic partition. BMC Genetics, 16:144.
- Henderson, C. R. (1976). A simple method for computing the inverse of a numerator relationship matrix. Biometrics, 32:69–83.
- Johnson, D. L. & Thompson, R. (1995). Restricted maximum likelihood estimation of variance components. J. Dairy Sci., 78:449–456.
- Kang, H. M. et al. (2010). Variance component model to account for sample structure in GWAS. Nat. Genet., 42:348–354.
Relationship matrices (matrix/)
matrix/mod.rs
Relationship-matrix construction.
Each submodule produces a positive semi-definite \(G\) (or \(A\) for pedigree) matrix that can be combined in a multi-component mixed model by the [crate::reml] kernels.
Shared infrastructure:
- [
GMatrix] — owned \(n \times n\) matrix wrapper used as the common return type so REML can iterate over heterogeneous components. - [
compute_k] — computes \(\mathrm{tr}(WW^\top)/n\) directly from \(W\) without forming \(WW^\top\), the scaling constant for normalising marker-based \(G\) matrices. - [
validate_w] — defensive checks (non-empty, finite values) before building \(G\); returns a structuredMatrixErrorfor upstream R-side error reporting. - [
MatrixError] — error enum bridged toextendr_api::Errorso R surfaces a clean error message instead of a Rust panic.
matrix/snp_additive.rs
🦀 masreml/matrix/snp_additive.rs
SNP additive relationship matrix (VanRaden 2008).
Builds the additive genomic relationship matrix
G = W W' / k, k = 2 · sum_j p_j (1 - p_j)
where \(W\) is the column-centered SNP matrix \(W_{ij} = X_{ij} - 2 p_j\) with \(X_{ij} \in \{0, 1, 2\}\) allele dosages and \(p_j\) the reference (training-set) allele frequency at marker \(j\).
GWABLUP support
Optionally accepts a per-marker weight vector \(d = (d_1, \dots, d_m)\) to construct a GWAS-weighted \(G\) matrix:
G_wa = sum_j d_j · w_j w_j' / sum_j d_j · 2 p_j (1 - p_j)
When weights = None this collapses to the standard VanRaden \(G\). Weights typically come from a smoothed likelihood-ratio statistic (see [crate::gwas::smoother]).
Parallelism
The \(WW^\top\) product is computed with ndarray::parallel row-chunked across rayon threads. The scaling constant \(k\) is computed via [crate::matrix::compute_k] without materialising \(WW^\top\) first.
matrix/snp_dominance.rs
🦀 masreml/matrix/snp_dominance.rs
SNP dominance relationship matrix.
Implements the dominance coding of Da et al. (2014) and Wang & Da (2014):
w_δ_ij = -2 · p_j² if genotype AA (homozygous reference)
w_δ_ij = 2 · p_j (1 - p_j) if genotype Aa (heterozygous)
w_δ_ij = -2 · (1 - p_j)² if genotype aa (homozygous alternate)
with \(p_j\) the reference allele frequency. The dominance G matrix is then
D = W_δ W_δ' / k_δ, k_δ = tr(W_δ W_δ') / n
Under HWE, the expected value of \(W_\delta^{(j)}\) is zero, so the coding is mean-zero by construction — analogous to the frequency-centered additive coding in [super::snp_additive].
Why a separate dominance G?
Additive and dominance variance components are orthogonal under HWE. In a mixed model
y = X β + Z g + Z d + ε
Var(g) = G · σ²_a Var(d) = D · σ²_d
REML can estimate \(\sigma^2_a\) and \(\sigma^2_d\) separately. Dropping the dominance term biases the additive heritability estimate upward when dominance variance is non-trivial.
matrix/mh_additive.rs
🦀 masreml/matrix/mh_additive.rs
Multi-allelic additive relationship matrix \(G_{\alpha h}\) (Da, 2015).
Builds the \(n \times n\) additive G matrix from phased microhaplotype data, locus by locus:
- Compute training-set allele frequencies \(p_k\) per locus.
- Choose the most frequent microhaplotype as the baseline and drop it from the design matrix (identifiability constraint).
- Encode the remaining \(h - 1\) microhaplotypes with the Da (2015) three-value rule into a per-locus \(W_l\) matrix.
- Apply the per-locus [
frequency_weighted_row_shrinkage] — see its docstring for the precise definition; identical to the equivalent step inmasbayes. - Optionally scale each locus by
sqrt(weight_l)for GWABLUP. - Accumulate \(G_{\alpha h} = \sum_l W_l W_l^\top\) and the scaling constant \(k_{\alpha h} = \mathrm{tr}(G_{\alpha h}) / n\), returning the normalised matrix \(G_{\alpha h} / k_{\alpha h}\).
Train / test alignment
Allele frequencies and baseline microhaplotypes are taken from the reference (training) data when ref_hap1 / ref_hap2 are supplied; otherwise they are computed from the data being encoded. Always pass training references when encoding a test set — see the discussion in [crate::matrix::snp_additive] for the analogous SNP case.
Parallelism
Locus-level \(W_l\) encoding is parallelised over individuals via rayon. The outer loop over loci is sequential because each locus accumulates into the shared \(G\) matrix.
Reference
Da, Y. (2015). Multi-allelic haplotype model based on genetic partition for genomic prediction and variance component estimation using SNP markers. BMC Genetics, 16:144.
matrix/pedigree.rs
Pedigree-based numerator relationship matrix \(A\) (Henderson 1976).
Builds the \(n \times n\) expected relationship matrix from a recoded pedigree using Henderson’s recursive formula:
a_ii = 1 + F_i, F_i = 0.5 · a[sire_i, dam_i]
a_ij = 0.5 · (a[i, sire_j] + a[i, dam_j]) for i < j
Founders (sire = dam = 0) contribute \(a_{ii} = 1\) and zero off-diagonals to ancestors not in the recoded pedigree.
Input format
Animals must be topologically pre-sorted so that any parent appears before its offspring. Sire and dam indices use 0 for unknown parent; all other indices are 1-based references into the same pedigree array. Validation of this ordering is the caller’s responsibility — the recursion silently produces wrong results if the order is violated.
Memory
The full \(A\) matrix is dense \(n \times n\) (the recursion does not exploit sparsity). For large pedigrees, consider \(A^{-1}\) sparse forms (Quaas, 1976) — not implemented here but a natural extension.
Reference
Henderson, C. R. (1976). A simple method for computing the inverse of a numerator relationship matrix used in prediction of breeding values. Biometrics, 32:69–83.
Variance-component estimation (reml/)
reml/mod.rs
Variance-component estimation via REML.
Implements three REML algorithms and an adaptive dispatcher:
- [
he_regression] — Haseman–Elston regression on phenotype products. Closed-form, single-shot, used for starting values and sanity checks. - [
ai_reml] — Average-Information REML (Johnson & Thompson, 1995). Newton-style updates with the AI matrix; fast but can diverge or hit negative variances on ill-conditioned problems. - [
em_reml] — Expectation-Maximization REML (Dempster et al., 1977; Meyer, 1989). Always produces non-negative variance components but converges slowly. - [
adaptive] — auto-selector that runs HE first for starting values, tries AI, and falls back to EM if AI fails to converge or produces negative components.
Common interfaces
- [
RemlData] — bundles the response, fixed-effect design \(X\), and the list of random-effect \(G\) matrices. Constructed once by the adaptive entry point. - [
VarianceComponents] — return type containing \(\sigma^2\) per component plus convergence diagnostics (iterations, log-likelihood). - [
RemlError] — structured error type bridged toextendr_api::Errorfor clean R-side reporting.
Numerical infrastructure
Inversion of the working matrix \(V = \sum_i G_i \sigma^2_i + I \sigma^2_e\) is delegated to [crate::utils::linalg] helpers, which use ndarray_linalg LAPACK bindings. Threading is controlled by set_num_threads from the same module.
reml/he_regression.rs
🦀 masreml/reml/he_regression.rs
Haseman–Elston regression for variance-component estimation.
Model
Given adjusted phenotypes \(\hat{y}_i\) (residuals from fixed effects), HE regresses pairwise products on entries of the relationship matrices:
ŷ_i · ŷ_j = μ + Σ_k K_k[i,j] · h²_k + e_ij
where \(K_k\) is the \(k\)-th relationship matrix (additive G, dominance D, pedigree A, etc.) and \(h^2_k\) is the corresponding heritability component. The off-diagonal regression is closed form (one OLS solve) and yields unbiased starting values for AI/EM-REML.
When to use
- Starting values for AI / EM iterations (default in [
super::adaptive]). - Quick check on a new dataset to see whether the variance structure is identifiable before committing to slow iterative REML.
- Sanity baseline when AI/EM produce suspicious estimates.
Limitations
HE assumes the residuals are well-modelled by additive variance components, and the closed-form OLS step does not account for correlation between pairwise products. The resulting estimates are consistent but inefficient (higher variance than ML/REML).
Numerical details
Pairwise products are accumulated in parallel via rayon. The compute_reml_loglik helper exported from this module is reused by AI-REML and EM-REML for convergence monitoring.
reml/ai_reml.rs
Average-Information REML (AI-REML).
Implements the Newton-style update of Johnson & Thompson (1995):
σ²⁽ᵗ⁺¹⁾ = σ²⁽ᵗ⁾ + AI⁻¹ · score
where the score and AI matrix are computed from \(P = V^{-1} - V^{-1}X(X'V^{-1}X)^{-1}X'V^{-1}\) applied to the response. AI averages the observed and expected information matrices, which tends to be better-conditioned than either alone and converges at a near-Newton rate close to the optimum.
Algorithm
- Initialise \(\sigma^2\) from HE (or user-supplied values).
- Compute \(V\), factor it once via Cholesky.
- Compute \(P y\) and the AI matrix.
- Solve \(\Delta = AI^{-1} \cdot \mathrm{score}\) and update \(\sigma^2\).
- Iterate until \(\|\Delta\| < \mathrm{tol}\) or
max_iter.
When AI fails
- Variance components heading negative — AI does not enforce non-negativity; the [
super::adaptive] dispatcher catches this and falls back to EM. - Singular AI matrix — happens when components are non-identifiable.
- Slow convergence with badly scaled \(G\) matrices — sometimes solved by scaling \(G\) to have \(\mathrm{tr}(G)/n = 1\) before passing in.
Reference
Johnson, D. L. & Thompson, R. (1995). Restricted maximum likelihood estimation of variance components for univariate animal models using sparse matrix techniques and average information. J. Dairy Sci., 78:449–456.
reml/em_reml.rs
Expectation-Maximization REML (EM-REML).
Multiplicative variance-component update of Dempster et al. (1977) and Meyer (1989):
σ²_i⁽ᵗ⁺¹⁾ = σ²_i⁽ᵗ⁾² / n_i · ( y'P G_i P y + tr(G_i · C_ii) )
σ²_e⁽ᵗ⁺¹⁾ = σ²_e⁽ᵗ⁾² / n · ( y'PPy + tr(C_ee) )
Each update is a ratio of non-negative quantities, so \(\sigma^2 \ge 0\) is preserved automatically — no constraint handling needed. The price is linear (not super-linear) convergence; EM typically needs 10–100× more iterations than AI to reach the same tolerance.
When to use
- Always non-negative — preferred when AI returns negative components or stalls.
- As the fallback path in [
super::adaptive] when AI fails. - For verification — running EM to convergence and comparing with AI confirms both are at the same optimum.
Performance
Each iteration computes \(Py\) and \(P G_i P y\), which dominates the runtime; the helpers in [crate::utils::linalg] amortise the Cholesky factorisation of \(V\) across all components.
References
- Dempster, A. P., Laird, N. M. & Rubin, D. B. (1977). Maximum likelihood from incomplete data via the EM algorithm. J. R. Stat. Soc. B, 39:1–38.
- Meyer, K. (1989). Restricted maximum likelihood to estimate variance components for animal models with several random effects using a derivative-free algorithm. Genet. Sel. Evol., 21:317–340.
reml/adaptive.rs
Adaptive REML algorithm dispatcher.
Single entry point for the R-side run_reml(...) call. Internally cascades through the three REML algorithms in [super]:
- HE regression ([
super::he_regression]) for starting values. Always run; cheap, closed-form, robust. - AI-REML ([
super::ai_reml]) for fast Newton-style convergence near the optimum. Aborted if it produces a negative variance component or fails to converge withinmax_iter. - EM-REML ([
super::em_reml]) as the always-non-negative fallback when AI fails.
Dispatcher logic
- If the user passes
method = "ai","em", or"he"explicitly, that algorithm runs alone (no fallback). Useful for diagnostic comparisons. - With
method = "auto"(default), HE → AI → EM cascade applies. - Thread count for parallel sections (HE pairwise products, parallel \(G\) matrix-vector products) is set via [
crate::utils::linalg::set_num_threads] usingn_threadsfrom the caller; default uses all available cores.
Return value
[VarianceComponents] carries the final \(\sigma^2\) vector, the method actually used ("ai", "em", or "he"), iteration count, and the REML log-likelihood — enough for downstream R code to surface a clean convergence diagnostic to the user.
EBV / BLUP solvers (solver/)
solver/mod.rs
EBV/BLUP solvers given fitted variance components.
Given \(V = \sum_i G_i \sigma^2_i + I \sigma^2_e\) from [crate::reml], these submodules solve the mixed-model equations to produce estimated breeding values (EBVs) and fixed-effect estimates.
Available solvers
- [
cholesky] — direct Cholesky factorisation of \(V\). Wraps [factorized::FactorizedV] for reuse across multiple right-hand sides. Default for \(n < 10{,}000\). - [
pcg] — preconditioned conjugate gradient. Avoids forming \(V^{-1}\) explicitly; the diagonal of \(V\) serves as the preconditioner. Default for \(n \ge 10{,}000\). - [
factorized] — shared Cholesky cache used bycholeskyand re-exported to [crate::gwas::emmax] so EMMAX does not re-factor \(V\) per marker.
Dispatcher
[solve_ebv] is the public entry point. It honours an explicit solver = "cholesky" | "pcg" argument or auto-selects via [auto_select_solver] based on \(n\). The returned [BlupResult] contains per-random-effect EBVs, fixed-effect estimates, and diagnostic info (which solver ran, iteration count for PCG).
solver/cholesky.rs
Cholesky-based EBV solver.
Thin wrapper around [super::factorized::FactorizedV] that exposes a solve_cholesky_internal(...) entry point matching the signature expected by [super::solve_ebv]. The heavy lifting (factorisation + triangular solves) is in the factorized module so it can be reused by other consumers such as [crate::gwas::emmax].
Direct factorisation is exact (no convergence tolerance) and produces BLUPs in a single pass, but memory scales as \(O(n^2)\) for the factor itself; for \(n \gtrsim 10{,}000\) prefer the iterative [super::pcg] solver.
solver/pcg.rs
Preconditioned Conjugate Gradient (PCG) solver.
Iterative solver for the mixed-model equations \(V x = b\) that avoids forming or factoring \(V\) explicitly. Recommended for \(n \ge 10{,}000\) where the \(O(n^2)\) memory cost of [super::cholesky] becomes prohibitive.
Algorithm
Standard PCG with the diagonal of \(V\) as preconditioner (Jacobi). Iteration stops when the relative residual \(\|r_k\| / \|b\|\) falls below tol (default 1e-8) or when max_iter is reached.
When to use vs. Cholesky
- \(n < 10{,}000\) → Cholesky is faster (one factorisation, then triangular solves).
- \(n \ge 10{,}000\) → PCG scales as \(O(n^2)\) per iteration via the \(G\) matrix-vector products, never materialising \(V^{-1}\).
- Very ill-conditioned \(V\) → PCG may need many iterations; consider tighter convergence thresholds or scaling \(G\) before passing in.
Output
Same [BlupResult] shape as [super::cholesky], with solver = "pcg" and n_iter set to the actual iteration count for diagnostics.
solver/factorized.rs
🦀 masreml/solver/factorized.rs
Cached Cholesky factorisation of \(V\).
Many downstream operations need to solve \(V x = b\) with the same \(V\) but many different \(b\):
- [
super::cholesky] — one solve per right-hand side for EBV. - [
crate::gwas::emmax] — one solve per marker for the per-SNP Wald statistic.
[FactorizedV] packages the Cholesky factor of \(V = \sum_i G_i \sigma^2_i + I \sigma^2_e\) together with the dependent quantities (factor of \(X' V^{-1} X\), residual projection) that all consumers need. Computing them once and passing the handle around eliminates redundant factorisation work.
Lifetime / mutability
FactorizedV is constructed from owned matrices (it stores its own Cholesky factors). It is &self-immutable on consumption, so multiple threads can call solve(...) concurrently — needed by the parallel per-marker loop in EMMAX.
GWAS (gwas/)
gwas/mod.rs
Single-marker association testing.
Two algorithms with complementary roles in the GWABLUP pipeline:
- [
emmax] — EMMAX (Kang et al., 2010). Per-marker mixed-model association test that reuses a single pre-factorised \(V\) across all markers, giving GWAS-quality \(p\)-values at a fraction of the per-SNP REML cost. - [
smoother] — moving-average smoother for likelihood-ratio statistics across adjacent markers/blocks, used to construct the GWAS-weighted \(G\) matrix (G_wa) in [crate::matrix::snp_additive] and [crate::matrix::mh_additive].
Output convention
[GwasResult] uses per-marker (SNP) or per-block (MH) aggregated statistics: LR test statistic, effect estimate \(\hat\beta\), standard error, and \(p\)-value. For multi-allelic MH blocks, \(\hat\beta\) and SE are aggregated across alleles within a block, and the \(p\)-value uses a chi-squared with \(\mathrm{df} = h - 1\) (number of non-baseline microhaplotypes).
gwas/emmax.rs
EMMAX — efficient single-marker mixed-model association.
Implements the EMMAX algorithm of Kang et al. (2010): instead of re-estimating variance components for every marker, factor the null-model \(V = G \sigma^2_g + I \sigma^2_e\) once and reuse the factorisation to compute Wald/likelihood-ratio tests for every marker in turn.
Algorithm
- Fit the null mixed model \(y = X\beta + u + \varepsilon\) via REML (handled upstream by [
crate::reml]). - Cholesky-factor \(V\) once — encapsulated in [
crate::solver::factorized::FactorizedV] for reuse. - For each marker \(j\):
- Solve \(V^{-1} x_j\) with the cached factorisation.
- Compute \(\hat\beta_j\), its standard error, and the LR statistic.
- Derive a \(\chi^2\) \(p\)-value (df=1 for SNP, df=\(h-1\) for MH).
- Loop over markers in parallel via rayon.
Multi-allelic extension
For MH blocks, the per-allele test statistics are aggregated to a single per-block LR with df=\(h-1\). Effect-size aggregation follows the weighted average used elsewhere in the package, keeping the output shape consistent with the SNP path.
Reference
Kang, H. M., Sul, J. H., Service, S. K., et al. (2010). Variance component model to account for sample structure in genome-wide association studies. Nat. Genet., 42:348–354.
gwas/smoother.rs
Local moving-average smoother for GWAS likelihood-ratio statistics.
Produces the per-marker weights \(d_j\) used to construct the GWAS-weighted relationship matrix \(G_{wa}\) in [crate::matrix::snp_additive] and [crate::matrix::mh_additive].
Smoother
d_j = mean( LR[j - w/2 : j + w/2] )
with a centred window of size window (caller-specified). At chromosome / array edges the window is shrunk rather than padded — no edge artefacts from zero-padding.
Why smooth?
Per-SNP LR statistics are noisy; smoothing averages signal over a local LD neighbourhood so that markers in linkage with a causal variant share its weight contribution. The result is a less variable \(G_{wa}\) that still up-weights regions of true association without over-fitting to single-SNP noise.
Choice of window
Typical values are 5–20 SNPs (or microhaplotype blocks) depending on marker density and local LD; the R-side wrapper exposes this as a window argument with default 5.
Utilities (utils/)
utils/mod.rs
Shared utilities.
Currently a single submodule, [linalg], collecting small linear algebra primitives that are used across [crate::matrix], [crate::reml], [crate::solver], and [crate::gwas].
Kept as a top-level module so future additions (RNG helpers, benchmarking shims, etc.) have an obvious home without disturbing the existing public surface.
utils/linalg.rs
Shared linear algebra primitives.
Thin wrappers around ndarray_linalg LAPACK bindings plus a few REML-specific helpers that appear in more than one place:
solve_matrix(V, B)— solve \(V X = B\) for matrix \(B\) via LU/Cholesky under the hood; returns aMatrixErroron failure so REML can classify the cause (singular \(V\), dimension mismatch, etc.).compute_py(V, X, y)— compute \(Py\) where \(P = V^{-1} - V^{-1} X (X' V^{-1} X)^{-1} X' V^{-1}\), the projection used inside HE, AI, and EM-REML scoring functions.log_det_cholesky(L)— log-determinant via Cholesky factor, used to evaluate the REML log-likelihood.set_num_threads(n)— global rayon thread-pool configuration. Called once at the top of [crate::reml::adaptive] so all parallel sections honour the same thread count.
Functions are deliberately small and unit-testable in isolation; they exist here rather than inlined into their callers so all REML algorithms share the same numerically tested implementations.