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, ancestry=None):
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        if ancestry is not None:
46            empty_ancestry = np.empty(shape=(0, *ancestry.shape[1:]), dtype=ancestry.dtype)
47            return children, ages[females], muta_prob[females], empty_ancestry
48        return children, ages[females], muta_prob[females]
49
50    # Random gamete selection (chromatid 0 or 1 per parent)
51    male_gamete_idx = (variables.rng.random(n_pairs) < 0.5).astype(np.int32)
52    female_gamete_idx = (variables.rng.random(n_pairs) < 0.5).astype(np.int32)
53
54    # Assemble children directly from genome array — no intermediate copies
55    children = _assemble_children(
56        genomes.array, males, females, male_gamete_idx, female_gamete_idx,
57    )
58
59    if ancestry is not None:
60        # Assemble offspring ancestry using same parent/gamete selections
61        offspring_ancestry = np.empty((n_pairs, *ancestry.shape[1:]), dtype=ancestry.dtype)
62        offspring_ancestry[:, 0] = ancestry[males, male_gamete_idx]
63        offspring_ancestry[:, 1] = ancestry[females, female_gamete_idx]
64        # TODO fix splitting of ages and muta_prob
65        return children, ages[females], muta_prob[females], offspring_ancestry
66
67    # TODO fix splitting of ages and muta_prob
68    return children, ages[females], muta_prob[females]