def test_nlineages(self): """ Test lineage counting """ k = 4 n = 1e4 rho = 1.5e-8 * 1 length = 1000 arg = arglib.sample_arg(k, n, rho, start=0, end=length) times = arghmm.get_time_points(ntimes=6) arghmm.discretize_arg(arg, times) tree = arg.get_marginal_tree(0) nlineages, nrecombs, ncoals = arghmm.get_nlineages_recomb_coal( tree, times) treelib.draw_tree_names(tree.get_tree(), scale=4e-3) print list(arghmm.iter_coal_states(tree, times)) print nlineages self.assert_(nlineages == sorted(nlineages, reverse=True)) print nlineages print nrecombs print ncoals
def test_trans_single(self): """ Calculate transition probabilities Only calculate a single matrix """ k = 4 n = 1e4 rho = 1.5e-8 mu = 2.5e-8 length = 1000 arg = arglib.sample_arg(k, n, rho, start=0, end=length) muts = arglib.sample_arg_mutations(arg, mu) seqs = arglib.make_alignment(arg, muts) times = arghmm.get_time_points(10) arghmm.discretize_arg(arg, times) print "recomb", arglib.get_recomb_pos(arg) new_name = "n%d" % (k-1) arg = arghmm.remove_arg_thread(arg, new_name) model = arghmm.ArgHmm(arg, seqs, new_name=new_name, times=times) pos = 10 tree = arg.get_marginal_tree(pos) mat = arghmm.calc_transition_probs( tree, model.states[pos], model.nlineages, model.times, model.time_steps, model.popsizes, rho) print model.states[pos] pc(mat) for row in mat: print sum(map(exp, row))
def test_prior(self): """ Calculate state priors """ k = 10 n = 1e4 rho = 1.5e-8 mu = 2.5e-8 length = 1000 arg = arglib.sample_arg(k, n, rho, start=0, end=length) muts = arglib.sample_arg_mutations(arg, mu) seqs = arglib.make_alignment(arg, muts) times = arghmm.get_time_points() arghmm.discretize_arg(arg, times) new_name = "n%d" % (k - 1) arg = arghmm.remove_arg_thread(arg, new_name) model = arghmm.ArgHmm(arg, seqs, new_name=new_name, times=times) prior = [ model.prob_prior(0, j) for j in xrange(model.get_num_states(0)) ] print prior print sum(map(exp, prior)) fequal(sum(map(exp, prior)), 1.0, rel=.01)
def test_prior(self): """ Calculate state priors """ k = 10 n = 1e4 rho = 1.5e-8 mu = 2.5e-8 length = 1000 arg = arglib.sample_arg(k, n, rho, start=0, end=length) muts = arglib.sample_arg_mutations(arg, mu) seqs = arglib.make_alignment(arg, muts) times = arghmm.get_time_points() arghmm.discretize_arg(arg, times) new_name = "n%d" % (k-1) arg = arghmm.remove_arg_thread(arg, new_name) model = arghmm.ArgHmm(arg, seqs, new_name=new_name, times=times) prior = [model.prob_prior(0, j) for j in xrange(model.get_num_states(0))] print prior print sum(map(exp, prior)) fequal(sum(map(exp, prior)), 1.0, rel=.01)
def test_trans_single(self): """ Calculate transition probabilities Only calculate a single matrix """ k = 4 n = 1e4 rho = 1.5e-8 mu = 2.5e-8 length = 1000 arg = arglib.sample_arg(k, n, rho, start=0, end=length) muts = arglib.sample_arg_mutations(arg, mu) seqs = arglib.make_alignment(arg, muts) times = arghmm.get_time_points(10) arghmm.discretize_arg(arg, times) print "recomb", arglib.get_recomb_pos(arg) new_name = "n%d" % (k - 1) arg = arghmm.remove_arg_thread(arg, new_name) model = arghmm.ArgHmm(arg, seqs, new_name=new_name, times=times) pos = 10 tree = arg.get_marginal_tree(pos) mat = arghmm.calc_transition_probs(tree, model.states[pos], model.nlineages, model.times, model.time_steps, model.popsizes, rho) print model.states[pos] pc(mat) for row in mat: print sum(map(exp, row))
def test_post(self): k = 6 n = 1e4 rho = 1.5e-8 * 10 mu = 2.5e-8 * 10 length = 10000 arg = arglib.sample_arg(k, n, rho, start=0, end=length) muts = arglib.sample_arg_mutations(arg, mu) seqs = arglib.make_alignment(arg, muts) print "muts", len(muts) print "recombs", len(arglib.get_recomb_pos(arg)) times = arghmm.get_time_points(ntimes=10) arghmm.discretize_arg(arg, times) tree = arg.get_marginal_tree(0) treelib.draw_tree_names(tree.get_tree(), minlen=5, scale=4e-4) # remove chrom new_name = "n%d" % (k - 1) keep = set(arg.leaf_names()) - set([new_name]) arglib.subarg_by_leaf_names(arg, keep) arg = arglib.smcify_arg(arg) model = arghmm.ArgHmm(arg, seqs, new_name=new_name, times=times, rho=rho, mu=mu) print "states", len(model.states[0]) probs = arghmm.get_posterior_probs(model, length, verbose=True) for pcol in probs: p = sum(map(exp, pcol)) print p, " ".join("%.3f" % f for f in map(exp, pcol)) fequal(p, 1.0, rel=1e-2)
def test_post_plot(self): k = 6 n = 1e4 rho = 1.5e-8 * 50 mu = 2.5e-8 * 50 length = 10000 arg = arglib.sample_arg(k, n, rho, start=0, end=length) muts = arglib.sample_arg_mutations(arg, mu) seqs = arglib.make_alignment(arg, muts) times = arghmm.get_time_points(ntimes=30) arghmm.discretize_arg(arg, times) pause() # save #arglib.write_arg("test/data/k4.arg", arg) #fasta.write_fasta("test/data/k4.fa", seqs) new_name = "n%d" % (k - 1) thread = list( arghmm.iter_chrom_thread(arg, arg[new_name], by_block=False)) p = plot(cget(thread, 1), style="lines", ymin=times[1], ylog=10) # remove chrom new_name = "n%d" % (k - 1) arg = arghmm.remove_arg_thread(arg, new_name) model = arghmm.ArgHmm(arg, seqs, new_name=new_name, times=times, rho=rho, mu=mu) print "states", len(model.states[0]) print "muts", len(muts) print "recomb", len(model.recomb_pos) - 2, model.recomb_pos[1:-1] p.plot(model.recomb_pos, [10000] * len(model.recomb_pos), style="points") probs = arghmm.get_posterior_probs(model, length, verbose=True) print "done" high = list(arghmm.iter_posterior_times(model, probs, .95)) low = list(arghmm.iter_posterior_times(model, probs, .05)) p.gnuplot("set linestyle 2") p.plot(high, style="lines") p.gnuplot("set linestyle 2") p.plot(low, style="lines") #write_list("test/data/post_real.txt", cget(thread, 1)) #write_list("test/data/post_high.txt", high) #write_list("test/data/post_low.txt", low) pause()
def test_post_plot(self): k = 6 n = 1e4 rho = 1.5e-8 * 50 mu = 2.5e-8 * 50 length = 10000 arg = arglib.sample_arg(k, n, rho, start=0, end=length) muts = arglib.sample_arg_mutations(arg, mu) seqs = arglib.make_alignment(arg, muts) times = arghmm.get_time_points(ntimes=30) arghmm.discretize_arg(arg, times) pause() # save #arglib.write_arg("test/data/k4.arg", arg) #fasta.write_fasta("test/data/k4.fa", seqs) new_name = "n%d" % (k-1) thread = list(arghmm.iter_chrom_thread(arg, arg[new_name], by_block=False)) p = plot(cget(thread, 1), style="lines", ymin=times[1], ylog=10) # remove chrom new_name = "n%d" % (k-1) arg = arghmm.remove_arg_thread(arg, new_name) model = arghmm.ArgHmm(arg, seqs, new_name=new_name, times=times, rho=rho, mu=mu) print "states", len(model.states[0]) print "muts", len(muts) print "recomb", len(model.recomb_pos) - 2, model.recomb_pos[1:-1] p.plot(model.recomb_pos, [10000] * len(model.recomb_pos), style="points") probs = arghmm.get_posterior_probs(model, length, verbose=True) print "done" high = list(arghmm.iter_posterior_times(model, probs, .95)) low = list(arghmm.iter_posterior_times(model, probs, .05)) p.gnuplot("set linestyle 2") p.plot(high, style="lines") p.gnuplot("set linestyle 2") p.plot(low, style="lines") #write_list("test/data/post_real.txt", cget(thread, 1)) #write_list("test/data/post_high.txt", high) #write_list("test/data/post_low.txt", low) pause()
def test_norecomb_plot(self): k = 50 n = 1e4 rho = 1.5e-8 * .0001 rho2 = 1.5e-8 * 10 mu = 2.5e-8 * 100 length = 1000 arg = arglib.sample_arg(k, n, rho, start=0, end=length) muts = arglib.sample_arg_mutations(arg, mu) seqs = arglib.make_alignment(arg, muts) times = arghmm.get_time_points(ntimes=20) arghmm.discretize_arg(arg, times) tree = arg.get_marginal_tree(0) treelib.draw_tree_names(tree.get_tree(), minlen=5, scale=4e-4) # get thread new_name = "n%d" % (k - 1) keep = ["n%d" % i for i in range(k - 1)] arglib.subarg_by_leaf_names(arg, keep) arg.set_ancestral() arg.prune() model = arghmm.ArgHmm(arg, seqs, new_name=new_name, times=times, rho=rho2, mu=mu) print "states", len(model.states[0]) print "muts", len(muts) # simulate a new thread states = list(islice(hmm.sample_hmm_states(model), 0, arg.end)) data = list(hmm.sample_hmm_data(model, states)) seqs[new_name] = "".join(data) #alignlib.print_align(seqs) thread = [ model.times[model.states[i][s][1]] for i, s in enumerate(states) ] p = plot(thread, style="lines") probs = arghmm.get_posterior_probs(model, length, verbose=True) print "done" high = list(arghmm.iter_posterior_times(model, probs, .75)) low = list(arghmm.iter_posterior_times(model, probs, .25)) p.plot(high, style="lines") p.plot(low, style="lines") pause()
def test_post_real(self): k = 3 n = 1e4 rho = 1.5e-8 mu = 2.5e-8 length = 100000 arg = arglib.sample_arg(k, n, rho, start=0, end=length) muts = arglib.sample_arg_mutations(arg, mu) seqs = arglib.make_alignment(arg, muts) #arg = arglib.read_arg("test/data/real.arg") #seqs = fasta.read_fasta("test/data/real.fa") #arglib.write_arg("test/data/real.arg", arg) #fasta.write_fasta("test/data/real.fa", seqs) times = arghmm.get_time_points(maxtime=50000, ntimes=20) arghmm.discretize_arg(arg, times) new_name = "n%d" % (k - 1) thread = list( arghmm.iter_chrom_thread(arg, arg[new_name], by_block=False)) tree = arg.get_marginal_tree(0) print tree.root.age treelib.draw_tree_names(tree.get_tree(), minlen=5, scale=4e-4) p = plot(cget(thread, 1), style="lines", ymin=10, ylog=10) #alignlib.print_align(seqs) # remove chrom keep = ["n%d" % i for i in range(k - 1)] arglib.subarg_by_leaf_names(arg, keep) arg = arglib.smcify_arg(arg) model = arghmm.ArgHmm(arg, seqs, new_name=new_name, times=times, rho=rho, mu=mu) print "states", len(model.states[0]) #print "muts", len(muts) print "recomb", len(model.recomb_pos) - 2, model.recomb_pos[1:-1] probs = arghmm.get_posterior_probs(model, length, verbose=True) high = list(arghmm.iter_posterior_times(model, probs, .95)) low = list(arghmm.iter_posterior_times(model, probs, .05)) p.plot(high, style="lines") p.plot(low, style="lines") pause()
def test_trans_switch_single(self): """ Calculate transitions probabilities for switching between blocks Only calculate a single matrix """ k = 5 n = 1e4 rho = 1.5e-8 * 100 mu = 2.5e-8 length = 1000 arg = arglib.sample_arg(k, n, rho, start=0, end=length) #arglib.write_arg("tmp/a.arg", arg) #arg = arglib.read_arg("tmp/a.arg") #arg.set_ancestral() muts = arglib.sample_arg_mutations(arg, mu) seqs = arglib.make_alignment(arg, muts) times = arghmm.get_time_points(5) arghmm.discretize_arg(arg, times) new_name = "n%d" % (k-1) arg = arghmm.remove_arg_thread(arg, new_name) model = arghmm.ArgHmm(arg, seqs, new_name=new_name, times=times) # get recombs recombs = list(x.pos for x in arghmm.iter_visible_recombs(arg)) print "recomb", recombs pos = recombs[0] + 1 tree = arg.get_marginal_tree(pos-.5) last_tree = arg.get_marginal_tree(pos-1-.5) print "states1>>", model.states[pos-1] print "states2>>", model.states[pos] treelib.draw_tree_names(last_tree.get_tree(), minlen=5, maxlen=5) treelib.draw_tree_names(tree.get_tree(), minlen=5, maxlen=5) print "pos>>", pos recomb = [x for x in tree if x.event == "recomb" and x.pos+1 == pos][0] mat = arghmm.calc_transition_probs_switch( tree, last_tree, recomb.name, model.states[pos-1], model.states[pos], model.nlineages, model.times, model.time_steps, model.popsizes, rho) pc(mat)
def test_norecomb_plot(self): k = 50 n = 1e4 rho = 1.5e-8 * .0001 rho2 = 1.5e-8 * 10 mu = 2.5e-8 * 100 length = 1000 arg = arglib.sample_arg(k, n, rho, start=0, end=length) muts = arglib.sample_arg_mutations(arg, mu) seqs = arglib.make_alignment(arg, muts) times = arghmm.get_time_points(ntimes=20) arghmm.discretize_arg(arg, times) tree = arg.get_marginal_tree(0) treelib.draw_tree_names(tree.get_tree(), minlen=5, scale=4e-4) # get thread new_name = "n%d" % (k-1) keep = ["n%d" % i for i in range(k-1)] arglib.subarg_by_leaf_names(arg, keep) arg.set_ancestral() arg.prune() model = arghmm.ArgHmm(arg, seqs, new_name=new_name, times=times, rho=rho2, mu=mu) print "states", len(model.states[0]) print "muts", len(muts) # simulate a new thread states = list(islice(hmm.sample_hmm_states(model), 0, arg.end)) data = list(hmm.sample_hmm_data(model, states)) seqs[new_name] = "".join(data) #alignlib.print_align(seqs) thread = [model.times[model.states[i][s][1]] for i, s in enumerate(states)] p = plot(thread, style="lines") probs = arghmm.get_posterior_probs(model, length, verbose=True) print "done" high = list(arghmm.iter_posterior_times(model, probs, .75)) low = list(arghmm.iter_posterior_times(model, probs, .25)) p.plot(high, style="lines") p.plot(low, style="lines") pause()
def test_trans_switch_single(self): """ Calculate transitions probabilities for switching between blocks Only calculate a single matrix """ k = 5 n = 1e4 rho = 1.5e-8 * 100 mu = 2.5e-8 length = 1000 arg = arglib.sample_arg(k, n, rho, start=0, end=length) #arglib.write_arg("tmp/a.arg", arg) #arg = arglib.read_arg("tmp/a.arg") #arg.set_ancestral() muts = arglib.sample_arg_mutations(arg, mu) seqs = arglib.make_alignment(arg, muts) times = arghmm.get_time_points(5) arghmm.discretize_arg(arg, times) new_name = "n%d" % (k - 1) arg = arghmm.remove_arg_thread(arg, new_name) model = arghmm.ArgHmm(arg, seqs, new_name=new_name, times=times) # get recombs recombs = list(x.pos for x in arghmm.iter_visible_recombs(arg)) print "recomb", recombs pos = recombs[0] + 1 tree = arg.get_marginal_tree(pos - .5) last_tree = arg.get_marginal_tree(pos - 1 - .5) print "states1>>", model.states[pos - 1] print "states2>>", model.states[pos] treelib.draw_tree_names(last_tree.get_tree(), minlen=5, maxlen=5) treelib.draw_tree_names(tree.get_tree(), minlen=5, maxlen=5) print "pos>>", pos recomb = [x for x in tree if x.event == "recomb" and x.pos + 1 == pos][0] mat = arghmm.calc_transition_probs_switch(tree, last_tree, recomb.name, model.states[pos - 1], model.states[pos], model.nlineages, model.times, model.time_steps, model.popsizes, rho) pc(mat)
def test_post_c(self): k = 3 n = 1e4 rho = 1.5e-8 * 30 mu = 2.5e-8 * 100 length = 100 arg = arglib.sample_arg(k, n, rho, start=0, end=length) arg.prune() muts = arglib.sample_arg_mutations(arg, mu) seqs = arglib.make_alignment(arg, muts) print arglib.get_recomb_pos(arg) print "muts", len(muts) print "recomb", len(arglib.get_recomb_pos(arg)) times = arghmm.get_time_points(ntimes=10) arghmm.discretize_arg(arg, times) tree = arg.get_marginal_tree(0) print tree.root.age treelib.draw_tree_names(tree.get_tree(), minlen=5, scale=4e-4) # remove chrom keep = ["n%d" % i for i in range(k - 1)] arglib.subarg_by_leaf_names(arg, keep) model = arghmm.ArgHmm(arg, seqs, new_name="n%d" % (k - 1), times=times, rho=rho, mu=mu) print "states", len(model.states[0]) util.tic("C") probs1 = list(arghmm.get_posterior_probs(model, length, verbose=True)) util.toc() util.tic("python") probs2 = list(hmm.get_posterior_probs(model, length, verbose=True)) util.toc() print "probs1" pc(probs1) print "probs2" pc(probs2) for col1, col2 in izip(probs1, probs2): for a, b in izip(col1, col2): fequal(a, b)
def test_post_real(self): k = 3 n = 1e4 rho = 1.5e-8 mu = 2.5e-8 length = 100000 arg = arglib.sample_arg(k, n, rho, start=0, end=length) muts = arglib.sample_arg_mutations(arg, mu) seqs = arglib.make_alignment(arg, muts) #arg = arglib.read_arg("test/data/real.arg") #seqs = fasta.read_fasta("test/data/real.fa") #arglib.write_arg("test/data/real.arg", arg) #fasta.write_fasta("test/data/real.fa", seqs) times = arghmm.get_time_points(maxtime=50000, ntimes=20) arghmm.discretize_arg(arg, times) new_name = "n%d" % (k - 1) thread = list(arghmm.iter_chrom_thread(arg, arg[new_name], by_block=False)) tree = arg.get_marginal_tree(0) print tree.root.age treelib.draw_tree_names(tree.get_tree(), minlen=5, scale=4e-4) p = plot(cget(thread, 1), style="lines", ymin=10, ylog=10) #alignlib.print_align(seqs) # remove chrom keep = ["n%d" % i for i in range(k-1)] arglib.subarg_by_leaf_names(arg, keep) arg = arglib.smcify_arg(arg) model = arghmm.ArgHmm(arg, seqs, new_name=new_name, times=times, rho=rho, mu=mu) print "states", len(model.states[0]) #print "muts", len(muts) print "recomb", len(model.recomb_pos) - 2, model.recomb_pos[1:-1] probs = arghmm.get_posterior_probs(model, length, verbose=True) high = list(arghmm.iter_posterior_times(model, probs, .95)) low = list(arghmm.iter_posterior_times(model, probs, .05)) p.plot(high, style="lines") p.plot(low, style="lines") pause()
def test_post_c(self): k = 3 n = 1e4 rho = 1.5e-8 * 30 mu = 2.5e-8 * 100 length = 100 arg = arglib.sample_arg(k, n, rho, start=0, end=length) arg.prune() muts = arglib.sample_arg_mutations(arg, mu) seqs = arglib.make_alignment(arg, muts) print arglib.get_recomb_pos(arg) print "muts", len(muts) print "recomb", len(arglib.get_recomb_pos(arg)) times = arghmm.get_time_points(ntimes=10) arghmm.discretize_arg(arg, times) tree = arg.get_marginal_tree(0) print tree.root.age treelib.draw_tree_names(tree.get_tree(), minlen=5, scale=4e-4) # remove chrom keep = ["n%d" % i for i in range(k-1)] arglib.subarg_by_leaf_names(arg, keep) model = arghmm.ArgHmm(arg, seqs, new_name="n%d" % (k-1), times=times, rho=rho, mu=mu) print "states", len(model.states[0]) util.tic("C") probs1 = list(arghmm.get_posterior_probs(model, length, verbose=True)) util.toc() util.tic("python") probs2 = list(hmm.get_posterior_probs(model, length, verbose=True)) util.toc() print "probs1" pc(probs1) print "probs2" pc(probs2) for col1, col2 in izip(probs1, probs2): for a, b in izip(col1, col2): fequal(a, b)
def test_post3(self): k = 3 n = 1e4 rho = 1.5e-8 * 3 mu = 2.5e-8 * 100 length = 10000 arg = arglib.sample_arg(k, n, rho, start=0, end=length) arg.prune() muts = arglib.sample_arg_mutations(arg, mu) seqs = arglib.make_alignment(arg, muts) times = arghmm.get_time_points(ntimes=10) arghmm.discretize_arg(arg, times) tree = arg.get_marginal_tree(0) treelib.draw_tree_names(tree.get_tree(), minlen=5, scale=4e-4) thread = list(arghmm.iter_chrom_thread(arg, arg["n2"], by_block=False)) p = plot(cget(thread, 1), style="lines", ymin=0) # remove chrom keep = ["n0", "n1"] arglib.subarg_by_leaf_names(arg, keep) arg.set_ancestral() arg.prune() model = arghmm.ArgHmm(arg, seqs, new_name="n2", times=times, rho=rho, mu=mu) print "states", len(model.states[0]) print "muts", len(muts) print "recomb", len(model.recomb_pos) - 2, model.recomb_pos[1:-1] p.plot(model.recomb_pos, [1000] * len(model.recomb_pos), style="points") probs = arghmm.get_posterior_probs(model, length, verbose=True) high = list(arghmm.iter_posterior_times(model, probs, .95)) low = list(arghmm.iter_posterior_times(model, probs, .05)) p.plot(high, style="lines") p.plot(low, style="lines") pause()
def test_post2(self): k = 2 n = 1e4 rho = 1.5e-8 * 10 mu = 2.5e-8 * 10 length = 10000 arg = arglib.sample_arg(k, n, rho, start=0, end=length) muts = arglib.sample_arg_mutations(arg, mu) seqs = arglib.make_alignment(arg, muts) print "muts", len(muts) times = arghmm.get_time_points() arghmm.discretize_arg(arg, times) thread = list(arghmm.iter_chrom_thread(arg, arg["n1"], by_block=False)) tree = arg.get_marginal_tree(0) print tree.root.age treelib.draw_tree_names(tree.get_tree(), minlen=5, scale=4e-4) p = plot(cget(thread, 1), style="lines", ymin=0) #alignlib.print_align(seqs) # remove chrom keep = ["n0"] arglib.subarg_by_leaf_names(arg, keep) arg = arglib.smcify_arg(arg) model = arghmm.ArgHmm(arg, seqs, new_name="n1", times=times, rho=rho, mu=mu) print "states", len(model.states[0]) probs = arghmm.get_posterior_probs(model, length, verbose=True) high = list(arghmm.iter_posterior_times(model, probs, .95)) low = list(arghmm.iter_posterior_times(model, probs, .05)) p.plot(high, style="lines") p.plot(low, style="lines") pause()
def test_states(self): """ Test state enumeration """ k = 2 n = 1e4 rho = 1.5e-8 * 100 length = 1000 for i in xrange(20): arg = arglib.sample_arg(k, n, rho, start=0, end=length) times = arghmm.get_time_points(10) arghmm.discretize_arg(arg, times) tree = arg.get_marginal_tree(0) states = list(arghmm.iter_coal_states(tree, times)) treelib.draw_tree_names(tree.get_tree(), scale=4e-4, minlen=6, maxlen=6) print states
def test_backward(self): """ Run backward algorithm """ k = 3 n = 1e4 rho = 1.5e-8 * 100 mu = 2.5e-8 * 100 length = 10000 arg = arglib.sample_arg(k, n, rho, start=0, end=length) muts = arglib.sample_arg_mutations(arg, mu) seqs = arglib.make_alignment(arg, muts) times = arghmm.get_time_points(ntimes=10) arghmm.discretize_arg(arg, times) tree = arg.get_marginal_tree(0) print tree.root.age treelib.draw_tree_names(tree.get_tree(), minlen=5, scale=4e-4) # remove chrom new_name = "n%d" % (k - 1) arg = arghmm.remove_arg_thread(arg, new_name) model = arghmm.ArgHmm(arg, seqs, new_name=new_name, times=times, rho=rho, mu=mu) print "states", len(model.states[0]) print "recomb", model.recomb_pos print "muts", len(muts) probs = hmm.backward_algorithm(model, length, verbose=True) for pcol in probs: p = sum(map(exp, pcol)) print p, " ".join("%.3f" % f for f in map(exp, pcol))
def test_emit_argmax(self): """ Calculate emission probabilities """ k = 10 n = 1e4 rho = 0.0 mu = 2.5e-8 * 100 length = 10000 arg = arglib.sample_arg(k, n, rho, start=0, end=length) muts = arglib.sample_arg_mutations(arg, mu) seqs = arglib.make_alignment(arg, muts) times = arghmm.get_time_points(10) arghmm.discretize_arg(arg, times) new_name = "n%d" % (k - 1) thread = list(arghmm.iter_chrom_thread(arg, arg[new_name])) arg = arghmm.remove_arg_thread(arg, new_name) model = arghmm.ArgHmm(arg, seqs, new_name=new_name, times=times) nstates = model.get_num_states(1) probs = [0.0 for j in xrange(nstates)] for i in xrange(1, length): if i % 100 == 0: print i for j in xrange(nstates): probs[j] += model.prob_emission(i, j) print # is the maximum likelihood emission matching truth data = sorted(zip(probs, model.states[0]), reverse=True) pc(data[:20]) state = (thread[0][0], times.index(thread[0][1])) print data[0][1], state assert data[0][1] == state
def test_emit_argmax(self): """ Calculate emission probabilities """ k = 10 n = 1e4 rho = 0.0 mu = 2.5e-8 * 100 length = 10000 arg = arglib.sample_arg(k, n, rho, start=0, end=length) muts = arglib.sample_arg_mutations(arg, mu) seqs = arglib.make_alignment(arg, muts) times = arghmm.get_time_points(10) arghmm.discretize_arg(arg, times) new_name = "n%d" % (k-1) thread = list(arghmm.iter_chrom_thread(arg, arg[new_name])) arg = arghmm.remove_arg_thread(arg, new_name) model = arghmm.ArgHmm(arg, seqs, new_name=new_name, times=times) nstates = model.get_num_states(1) probs = [0.0 for j in xrange(nstates)] for i in xrange(1, length): if i % 100 == 0: print i for j in xrange(nstates): probs[j] += model.prob_emission(i, j) print # is the maximum likelihood emission matching truth data = sorted(zip(probs, model.states[0]), reverse=True) pc(data[:20]) state = (thread[0][0], times.index(thread[0][1])) print data[0][1], state assert data[0][1] == state
def test_trans(self): """ Calculate transition probabilities for k=2 Only calculate a single matrix """ k = 4 n = 1e4 rho = 1.5e-8 * 20 mu = 2.5e-8 * 20 length = 1000 times = arghmm.get_time_points(ntimes=4, maxtime=200000) popsizes = [n] * len(times) arg = arglib.sample_arg(k, 2 * n, rho, start=0, end=length) arghmm.discretize_arg(arg, times) pos = 10 tree = arg.get_marginal_tree(pos) assert arghmm.assert_transition_probs(tree, times, popsizes, rho)
def test_trans(self): """ Calculate transition probabilities for k=2 Only calculate a single matrix """ k = 4 n = 1e4 rho = 1.5e-8 * 20 mu = 2.5e-8 * 20 length = 1000 times = arghmm.get_time_points(ntimes=4, maxtime=200000) popsizes = [n] * len(times) arg = arglib.sample_arg(k, 2*n, rho, start=0, end=length) arghmm.discretize_arg(arg, times) pos = 10 tree = arg.get_marginal_tree(pos) assert arghmm.assert_transition_probs(tree, times, popsizes, rho)
def test_recomb(self): """ Investigate the fact that some recombinations are not visible """ k = 3 n = 1e4 rho = 1.5e-8 * 20 mu = 2.5e-8 length = 1000 arg = arglib.sample_arg(k, n, rho, start=0, end=length) times = arghmm.get_time_points(10) arghmm.discretize_arg(arg, times) arg.set_ancestral() arg.prune() recombs = arglib.get_recomb_pos(arg) # find recombs by walking recombs2 = [] i = 0 while True: tree = arg.get_marginal_tree(i - .5) recomb = arghmm.find_tree_next_recomb(tree, i + 1, tree=True) if recomb: recombs2.append(recomb.pos) i = recomb.pos else: break # these are suppose to differ because some recombination occur # in the hole of ancestral sequence intervals print recombs print recombs2 arglib.write_arg("tmp/b.arg", arg)
def test_recomb(self): """ Investigate the fact that some recombinations are not visible """ k = 3 n = 1e4 rho = 1.5e-8 * 20 mu = 2.5e-8 length = 1000 arg = arglib.sample_arg(k, n, rho, start=0, end=length) times = arghmm.get_time_points(10) arghmm.discretize_arg(arg, times) arg.set_ancestral() arg.prune() recombs = arglib.get_recomb_pos(arg) # find recombs by walking recombs2 = [] i = 0 while True: tree = arg.get_marginal_tree(i-.5) recomb = arghmm.find_tree_next_recomb(tree, i+1, tree=True) if recomb: recombs2.append(recomb.pos) i = recomb.pos else: break # these are suppose to differ because some recombination occur # in the hole of ancestral sequence intervals print recombs print recombs2 arglib.write_arg("tmp/b.arg", arg)
def test_remove_thread(self): """ Remove a leaf from an ARG """ k = 3 n = 1e4 rho = 1.5e-8 * 20 mu = 2.5e-8 length = 1000 arg = arglib.sample_arg(k, n, rho, start=0, end=length) times = arghmm.get_time_points(10) arghmm.discretize_arg(arg, times) chrom = "n%d" % (k - 1) arg = arghmm.remove_arg_thread(arg, chrom) recomb = arglib.get_recomb_pos(arg) print "recomb", recomb tree = arg.get_marginal_tree(-.5) draw_tree_names(tree.get_tree(), minlen=5, maxlen=5) print sorted([(x.pos, x.event) for x in tree])
def test_remove_thread(self): """ Remove a leaf from an ARG """ k = 3 n = 1e4 rho = 1.5e-8 * 20 mu = 2.5e-8 length = 1000 arg = arglib.sample_arg(k, n, rho, start=0, end=length) times = arghmm.get_time_points(10) arghmm.discretize_arg(arg, times) chrom = "n%d" % (k-1) arg = arghmm.remove_arg_thread(arg, chrom) recomb = arglib.get_recomb_pos(arg) print "recomb", recomb tree = arg.get_marginal_tree(-.5) draw_tree_names(tree.get_tree(), minlen=5, maxlen=5) print sorted([(x.pos, x.event) for x in tree])
def test_backward(self): """ Run backward algorithm """ k = 3 n = 1e4 rho = 1.5e-8 * 100 mu = 2.5e-8 * 100 length = 10000 arg = arglib.sample_arg(k, n, rho, start=0, end=length) muts = arglib.sample_arg_mutations(arg, mu) seqs = arglib.make_alignment(arg, muts) times = arghmm.get_time_points(ntimes=10) arghmm.discretize_arg(arg, times) tree = arg.get_marginal_tree(0) print tree.root.age treelib.draw_tree_names(tree.get_tree(), minlen=5, scale=4e-4) # remove chrom new_name = "n%d" % (k-1) arg = arghmm.remove_arg_thread(arg, new_name) model = arghmm.ArgHmm(arg, seqs, new_name=new_name, times=times, rho=rho, mu=mu) print "states", len(model.states[0]) print "recomb", model.recomb_pos print "muts", len(muts) probs = hmm.backward_algorithm(model, length, verbose=True) for pcol in probs: p = sum(map(exp, pcol)) print p, " ".join("%.3f" % f for f in map(exp, pcol))
n = 1e4 rho = 1.5e-8 * 100 mu = 2.5e-8 length = 1000 arg = arglib.sample_arg(k, n, rho, start=0, end=length) arglib.write_arg("tmp/a.arg", arg) muts = arglib.sample_arg_mutations(arg, mu) seqs = arglib.make_alignment(arg, muts) #arg = arglib.read_arg("tmp/a.arg") #arg.set_ancestral() #find_recomb_coal(tree, last_tree, recomb_name=None, pos=None) times = arghmm.get_time_points(30, maxtime=60e3) arghmm.discretize_arg(arg, times) # get recombs recombs = list(x.pos for x in arghmm.iter_visible_recombs(arg)) print "recomb", recombs pos = recombs[0] + 1 tree = arg.get_marginal_tree(pos-.5) last_tree = arg.get_marginal_tree(pos-1-.5) r, c = arghmm.find_recomb_coal(tree, last_tree, pos=pos) #treelib.draw_tree_names(last_tree.get_tree(), minlen=5, maxlen=5) #treelib.draw_tree_names(tree.get_tree(), minlen=5, maxlen=5) model = arghmm.ArgHmm(arg, seqs, new_name="n%d" % (k-1), times=times)
n = 1e4 rho = 1.5e-8 * 100 mu = 2.5e-8 length = 1000 arg = arglib.sample_arg(k, n, rho, start=0, end=length) arglib.write_arg("tmp/a.arg", arg) muts = arglib.sample_arg_mutations(arg, mu) seqs = arglib.make_alignment(arg, muts) #arg = arglib.read_arg("tmp/a.arg") #arg.set_ancestral() #find_recomb_coal(tree, last_tree, recomb_name=None, pos=None) times = arghmm.get_time_points(30, maxtime=60e3) arghmm.discretize_arg(arg, times) # get recombs recombs = list(x.pos for x in arghmm.iter_visible_recombs(arg)) print "recomb", recombs pos = recombs[0] + 1 tree = arg.get_marginal_tree(pos - .5) last_tree = arg.get_marginal_tree(pos - 1 - .5) r, c = arghmm.find_recomb_coal(tree, last_tree, pos=pos) #treelib.draw_tree_names(last_tree.get_tree(), minlen=5, maxlen=5) #treelib.draw_tree_names(tree.get_tree(), minlen=5, maxlen=5) model = arghmm.ArgHmm(arg, seqs, new_name="n%d" % (k - 1), times=times)
def test_determ(self): k = 8 n = 1e4 rho = 1.5e-8 mu = 2.5e-8 length = 100000 arg = arglib.sample_arg(k, n, rho, start=0, end=length) muts = arglib.sample_arg_mutations(arg, mu) seqs = arglib.make_alignment(arg, muts) times = arghmm.get_time_points(maxtime=50000, ntimes=20) arghmm.discretize_arg(arg, times) new_name = "n%d" % (k - 1) thread = list(arghmm.iter_chrom_thread(arg, arg[new_name], by_block=False)) thread_clades = list(arghmm.iter_chrom_thread( arg, arg[new_name], by_block=True, use_clades=True)) # remove chrom keep = ["n%d" % i for i in range(k-1)] arglib.subarg_by_leaf_names(arg, keep) arg = arglib.smcify_arg(arg) model = arghmm.ArgHmm(arg, seqs, new_name=new_name, times=times, rho=rho, mu=mu) for i, rpos in enumerate(model.recomb_pos[1:-1]): pos = rpos + 1 model.check_local_tree(pos, force=True) #recomb = arghmm.find_tree_next_recomb(arg, pos - 1) tree = arg.get_marginal_tree(pos-.5) last_tree = arg.get_marginal_tree(pos-1-.5) states1 = model.states[pos-1] states2 = model.states[pos] (recomb_branch, recomb_time), (coal_branch, coal_time) = \ arghmm.find_recomb_coal(tree, last_tree, pos=rpos) recomb_time = times.index(recomb_time) coal_time = times.index(coal_time) determ = arghmm.get_deterministic_transitions( states1, states2, times, tree, last_tree, recomb_branch, recomb_time, coal_branch, coal_time) leaves1, time1, block1 = thread_clades[i] leaves2, time2, block2 = thread_clades[i+1] if new_name in leaves1: leaves1.remove(new_name) if new_name in leaves2: leaves2.remove(new_name) node1 = arghmm.arg_lca(arg, leaves1, None, pos-1).name node2 = arghmm.arg_lca(arg, leaves2, None, pos).name state1 = (node1, times.index(time1)) state2 = (node2, times.index(time2)) print pos, state1, state2 try: statei1 = states1.index(state1) statei2 = states2.index(state2) except: print "states1", states1 print "states2", states2 raise statei3 = determ[statei1] print " ", statei1, statei2, statei3, states2[statei3] if statei2 != statei3 and statei3 != -1: tree = tree.get_tree() treelib.remove_single_children(tree) last_tree = last_tree.get_tree() treelib.remove_single_children(last_tree) print "block1", block1 print "block2", block2 print "r=", (recomb_branch, recomb_time) print "c=", (coal_branch, coal_time) print "tree" treelib.draw_tree_names(tree, minlen=8, maxlen=8) print "last_tree" treelib.draw_tree_names(last_tree, minlen=8, maxlen=8) assert False
def test_determ(self): k = 8 n = 1e4 rho = 1.5e-8 mu = 2.5e-8 length = 100000 arg = arglib.sample_arg(k, n, rho, start=0, end=length) muts = arglib.sample_arg_mutations(arg, mu) seqs = arglib.make_alignment(arg, muts) times = arghmm.get_time_points(maxtime=50000, ntimes=20) arghmm.discretize_arg(arg, times) new_name = "n%d" % (k - 1) thread = list( arghmm.iter_chrom_thread(arg, arg[new_name], by_block=False)) thread_clades = list( arghmm.iter_chrom_thread(arg, arg[new_name], by_block=True, use_clades=True)) # remove chrom keep = ["n%d" % i for i in range(k - 1)] arglib.subarg_by_leaf_names(arg, keep) arg = arglib.smcify_arg(arg) model = arghmm.ArgHmm(arg, seqs, new_name=new_name, times=times, rho=rho, mu=mu) for i, rpos in enumerate(model.recomb_pos[1:-1]): pos = rpos + 1 model.check_local_tree(pos, force=True) #recomb = arghmm.find_tree_next_recomb(arg, pos - 1) tree = arg.get_marginal_tree(pos - .5) last_tree = arg.get_marginal_tree(pos - 1 - .5) states1 = model.states[pos - 1] states2 = model.states[pos] (recomb_branch, recomb_time), (coal_branch, coal_time) = \ arghmm.find_recomb_coal(tree, last_tree, pos=rpos) recomb_time = times.index(recomb_time) coal_time = times.index(coal_time) determ = arghmm.get_deterministic_transitions( states1, states2, times, tree, last_tree, recomb_branch, recomb_time, coal_branch, coal_time) leaves1, time1, block1 = thread_clades[i] leaves2, time2, block2 = thread_clades[i + 1] if new_name in leaves1: leaves1.remove(new_name) if new_name in leaves2: leaves2.remove(new_name) node1 = arghmm.arg_lca(arg, leaves1, None, pos - 1).name node2 = arghmm.arg_lca(arg, leaves2, None, pos).name state1 = (node1, times.index(time1)) state2 = (node2, times.index(time2)) print pos, state1, state2 try: statei1 = states1.index(state1) statei2 = states2.index(state2) except: print "states1", states1 print "states2", states2 raise statei3 = determ[statei1] print " ", statei1, statei2, statei3, states2[statei3] if statei2 != statei3 and statei3 != -1: tree = tree.get_tree() treelib.remove_single_children(tree) last_tree = last_tree.get_tree() treelib.remove_single_children(last_tree) print "block1", block1 print "block2", block2 print "r=", (recomb_branch, recomb_time) print "c=", (coal_branch, coal_time) print "tree" treelib.draw_tree_names(tree, minlen=8, maxlen=8) print "last_tree" treelib.draw_tree_names(last_tree, minlen=8, maxlen=8) assert False
def test_trans2(self): """ Calculate transition probabilities for k=2 Only calculate a single matrix """ k = 2 n = 1e4 rho = 1.5e-8 * 20 mu = 2.5e-8 * 20 length = 1000 times = arghmm.get_time_points(ntimes=5, maxtime=200000) arg = arglib.sample_arg(k, 2*n, rho, start=0, end=length) muts = arglib.sample_arg_mutations(arg, mu) seqs = arglib.make_alignment(arg, muts) arghmm.discretize_arg(arg, times) print "recomb", arglib.get_recomb_pos(arg) new_name = "n%d" % (k-1) arg = arghmm.make_trunk_arg(0, length, "n0") model = arghmm.ArgHmm(arg, seqs, new_name=new_name, popsize=n, rho=rho, mu=mu, times=times) pos = 10 tree = arg.get_marginal_tree(pos) model.check_local_tree(pos, force=True) mat = arghmm.calc_transition_probs( tree, model.states[pos], model.nlineages, model.times, model.time_steps, model.popsizes, rho) states = model.states[pos] nstates = len(states) def coal(j): return 1.0 - exp(-model.time_steps[j]/(2.0 * n)) def recoal2(k, j): p = coal(j) for m in range(k, j): p *= 1.0 - coal(m) return p def recoal(k, j): if j == nstates-1: return exp(- sum(model.time_steps[m] / (2.0 * n) for m in range(k, j))) else: return ((1.0 - exp(-model.time_steps[j]/(2.0 * n))) * exp(- sum(model.time_steps[m] / (2.0 * n) for m in range(k, j)))) def isrecomb(i): return 1.0 - exp(-max(rho * 2.0 * model.times[i], rho)) def recomb(i, k): treelen = 2*model.times[i] + model.time_steps[i] if k < i: return 2.0 * model.time_steps[k] / treelen / 2.0 else: return model.time_steps[k] / treelen / 2.0 def trans(i, j): a = states[i][1] b = states[j][1] p = sum(recoal(k, b) * recomb(a, k) for k in range(0, min(a, b)+1)) p += sum(recoal(k, b) * recomb(a, k) for k in range(0, min(a, b)+1)) p *= isrecomb(a) if i == j: p += 1.0 - isrecomb(a) return p for i in range(len(states)): for j in range(len(states)): print isrecomb(states[i][1]) print states[i], states[j], mat[i][j], log(trans(i, j)) fequal(mat[i][j], log(trans(i, j))) # recombs add up to 1 fequal(sum(recomb(i, k) for k in range(i+1)), 0.5) # recoal add up to 1 fequal(sum(recoal(i, j) for j in range(i, nstates)), 1.0) # recomb * recoal add up to .5 fequal(sum(sum(recoal(k, j) * recomb(i, k) for k in range(0, min(i, j)+1)) for j in range(0, nstates)), 0.5) fequal(sum(trans(i, j) for j in range(len(states))), 1.0)
def test_trans2(self): """ Calculate transition probabilities for k=2 Only calculate a single matrix """ k = 2 n = 1e4 rho = 1.5e-8 * 20 mu = 2.5e-8 * 20 length = 1000 times = arghmm.get_time_points(ntimes=5, maxtime=200000) arg = arglib.sample_arg(k, 2 * n, rho, start=0, end=length) muts = arglib.sample_arg_mutations(arg, mu) seqs = arglib.make_alignment(arg, muts) arghmm.discretize_arg(arg, times) print "recomb", arglib.get_recomb_pos(arg) new_name = "n%d" % (k - 1) arg = arghmm.make_trunk_arg(0, length, "n0") model = arghmm.ArgHmm(arg, seqs, new_name=new_name, popsize=n, rho=rho, mu=mu, times=times) pos = 10 tree = arg.get_marginal_tree(pos) model.check_local_tree(pos, force=True) mat = arghmm.calc_transition_probs(tree, model.states[pos], model.nlineages, model.times, model.time_steps, model.popsizes, rho) states = model.states[pos] nstates = len(states) def coal(j): return 1.0 - exp(-model.time_steps[j] / (2.0 * n)) def recoal2(k, j): p = coal(j) for m in range(k, j): p *= 1.0 - coal(m) return p def recoal(k, j): if j == nstates - 1: return exp(-sum(model.time_steps[m] / (2.0 * n) for m in range(k, j))) else: return ((1.0 - exp(-model.time_steps[j] / (2.0 * n))) * exp(-sum(model.time_steps[m] / (2.0 * n) for m in range(k, j)))) def isrecomb(i): return 1.0 - exp(-max(rho * 2.0 * model.times[i], rho)) def recomb(i, k): treelen = 2 * model.times[i] + model.time_steps[i] if k < i: return 2.0 * model.time_steps[k] / treelen / 2.0 else: return model.time_steps[k] / treelen / 2.0 def trans(i, j): a = states[i][1] b = states[j][1] p = sum( recoal(k, b) * recomb(a, k) for k in range(0, min(a, b) + 1)) p += sum( recoal(k, b) * recomb(a, k) for k in range(0, min(a, b) + 1)) p *= isrecomb(a) if i == j: p += 1.0 - isrecomb(a) return p for i in range(len(states)): for j in range(len(states)): print isrecomb(states[i][1]) print states[i], states[j], mat[i][j], log(trans(i, j)) fequal(mat[i][j], log(trans(i, j))) # recombs add up to 1 fequal(sum(recomb(i, k) for k in range(i + 1)), 0.5) # recoal add up to 1 fequal(sum(recoal(i, j) for j in range(i, nstates)), 1.0) # recomb * recoal add up to .5 fequal( sum( sum( recoal(k, j) * recomb(i, k) for k in range(0, min(i, j) + 1)) for j in range(0, nstates)), 0.5) fequal(sum(trans(i, j) for j in range(len(states))), 1.0)