def test_iter_sprs(self):
rho = 1.5e-8 # recomb/site/gen
l = 100000 # length of locus
k = 6 # number of lineages
n = 2*10000 # effective popsize
arg = arglib.sample_arg(k, n, rho, 0, l)
for a, b in izip(arglib.iter_arg_sprs(arg),
arglib.iter_arg_sprs_simple(arg)):
self.assertEqual(a, b)
def test_iter_sprs_remove_thread(self):
rho = 1.5e-8 # recomb/site/gen
l = 100000 # length of locus
k = 6 # number of lineages
n = 2*10000 # effective popsize
arg = arglib.sample_arg(k, n, rho, 0, l)
remove_chroms = set("n%d" % (k-1))
keep = [x for x in arg.leaf_names() if x not in remove_chroms]
arg = arg.copy()
arglib.subarg_by_leaf_names(arg, keep)
for a, b in izip(arglib.iter_arg_sprs(arg),
arglib.iter_arg_sprs_simple(arg)):
self.assertEqual(a, b)
def test_iter_sprs_leaves(self):
rho = 1.5e-8 # recomb/site/gen
l = 100000 # length of locus
k = 40 # number of lineages
n = 2*10000 # effective popsize
arg = arglib.sample_arg(k, n, rho, 0, l)
for a, b in izip(arglib.iter_arg_sprs(arg, use_leaves=True),
arglib.iter_arg_sprs_simple(arg, use_leaves=True)):
a[1][0].sort()
a[2][0].sort()
b[1][0].sort()
b[2][0].sort()
self.assertEqual(a, b)