remap_pedigree
remap_pedigree (full_pedigree:jax.Array)
Remaps IDs and returns a NumPy array for the sparse builder.
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.
build_g_matrix
build_g_matrix (geno:jax.Array)
Constructs the Genomic Relationship Matrix (G).
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.