aegis_sim.submodels.reproduction.pairing

 1import numpy as np
 2from numba import njit, prange
 3from aegis_sim import variables
 4from aegis_sim.dataclasses.genomes import Genomes
 5from aegis_sim import submodels
 6
 7
 8@njit(parallel=True)
 9def _assemble_children(genome_array, males, females, male_gamete_idx, female_gamete_idx):
10    """Assemble children genomes directly from parent genome array.
11
12    Reads each parent's selected chromatid and writes it into the children array
13    in one pass, parallelized over pairs. No intermediate arrays allocated.
14    """
15    n_pairs = len(males)
16    # genome_array shape: (n_individuals, ploidy, loci, bpl)
17    n_loci = genome_array.shape[2]
18    n_bpl = genome_array.shape[3]
19    children = np.empty((n_pairs, 2, n_loci, n_bpl), dtype=genome_array.dtype)
20
21    for p in prange(n_pairs):
22        m = males[p]
23        f = females[p]
24        mg = male_gamete_idx[p]
25        fg = female_gamete_idx[p]
26        for i in range(n_loci):
27            for j in range(n_bpl):
28                children[p, 0, i, j] = genome_array[m, mg, i, j]
29                children[p, 1, i, j] = genome_array[f, fg, i, j]
30
31    return children
32
33
34def pairing(genomes: Genomes, parental_sexes, ages, muta_prob):
35    """Return assorted chromatids."""
36
37    # Get pairs
38    males, females = submodels.matingmanager.pair_up_polygamously(parental_sexes)
39    assert len(males) == len(females)
40    n_pairs = len(males)
41
42    if n_pairs == 0:
43        gshape = genomes.shape()
44        children = np.empty(shape=(0, *gshape[1:]), dtype=np.bool_)
45        return children, ages[females], muta_prob[females]
46
47    # Random gamete selection (chromatid 0 or 1 per parent)
48    male_gamete_idx = (variables.rng.random(n_pairs) < 0.5).astype(np.int32)
49    female_gamete_idx = (variables.rng.random(n_pairs) < 0.5).astype(np.int32)
50
51    # Assemble children directly from genome array — no intermediate copies
52    children = _assemble_children(
53        genomes.array, males, females, male_gamete_idx, female_gamete_idx,
54    )
55
56    # TODO fix splitting of ages and muta_prob
57    return children, ages[females], muta_prob[females]