def main(args): # do some initialization ascii_states = ['AB', 'Ab', 'aB', 'ab'] k = len(ascii_states) N = args.N M = np.array(list(gen_states(N, k))) T = get_inverse_dict(M) lmcs = omnomnomial.get_lmcs(M) lps_neutral = np.log(M / float(N)) # write a file for each of the 2x2 option combos for rate_adjustment in (RAW_RATE, ADJUSTED_RATE): for mut_style in (DETERMINISTIC_MUT, STOCHASTIC_MUT): # define the mutation rate between adjacent states if rate_adjustment == RAW_RATE: mu = args.mu elif rate_adjustment == ADJUSTED_RATE: mu = get_adjusted_rate(args.mu) else: raise Exception # construct population genetic transition matrix P if mut_style == DETERMINISTIC_MUT: Q = get_micro_mut_rate_matrix(ascii_states) mmu = scipy.linalg.expm(mu*Q) lps = np.log(np.dot(M, mmu) / N) P = np.exp(omnomnomial.get_log_transition_matrix(M, lmcs, lps)) elif mut_style == STOCHASTIC_MUT: Q_mut = get_macro_mut_rate_matrix(M, T) P_mut = scipy.linalg.expm(mu*Q_mut) P_dft = np.exp(omnomnomial.get_log_transition_matrix( M, lmcs, lps_neutral)) P = np.dot(P_mut, P_dft) else: raise Exception # compute the equilibrium distribution v = get_stationary_distribution(P) # write the equilibrium distribution to the file file_description = '_'.join((rate_adjustment, mut_style)) filename = file_description + '.txt' with open(filename, 'w') as fout: print >> fout, '\t'.join(ascii_states + ['probability']) for i, (AB, Ab, aB, ab) in enumerate(M): row = (AB, Ab, aB, ab, v[i]) print >> fout, '\t'.join(str(x) for x in row)
def main(args): # initialize some variables N = args.N mu = args.mu ascii_states = ['AB', 'Ab', 'aB', 'ab'] k = len(ascii_states) # construct the microstate mutational transition matrix #mmu = get_jeff_mut(ascii_states, mu) Q = get_mut_rate_matrix(ascii_states) mmu = scipy.linalg.expm(mu*Q) # show the transition matrix print 'transition matrix:' print mmu print # construct the population genetic transition matrix M = np.array(list(gen_states(N, k))) lmcs = omnomnomial.get_lmcs(M) lps = np.log(np.dot(M, mmu) / N) P = np.exp(omnomnomial.get_log_transition_matrix(M, lmcs, lps)) # compute the equilibrium distribution v = get_stationary_distribution(P) # print the equilibrium distribution print '\t'.join(ascii_states + ['probability']) for i, (AB, Ab, aB, ab) in enumerate(M): print '\t'.join(str(x) for x in (AB, Ab, aB, ab, v[i]))