meiosis

Fill in a module description here

source

meiosis_for_one_cross

 meiosis_for_one_cross (key:<function PRNGKey>, mother_geno:jax.Array,
                        father_geno:jax.Array, n_chr:int,
                        gen_map:jax.Array, v_interference:float)

*Creates a single diploid progeny’s genotype from two parents’ genotypes by simulating meiosis for all chromosomes in parallel.

This function is a high-performance kernel designed to be compiled by JAX.

— JAX Implementation Notes — This function showcases the composition of JAX’s core transformations for maximum performance on parallel architectures (like GPUs).

  1. vmap for Parallelism: The core logic of creating a gamete for a single chromosome is defined in _create_gamete. vmap is used to automatically “vectorize” this function, applying it across all chromosomes of a parent simultaneously. The in_axes argument is critical for this:

    • in_axes=(0, 0, 0, None) tells vmap to map over the first axis (the chromosome axis) of the keys, parent_geno, and gen_map arrays, while broadcasting the single v_interference value to all parallel executions. This avoids unnecessary memory duplication.

  2. jit for Fused Compilation: The entire function is JIT-compiled. JAX is able to “fuse” the vmap operations and the final jnp.stack into a single, highly-optimized kernel. This minimizes overhead from launching separate computations and maximizes hardware utilization.

  3. Static Arguments: The number of chromosomes, n_chr, is used to determine the number of random keys to split via jax.random.split(key, n_chr). Because the shape of a JAX array must be known at compile time, n_chr cannot be a dynamic (traced) value. It is therefore marked as a static argument, meaning JAX will re-compile this function if n_chr changes.

Args: key: A JAX random key. mother_geno: The mother’s genotype. Shape: (nChr, ploidy, nLoci). father_geno: The father’s genotype. Shape: (nChr, ploidy, nLoci). n_chr: The number of chromosomes. Must be a static integer. gen_map: The genetic map defining locus positions for each chromosome. v_interference: The interference parameter for the Gamma process.

Returns: The progeny’s complete diploid genotype. Shape: (nChr, ploidy, nLoci).*