masbayes Rust modules

The Rust kernel powering masbayes is split into eight focused modules. 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 masbayes/src/rust && cargo doc --no-deps --document-private-items --open

Package overview

lib.rs

🦀 masbayes/lib.rs

Computational backend for the masbayes R package: Bayesian genomic prediction models for biallelic SNP and multi-allelic microhaplotype markers. All numerically heavy work (MCMC, stochastic EM, design matrix construction) lives here; the R side is a thin extendr wrapper.

Module map

  • [matrix] — design-matrix construction:
    • WMatrixBuilder implements the Da (2015) \(W_{\alpha h}\) encoding for phased multi-allelic haplotype data, followed by a per-locus frequency-weighted row shrinkage.
    • Helper routines for biallelic SNP encoding (VanRaden) are exposed via the FFI wrappers in this file.
  • [bayesa] — BayesA Gibbs sampler. Per-marker variance follows a scaled inverse chi-squared prior, leading to a t-shrunk effect distribution. Supports binary traits via Albert–Chib data augmentation.
  • [bayesr] — BayesR Gibbs sampler. Marker effects follow a four-component normal mixture (one spike at zero plus three slab components scaled relative to the genetic variance); mixture proportions are updated with a Dirichlet–Multinomial step.
  • [bayesa_em] / [bayesr_em] — Stochastic EM variants of the two samplers. Faster but discard posterior uncertainty.
  • [utils] — RNG helpers (inverse-gamma, Dirichlet, normal, tabulate) and R↔︎ndarray conversion shims used across the kernels.
  • [types] — POD result structs returned to R via the extendr wrappers.

Reproducibility

All samplers take a seed: u64 initialising a PCG64 RNG (rand_pcg). Given the same data, hyperparameters, and seed, the kernels reproduce posterior samples bit-for-bit on the same platform.

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.

Module reference

matrix.rs

🦀 masbayes/matrix.rs

Design-matrix construction for multi-allelic microhaplotype markers.

Pipeline

  1. Per-block encoding — for each haplotype locus (block) and each non-baseline microhaplotype \(k\), the entry \(W_{i,k}\) for individual \(i\) is filled with the Da (2015) three-value coding rule (Eqs. 22–24): \(-2(1-p_k)\) when the individual is homozygous for \(k\), \(-(1-2p_k)\) when heterozygous, \(+2p_k\) when \(k\) is absent.
  2. Baseline drop — the most frequent microhaplotype per locus is used as the contrast reference and excluded from the columns; its effect is absorbed by the intercept \(\mu\) in downstream models.
  3. Frequency-weighted row shrinkage — see [frequency_weighted_row_shrinkage]. Applied per locus after encoding; partially reduces the alignment between each individual row and the locus’s allele-frequency vector.

Equivalence with VanRaden

Da’s three-value rule can be written in closed form as \(W_{i,k} = 2 p_k - n_{i,k}\), where \(n_{i,k} \in \{0,1,2\}\) is the dosage of microhaplotype \(k\) in individual \(i\). This is the per-allele VanRaden centering with the sign flipped, stacked across all non-baseline microhaplotypes within the locus.

Train / test alignment

WMatrixBuilder accepts a ReferenceStructure for test-set encoding so that allele frequencies and the choice of baseline microhaplotype come from the training set, not the test set itself. Re-estimating \(p_k\) on test data would shift the centering and bias GEBVs — never do that.

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.

bayesa.rs

🦀 masbayes/bayesa.rs

BayesA Gibbs sampler.

Model

For a continuous trait,

y = 1·μ + X·α + W·β + ε,        ε ~ N(0, σ²_e · I)
β_j | σ²_j ~ N(0, σ²_j)
σ²_j      ~ InvChi2(ν, S²),     S² = σ²_β / L  (scaled)
σ²_e      ~ InvGamma(a₀_e, b₀_e)

