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)