def verify(self, ts, sample_sets):
A1 = self.simple_mean_sample_ancestry(ts, sample_sets)
A2 = tsinfer.mean_sample_ancestry(ts, sample_sets)
self.assertEqual(A1.shape, A2.shape)
# for tree in ts.trees():
# print(tree.interval)
# print(tree.draw(format="unicode"))
# print()
# for node in ts.nodes():
# print(node.id, np.sum(A2[:, node.id]), A2[:, node.id], sep="\t")
self.assertTrue(np.allclose(A1, A2))
S = np.sum(A1, axis=0)
self.assertTrue(np.allclose(S[S != 0], 1))
return A1
def test_span_zero_nodes(self):
ts = msprime.simulate(10, random_seed=1)
tables = ts.dump_tables()
# Add in a few unreferenced nodes.
u1 = tables.nodes.add_row(flags=0, time=1234)
u2 = tables.nodes.add_row(flags=1, time=1234)
ts = tables.tree_sequence()
sample_sets = [[j] for j in range(10)]
A1 = self.simple_mean_sample_ancestry(ts, sample_sets)
A2 = tsinfer.mean_sample_ancestry(ts, sample_sets)
S = np.sum(A1, axis=0)
self.assertTrue(np.allclose(A1, A2))
self.assertTrue(np.allclose(S[:u1], 1))
self.assertTrue(np.all(A1[:, u1] == 0))
self.assertTrue(np.all(A1[:, u2] == 0))
def test_two_populations_overlapping_samples(self):
ts = self.two_populations_high_migration_example()
with self.assertRaises(ValueError):
tsinfer.mean_sample_ancestry(ts, [[1], [1]])
with self.assertRaises(ValueError):
tsinfer.mean_sample_ancestry(ts, [[1, 1], [2]])