Marginalising the per-marker variance yields a \(t_\nu\)-shrunk effect distribution, which is the defining feature of BayesA relative to ridge regression / BayesC. The prior on σ²_j is informative (default ν = 4.5) so the posterior contracts toward small effects unless data demand otherwise.

Sampling steps (per Gibbs iteration)

  1. Update marker effects \(\beta_j\) one at a time from their full conditionals (Normal with mean and variance involving wtw_diag[j], σ²_e, σ²_j, and the working residual yadj).
  2. Update per-marker variances \(\sigma²_j\) from \(\mathrm{InvChi2}(\nu + 1, (\nu S² + \beta_j²) / (\nu + 1))\).
  3. Update residual variance \(\sigma²_e\) from \(\mathrm{InvGamma}(a_0 + n/2,\, b_0 + \|y - \hat{y}\|^2 / 2)\).
  4. Update intercept \(\mu\) from its normal full conditional given the residual mean.
  5. (Optional) Update fixed effects \(\alpha\) component-wise from normal full conditionals using xtx_diag.
  6. (Binary trait) Albert–Chib step: sample latent liabilities \(z\) from truncated normals consistent with the observed binary response.

State management

BayesARunner keeps a single working residual vector yadj = y - μ - X·α - W·β (or z - … for binary traits) updated incrementally after every coordinate move. This avoids recomputing the full \(W\beta\) product at each iteration, which is the main reason the kernel is fast.

Output

BayesARunner::run returns [BayesAResults] containing posterior samples (after burn-in / thinning) of all parameters, plus posterior means and derived quantities (sigma2_g, h2).

bayesa_em.rs

🦀 masbayes/bayesa_em.rs

BayesA — stochastic EM variant.

Replaces the full Gibbs sweep of [crate::bayesa] with an EM-style coordinate-ascent update that uses posterior means (rather than draws) of the per-marker variances \(\sigma²_j\). Concretely, instead of sampling \(\sigma²_j\) each iteration, the E-step plugs in its conditional expectation and the M-step closes the loop by updating \(\beta\) and \(\sigma²_e\).

When to use

  • Genome-wide screens or large runs where full posterior uncertainty is not needed and only point estimates of marker effects (\(\hat\beta\)) are reported.
  • Cross-validation folds where many BayesA fits are needed and runtime is the bottleneck.

When not to use

  • Reporting credible intervals, posterior distributions of variance components, or heritability uncertainty — these require the full MCMC in [crate::bayesa].
  • Datasets with strong multimodality in the marker-effect posterior; EM only finds a local mode.

Output

Reuses the [BayesAResults] struct so downstream R code does not need to branch on which estimator was used. Sample arrays in the result contain point estimates rather than MCMC traces.

bayesr.rs

🦀 masbayes/bayesr.rs

BayesR Gibbs sampler.

Model

Marker effects follow a four-component mixture of normals (a spike at zero plus three slab components scaled relative to the total genetic variance):

y = 1·μ + X·α + W·β + ε,             ε ~ N(0, σ²_e · I)
β_j | γ_j = c ~ N(0, vᶜ · σ²_g),     vᶜ ∈ {0, 1e-4, 1e-3, 1e-2}
γ_j        ~ Categorical(π),          γ_j ∈ {1, 2, 3, 4}
π          ~ Dirichlet(α₀)
σ²_e, σ²_g ~ InvGamma(·, ·)

Component 1 is the spike (vᶜ = 0, exactly zero effect); components 2–4 are slabs of increasing variance. Sparsity comes from the high prior mass on the spike under the Dirichlet, while large-effect markers escape into one of the slabs.

