foun.py

import numpy as np
from scipy.misc import comb
from math import log10

with open("rosalind_foun.txt","rU") as f:
    line1=f.readline().strip()
    line2=f.readline().strip()

e=map(float,line1.split())
popsize,generations=e[0],e[1]
alleles=map(float,line2.split())
print popsize,generations
print alleles
    
N=2*popsize
# m is the transition matrix for the wright-fisher markov chain
m=np.zeros((N+1,N+1))
for i in np.arange(N+1):
    for j in np.arange(N+1):
        m[j,i]=comb(N,j)*(i/N)**j*((N-i)/N)**(N-j)

b=np.zeros((int(generations),len(alleles)))
for j,x in enumerate(alleles):
    pop=np.zeros(N+1)
    pop[x]=1
    newpop=pop
    for i in np.arange(generations):
        newpop=m.dot(newpop)
        b[i,j]= log10(newpop[0])
    
for i in np.arange(generations):
    for j in range(len(alleles)):
        print b[i,j],
    print