source

remap_pedigree

 remap_pedigree (full_pedigree:jax.Array)

Remaps IDs and returns a NumPy array for the sparse builder.


source

solve_multi_trait_mme_iterative

 solve_multi_trait_mme_iterative (y:jax.Array,
                                  relationship_matrix_inv:jax.experimental
                                  .sparse.bcoo.BCOO, G0_inv:jax.Array,
                                  R0_inv:jax.Array, n_traits:int)

Solves a multi-trait MME using an iterative Conjugate Gradient solver.


source

build_g_matrix

 build_g_matrix (geno:jax.Array)

Constructs the Genomic Relationship Matrix (G).


source

build_a_inverse_sparse

 build_a_inverse_sparse (pedigree_meta:numpy.ndarray)

Constructs the inverse of the Numerator Relationship Matrix (A_inv) directly from a pedigree using the correct, numerically stable version of Henderson’s rules. Returns the matrix in JAX’s BCOO sparse format.