# CS150 Population genomics example
from random import choice
def pop_gen(pop_size):
"""
Generate a random population containing two alleles
Args:
pop_size: Size of the population
Returns:
Population as a string
"""
pop = ""
for i in range(pop_size):
# choice randomly chooses from a sequence, e.g. a string
pop += choice("aA")
return pop
def next_gen(pop):
"""
Generate the next generation by randomly sampling from the
current population
Args:
pop: Current population as a string
Returns:
Next generation as a string
"""
next_pop = ""
for i in range(len(pop)):
next_pop += choice(pop)
return next_pop
def pop_sim(pop):
"""
Simulate allele fixation in a population
Args:
pop: Initial population as a string
Returns:
Integer number of generations need to achieve fixation
"""
print(pop)
generations = 0
while "a" in pop and "A" in pop:
pop = next_gen(pop)
print(pop)
generations += 1
return generations