aegis_sim.submodels.reproduction.recombination

  1import numpy as np
  2from numba import njit
  3from aegis_sim import variables
  4from aegis_sim.utilities.funcs import profile_time
  5
  6
  7def recombination(genomes, RECOMBINATION_RATE):
  8    """Return recombined chromatids."""
  9
 10    if RECOMBINATION_RATE == 0:
 11        return genomes
 12
 13    # Recombine two chromatids but pass only one;
 14    #   thus double the number of chromatids, recobine,
 15    #   then return only one chromatid from each chromatid pair
 16    genomes = genomes[np.repeat(np.arange(len(genomes)), 2)]
 17
 18    # Flatten loci and bits
 19    flat_genomes = genomes.reshape(len(genomes), 2, -1)
 20
 21    # Get chromatids
 22    chromatid1 = flat_genomes[:, 0]
 23    chromatid2 = flat_genomes[:, 1]
 24
 25    # Make choice array: when to take recombined and when to take original loci
 26    # -1 means synapse; +1 means clear
 27    rr = RECOMBINATION_RATE / 2  # / 2 because you are generating two random vectors (fwd and bkd)
 28    reco_fwd = (variables.rng.random(chromatid1.shape) < rr) * -2 + 1
 29    reco_bkd = (variables.rng.random(chromatid2.shape) < rr) * -2 + 1
 30
 31    # Propagate synapse
 32    reco_fwd_cum = np.cumprod(reco_fwd, axis=1)
 33    reco_bkd_cum = np.cumprod(reco_bkd[:, ::-1], axis=1)[:, ::-1]
 34
 35    # Recombine if both sites recombining
 36    reco_final = (reco_fwd_cum + reco_bkd_cum) == -2
 37
 38    # Choose bits from first or second chromatid
 39    # recombined = np.empty(flat_genomes.shape, bool)
 40    recombined = np.empty(flat_genomes.shape, dtype=np.bool_)
 41    recombined[:, 0] = np.where(reco_final, chromatid2, chromatid1)
 42    recombined[:, 1] = np.where(reco_final, chromatid1, chromatid2)
 43
 44    recombined = recombined.reshape(genomes.shape)
 45    recombined = recombined[::2]  # Look at first comment in the function
 46
 47    return recombined
 48
 49
 50# # Loop version of the vectorized function above
 51# def recombination_via_pairs(genomes, RECOMBINATION_RATE):
 52
 53#     if RECOMBINATION_RATE == 0:
 54#         return genomes
 55
 56#     flat_genomes = genomes.reshape(len(genomes), 2, -1)
 57
 58#     n_sites = flat_genomes.shape[-1]
 59
 60#     n_recombination_sites = np.random.binomial(
 61#         n=n_sites,
 62#         p=RECOMBINATION_RATE,
 63#         size=len(flat_genomes),
 64#     )
 65
 66#     # Produce all random numbers immediately
 67#     chiasmata_list = variables.rng.integers(
 68#         low=1,
 69#         high=n_sites,
 70#         size=(len(n_recombination_sites), max(n_recombination_sites)),
 71#         dtype=np.int32,
 72#     )  # [low, high)
 73
 74#     for i, (chiasmata, n) in enumerate(zip(chiasmata_list, n_recombination_sites)):
 75#         for chiasma in chiasmata[:n]:
 76#             flat_genomes[i, 0, :chiasma], flat_genomes[i, 1, :chiasma] = (
 77#                 flat_genomes[i, 1, :chiasma],
 78#                 flat_genomes[i, 0, :chiasma],
 79#             )
 80
 81#     unflattened_genomes = flat_genomes.reshape(genomes.shape)
 82
 83#     return unflattened_genomes
 84
 85
 86@njit
 87def recombination_via_pairs_numba(flat_genomes, n_recombination_sites, chiasmata_list):
 88    for i in range(len(flat_genomes)):
 89        for j in range(n_recombination_sites[i]):
 90            chiasma = chiasmata_list[i, j]
 91            flat_genomes[i, 0, :chiasma], flat_genomes[i, 1, :chiasma] = (
 92                flat_genomes[i, 1, :chiasma].copy(),
 93                flat_genomes[i, 0, :chiasma].copy(),
 94            )
 95    return flat_genomes
 96
 97
 98def recombination_via_pairs(genomes, RECOMBINATION_RATE):
 99    if RECOMBINATION_RATE == 0:
100        return genomes
101
102    flat_genomes = genomes.reshape(len(genomes), 2, -1).copy()
103    n_sites = flat_genomes.shape[-1]
104    n_recombination_sites = np.random.binomial(n=n_sites, p=RECOMBINATION_RATE, size=len(flat_genomes))
105
106    max_n = max(n_recombination_sites)
107    chiasmata_list = variables.rng.integers(low=1, high=n_sites, size=(len(flat_genomes), max_n), dtype=np.int32)
108
109    flat_genomes = recombination_via_pairs_numba(flat_genomes, n_recombination_sites, chiasmata_list)
110
111    return flat_genomes.reshape(genomes.shape)