Uni : AF-23 Simulation

"""
Effect of outbreeding on genotype richness under arrhenotokous haploidiploidy.

[] population
() individual

Example uses:

male = ('♂ a')
female = ('♀ ac')

offspring(male, female, True)

has_diploid_males(('♂ aa', '♀ ac', '♂ a', '♂ c'), True)

males = [('aa'), ('b')]
females = [('mn')]

run(*parse(males, females), True, False)

males = [('aa'), ('b')] + [('c')]
females = [('mn')]

run(*parse(males, females), True, False)
"""

def f(sex, alleles):
    return (f'{sex} {alleles}')

def offspring(male, female, sex_locus):
    def a(i):
        return i.split(' ')[-1]
    def s(alleles):
        sex = '♂' if alleles[0]==alleles[1] and sex_locus else '♀'
        return f(sex, ''.join(sorted(alleles)))
    male, female = (a(e) for e in (male, female))
    x = itertools.product(male, female)
    d = (s(i) for i in x)
    h = (f('♂', i) for i in female)
    return tuple(d) + tuple(h)

def has_diploid_males(offspring, count):
    d = (e[0]=='♂' and e[2]==e[3] for e in offspring if len(e)==4)
    return len([e for e in d if e]) if count else any(d)

def all_offspring(males, females, sex_locus, fertile_diploids):
    def o(male, female):
        return offspring(male, female, sex_locus)
    if not fertile_diploids:
        males = (e for e in males if len(e)<4)
    i = itertools.product(males, females)
    a = (o(male, female) for male, female in i)
    x = functools.reduce(lambda x, y: {*x} | {*y}, a, {})
    return tuple(x)

def sex_split(population):
    males = [e for e in population if '♂' in e]
    females = [e for e in population if '♀' in e]
    return males, females

def summary(population):
    m, f = (len(e) for e in sex_split(population))
    w = len(population)
    d = has_diploid_males(population, True)
    return f'genotypes: {m} male, {f} female, {d} diploid male, {w} total'

def run(males, females, sex_locus, fertile_diploids, generations=5):
    """
        Parameters:
        sex_locus: whether the alleles target the sex-determination locus
        fertile_diploids: whether homozygous diploid males can reproduce
    """
    print(f'Generation 0:', end=' ')
    population = tuple(males) + tuple(females)
    print(summary(population))
    print(population, end='\n\n')
    for generation in range(1, generations):
        print(f'Generation {generation}:', end=' ')
        males, females = sex_split(population)
        population = all_offspring(males, females, sex_locus, fertile_diploids) # !
        print(summary(population))
        print(population, end='\n\n')

def parse(males, females):
    males = (f('♂', i) for i in males)
    females = (f('♀', i) for i in females)
    return males, females