import simuPOP as sim
def checkNumOffspring(numOffspring, ops=[]):
    '''Check the number of offspring for each family using
       information field father_idx
    '''
    pop = sim.Population(size=[30], loci=1, infoFields=['father_idx', 'mother_idx'])
    pop.evolve(
        initOps=[
            sim.InitSex(),
            sim.InitGenotype(freq=[0.5, 0.5]),
        ],
        matingScheme=sim.RandomMating(ops=[
            sim.MendelianGenoTransmitter(),
            sim.ParentsTagger(),
            ] + ops,
            numOffspring=numOffspring),
        gen=1)
    # get the parents of each offspring
    parents = [(x, y) for x, y in zip(pop.indInfo('mother_idx'),
        pop.indInfo('father_idx'))]
    # Individuals with identical parents are considered as siblings.
    famSize = []
    lastParent = (-1, -1)
    for parent in parents:
        if parent == lastParent:
            famSize[-1] += 1
        else:
            lastParent = parent
            famSize.append(1)
    return famSize

# Case 1: produce the given number of offspring
checkNumOffspring(numOffspring=2)
# Case 2: Use a Python function
import random
def func(gen):
    return random.randint(5, 8)

checkNumOffspring(numOffspring=func)
# Case 3: A geometric distribution
checkNumOffspring(numOffspring=(sim.GEOMETRIC_DISTRIBUTION, 0.3))
# Case 4: A Possition distribution
checkNumOffspring(numOffspring=(sim.POISSON_DISTRIBUTION, 1.6))
# Case 5: A Binomial distribution
checkNumOffspring(numOffspring=(sim.BINOMIAL_DISTRIBUTION, 0.1, 10))
# Case 6: A uniform distribution
checkNumOffspring(numOffspring=(sim.UNIFORM_DISTRIBUTION, 2, 6))
# Case 7: With selection on offspring
checkNumOffspring(numOffspring=8,
    ops=[sim.MapSelector(loci=0, fitness={(0,0):1, (0,1):0.8, (1,1):0.5})])