def test_ploidy_2(self):
ts = msprime.simulate(10, random_seed=1)
assert ts.num_individuals == 0
ts = tsutil.insert_individuals(ts, ploidy=2)
assert ts.num_individuals == 5
for j, ind in enumerate(ts.individuals()):
assert list(ind.nodes) == [2 * j, 2 * j + 1]
def test_ploidy_2(self):
ts = msprime.simulate(10, random_seed=1)
self.assertEqual(ts.num_individuals, 0)
ts = tsutil.insert_individuals(ts, ploidy=2)
self.assertEqual(ts.num_individuals, 5)
for j, ind in enumerate(ts.individuals()):
self.assertEqual(list(ind.nodes), [2 * j, 2 * j + 1])
def test_bad_individuals(self):
ts = msprime.simulate(10, mutation_rate=0.1, random_seed=2)
ts = tsutil.insert_individuals(ts, ploidy=2)
with pytest.raises(ValueError):
ts.write_vcf(io.StringIO(), individuals=[0, -1])
with pytest.raises(ValueError):
ts.write_vcf(io.StringIO(), individuals=[1, 2, ts.num_individuals])
def test_ploidy_2_reversed(self):
ts = msprime.simulate(10, random_seed=1)
assert ts.num_individuals == 0
samples = ts.samples()[::-1]
ts = tsutil.insert_individuals(ts, nodes=samples, ploidy=2)
assert ts.num_individuals == 5
for j, ind in enumerate(ts.individuals()):
assert list(ind.nodes) == [samples[2 * j + 1], samples[2 * j]]
def test_ploidy_2_reversed(self):
ts = msprime.simulate(10, random_seed=1)
self.assertEqual(ts.num_individuals, 0)
samples = ts.samples()[::-1]
ts = tsutil.insert_individuals(ts, samples=samples, ploidy=2)
self.assertEqual(ts.num_individuals, 5)
for j, ind in enumerate(ts.individuals()):
self.assertEqual(list(ind.nodes),
[samples[2 * j + 1], samples[2 * j]])
def test_many_trees_infinite_sites(self):
ts = msprime.simulate(6,
recombination_rate=2,
mutation_rate=2,
random_seed=1)
assert ts.num_sites > 0
assert ts.num_trees > 2
ts = tsutil.insert_individuals(ts, ploidy=2)
self.verify(ts)
def test_simplify_keep_unary(self, ts, nsamples):
np.random.seed(123)
ts = tsutil.mark_metadata(ts, "nodes")
sub_samples = random.sample(list(ts.samples()),
min(nsamples, ts.num_samples))
random_nodes = np.random.choice(ts.num_nodes, ts.num_nodes // 2)
ts = tsutil.insert_individuals(ts, random_nodes)
ts = tsutil.mark_metadata(ts, "individuals")
for params in [{}, {
"keep_unary": True
}, {
"keep_unary_in_individuals": True
}]:
sts = ts.simplify(sub_samples, **params)
# check samples match
assert sts.num_samples == len(sub_samples)
for n, sn in zip(sub_samples, sts.samples()):
assert ts.node(n).metadata == sts.node(sn).metadata
# check that nodes are correctly retained: only nodes ancestral to
# retained samples, and: by default, only coalescent events; if
# keep_unary_in_individuals then also nodes in individuals; if
# keep_unary then all such nodes.
for t in ts.trees(tracked_samples=sub_samples):
st = sts.at(t.interval.left)
visited = [False for _ in sts.nodes()]
for n, sn in zip(sub_samples, sts.samples()):
last_n = t.num_tracked_samples(n)
while n != tskit.NULL:
ind = ts.node(n).individual
keep = False
if t.num_tracked_samples(n) > last_n:
# a coalescent node
keep = True
if "keep_unary_in_individuals" in params and ind != tskit.NULL:
keep = True
if "keep_unary" in params:
keep = True
if (n in sub_samples) or keep:
visited[sn] = True
assert sn != tskit.NULL
assert ts.node(n).metadata == sts.node(sn).metadata
assert t.num_tracked_samples(n) == st.num_samples(
sn)
if ind != tskit.NULL:
sind = sts.node(sn).individual
assert sind != tskit.NULL
assert (ts.individual(ind).metadata ==
sts.individual(sind).metadata)
sn = st.parent(sn)
last_n = t.num_tracked_samples(n)
n = t.parent(n)
st_nodes = list(st.nodes())
for k, v in enumerate(visited):
assert v == (k in st_nodes)
def test_many_trees_sequence_length_infinite_sites(self):
for L in [0.5, 1.5, 3.3333]:
ts = msprime.simulate(6,
length=L,
recombination_rate=2,
mutation_rate=1,
random_seed=1)
assert ts.num_sites > 0
ts = tsutil.insert_individuals(ts, ploidy=2)
self.verify(ts)
def test_wright_fisher_unsimplified_multiple_roots(self):
tables = wf.wf_sim(8,
15,
seed=1,
deep_history=False,
initial_generation_samples=False,
num_loci=20)
tables.sort()
ts = msprime.mutate(tables.tree_sequence(), rate=0.006, random_seed=2)
self.assertGreater(ts.num_sites, 0)
ts = tsutil.insert_individuals(ts, ploidy=2)
self.verify(ts)
def test_wright_fisher_initial_generation(self):
tables = wf.wf_sim(6,
5,
seed=3,
deep_history=True,
initial_generation_samples=True,
num_loci=2)
tables.sort()
tables.simplify()
ts = msprime.mutate(tables.tree_sequence(), rate=0.08, random_seed=2)
assert ts.num_sites > 0
ts = tsutil.insert_individuals(ts, ploidy=3)
self.verify(ts)
def test_wright_fisher_simplified(self):
tables = wf.wf_sim(9,
10,
seed=1,
deep_history=True,
initial_generation_samples=False,
num_loci=5)
tables.sort()
ts = tables.tree_sequence().simplify()
ts = msprime.mutate(ts, rate=0.01, random_seed=1234)
self.assertGreater(ts.num_sites, 0)
ts = tsutil.insert_individuals(ts, ploidy=3)
self.verify(ts)
def test_wright_fisher_unsimplified(self):
tables = wf.wf_sim(
4,
5,
seed=1,
deep_history=True,
initial_generation_samples=False,
num_loci=10,
)
tables.sort()
ts = msprime.mutate(tables.tree_sequence(), rate=0.05, random_seed=234)
assert ts.num_sites > 0
ts = tsutil.insert_individuals(ts, ploidy=4)
self.verify(ts)
def test_bad_length_individuals(self):
ts = msprime.simulate(6, mutation_rate=2, random_seed=1)
self.assertGreater(ts.num_sites, 0)
ts = tsutil.insert_individuals(ts, ploidy=2)
with self.assertRaises(ValueError):
ts.write_vcf(io.StringIO(), individual_names=[])
with self.assertRaises(ValueError):
ts.write_vcf(io.StringIO(),
individual_names=["x" for _ in range(4)])
with self.assertRaises(ValueError):
ts.write_vcf(
io.StringIO(),
individuals=list(range(ts.num_individuals)),
individual_names=["x" for _ in range(ts.num_individuals - 1)])
with self.assertRaises(ValueError):
ts.write_vcf(
io.StringIO(),
individuals=list(range(ts.num_individuals - 1)),
individual_names=["x" for _ in range(ts.num_individuals)])
def test_single_tree_multichar_mutations(self):
ts = msprime.simulate(6, random_seed=1, mutation_rate=1)
ts = tsutil.insert_multichar_mutations(ts)
ts = tsutil.insert_individuals(ts, ploidy=2)
self.verify(ts)
def test_simple_infinite_sites_ploidy_2_even_samples(self):
ts = msprime.simulate(20, mutation_rate=1, random_seed=2)
samples = ts.samples()[0::2]
ts = tsutil.insert_individuals(ts, nodes=samples, ploidy=2)
assert ts.num_sites > 2
self.verify(ts)
def test_simple_infinite_sites_ploidy_2(self):
ts = msprime.simulate(10, mutation_rate=1, random_seed=2)
ts = tsutil.insert_individuals(ts, ploidy=2)
assert ts.num_sites > 2
self.verify(ts)
def test_simple_infinite_sites_ploidy_2_reversed_samples(self):
ts = msprime.simulate(10, mutation_rate=1, random_seed=2)
samples = ts.samples()[::-1]
ts = tsutil.insert_individuals(ts, samples=samples, ploidy=2)
self.assertGreater(ts.num_sites, 2)
self.verify(ts)