Sampling steps (per Gibbs iteration)

  1. For each marker \(j\), draw the mixture label \(\gamma_j\) from its full conditional categorical (probabilities proportional to the marginal likelihood under each component × \(\pi_c\)).
  2. Conditional on \(\gamma_j = c\), draw \(\beta_j\) from its normal full conditional (zero when vᶜ = 0).
  3. Update \(\pi\) via Dirichlet–Multinomial conjugacy: \(\pi \mid \gamma \sim \mathrm{Dirichlet}(\alpha_0 + n_c)\) with \(n_c\) tabulated by tabulate(gamma, 4).
  4. Update \(\sigma²_g\) and \(\sigma²_e\) from their inverse-gamma full conditionals.
  5. Update the intercept \(\mu\) and (optional) fixed effects \(\alpha\).
  6. (Binary trait) Albert–Chib augmentation step on latent liabilities.

State management

Identical strategy to BayesA: keep an incremental working residual yadj updated after every coordinate move so \(W\beta\) is never recomputed from scratch within a sweep.

Output

BayesRRunner::run returns [BayesRResults] with posterior samples of all parameters, mixture allocations \(\gamma\), mixture weights \(\pi\), and derived quantities (genetic variance, heritability).

bayesr_em.rs

🦀 masbayes/bayesr_em.rs

BayesR — stochastic EM variant.

Replaces the full Gibbs sweep of [crate::bayesr] with an EM-style update where the mixture allocations \(\gamma_j\) are replaced by their posterior probabilities and the Dirichlet draws of \(\pi\) are replaced by their conditional expectation under the current allocation counts. Marker effects, variance components, and intercept retain closed-form coordinate updates.

When to use

  • Same niche as [crate::bayesa_em]: large screens where point estimates suffice and runtime per fit must be small.
  • Hybrid pipelines that first run EM to warm-start MCMC, then switch to the full [crate::bayesr] sampler for posterior inference.

When not to use

  • Reporting mixture-membership uncertainty (which markers belong to the spike vs. slab) — EM collapses this to soft probabilities at the mode and loses the underlying Bernoulli/Dirichlet posterior.
  • Variance-component inference; same caveats as the BayesA EM variant.

Output

Returns a [BayesRResults] populated with EM point estimates in the sample fields so the R-side API is identical to the MCMC runner.

utils.rs

🦀 masbayes/utils.rs

Shared utilities for the masbayes kernels.

Two unrelated groups of helpers live here:

  1. R ↔︎ ndarray conversion — rmatrix_to_array2* copy an extendr RMatrix into an owned ndarray::Array2 so the MCMC loops can work with native Rust types without going through the R FFI on every cell access. Copy cost is paid once per kernel invocation.
  2. Sampling primitives — rnorm, rinvgamma, rdirichlet, tabulate. These wrap rand_distr distributions so the call sites in bayesa.rs / bayesr.rs stay close to the manuscript notation (e.g. sigma2 ~ InvGamma(a, b) reads as rinvgamma(rng, a, b)).

All sampling functions take an &mut R: Rng so that the caller controls the seed; this is what makes the full MCMC reproducible from a single Pcg64::seed_from_u64(seed) at the start of each runner.

types.rs

🦀 masbayes/types.rs

POD result structs returned by the BayesR and BayesA kernels.

These types are written by the MCMC / EM runners in bayesr.rs, bayesa.rs, bayesr_em.rs, bayesa_em.rs and then marshalled to R lists by the extendr wrappers in lib.rs. Field naming matches the manuscript notation:

  • beta_samples / beta_hat: marker (allele) effects.
  • mu_samples / mu_hat: intercept (grand mean).
  • sigma2_e_*: residual variance.
  • sigma2_j_* (BayesA) and sigma2_{small,medium,large}_* (BayesR): per-marker / per-component variance components.
  • gamma_samples / pi_samples (BayesR only): mixture allocations and their Dirichlet posterior weights.
  • alpha_*: optional fixed-effect coefficients; None when no design matrix X was supplied.
  • z_hat: posterior mean of latent liabilities (Albert–Chib) for binary traits; None for continuous traits.
  • sigma2_g, h2: derived quantities — total additive variance and narrow-sense heritability computed from posterior means.