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            parent_positions=None, max_search_radius=0):
36    """Return assorted chromatids.
37
38    When `parent_positions` is given (LATTICE_MODE), the mating manager does
39    lattice-aware expanding-ring pairing. Females who can't find a male within
40    `max_search_radius` rings are not paired. Otherwise (default), classical
41    well-mixed pairing happens.
42
43    Returns `(children, child_ages, child_muta_prob, [child_ancestry], female_slots)`
44    where `female_slots` are the slot indices (into the reproducing pool) of the
45    paired mothers — the caller uses these to look up mother positions when
46    placing offspring on the lattice.
47    """
48
49    # Get pairs
50    males, females = submodels.matingmanager.pair_up_polygamously(
51        parental_sexes,
52        parent_positions=parent_positions,
53        max_search_radius=max_search_radius,
54    )
55    assert len(males) == len(females)
56    n_pairs = len(males)
57
58    if n_pairs == 0:
59        gshape = genomes.shape()
60        children = np.empty(shape=(0, *gshape[1:]), dtype=np.bool_)
61        empty_female_slots = np.empty(0, dtype=np.int64)
62        if ancestry is not None:
63            empty_ancestry = np.empty(shape=(0, *ancestry.shape[1:]), dtype=ancestry.dtype)
64            return children, ages[females], muta_prob[females], empty_ancestry, empty_female_slots
65        return children, ages[females], muta_prob[females], empty_female_slots
66
67    # Random gamete selection (chromatid 0 or 1 per parent)
68    male_gamete_idx = (variables.rng.random(n_pairs) < 0.5).astype(np.int32)
69    female_gamete_idx = (variables.rng.random(n_pairs) < 0.5).astype(np.int32)
70
71    # Assemble children directly from genome array — no intermediate copies
72    children = _assemble_children(
73        genomes.array, males, females, male_gamete_idx, female_gamete_idx,
74    )
75
76    if ancestry is not None:
77        offspring_ancestry = np.empty((n_pairs, *ancestry.shape[1:]), dtype=ancestry.dtype)
78        offspring_ancestry[:, 0] = ancestry[males, male_gamete_idx]
79        offspring_ancestry[:, 1] = ancestry[females, female_gamete_idx]
80        # TODO fix splitting of ages and muta_prob
81        return children, ages[females], muta_prob[females], offspring_ancestry, females
82
83    # TODO fix splitting of ages and muta_prob
84    return children, ages[females], muta_prob[females], females