Simulating genetic drift using python

drift.py

import random
import sys

generations = 5000
pop_size = 1000
gamete_no = 1000
pop = []

n = len(sys.argv)
if n == 2:
    pop_size = int(sys.argv[1])
elif n == 3:
    pop_size = int(sys.argv[1])
    generations = int(sys.argv[2])

print("Population size:", pop_size)
print("No of generations:", generations)

# Individuals are represented as a tuple pair with alleles at the different loci, e.g., (0,0), (1,1)

# Makes a pre defined number of gametes from the individual provided
def make_gametes(individual):
    g = []
    for i in range(gamete_no):
        g.append(random.choice(individual))
    return g

#  Initialise the population
def init():
    global pop
    for i in range(round(pop_size/2)):
        pop.append((0,0))
    for i in range(round(pop_size/2),pop_size):
        pop.append((1,1))

# Provide gametes from two individuals and it will mate them to prodce random offsprings
def mate(gam1, gam2):
    return (random.choice(gam1), random.choice(gam2))

# Chooses random individuals from the mated list to create the new population
def new_pop(pop_list):
    p_list = []
    for i in range(pop_size):
        p_list.append(random.choice(pop_list))
    return p_list

# Calculates the allele frequencies in the population
def allele_freq(pop_list):
    zero = 0
    one = 0
    for individual in pop_list:
        if individual[0] == 0:
            zero += 1
        else:
            one += 1 
        if individual[1] == 0:
            zero += 1
        else:
            one += 1
    total = zero + one
    zero = round(zero/total,2)
    one = round(one/total,2)
    return (zero,one)


# Main loop
init()
print('Initial allele freq:', str(allele_freq(pop)))
for i in range(generations):
    gamete_list = []
    for individual in pop:
        gamete_list.append(make_gametes(individual))
    pop = []
    for j in range(pop_size*5):
        child = mate(random.choice(gamete_list),random.choice(gamete_list))
        pop.append(child)
    pop = new_pop(pop)
    all_f = allele_freq(pop)
    if all_f[0] == 0 or all_f[1] == 0:
        print("Allele fixed at generation #", i+1)
        break
    else:
    	print('Generation #', i+1)
    	print(allele_freq(pop))