From a000f35a30bf89c9017a28936f6d0751924dfc0b Mon Sep 17 00:00:00 2001 From: Gertjan Bisschop Date: Thu, 8 Dec 2022 22:23:11 +0000 Subject: [PATCH] first pass --- algorithms.py | 37 ++++++++++++++--- lib/msprime.c | 71 ++++++++++++++++++++++++++----- lib/tests/test_ancestry.c | 45 ++++++++++++++++++++ lib/tests/test_pedigrees.c | 85 +++++++++++++++++++++++++++----------- tests/test_algorithms.py | 54 ++++++++++++++++++++++++ tests/test_ancestry.py | 28 +++++++++++++ tests/test_pedigree.py | 55 ++++++++++++++++++++++++ 7 files changed, 335 insertions(+), 40 deletions(-) diff --git a/algorithms.py b/algorithms.py index 14a81bd01..57efc4e2b 100644 --- a/algorithms.py +++ b/algorithms.py @@ -1140,6 +1140,10 @@ def dtwf_generation(self): for h in H: segments_to_merge = len(h) if segments_to_merge == 1: + if self.store_unary: + # FIXME: should be flagged as a PASS_TRHOUGH_EVENT + new_node_id = self.store_node(pop_idx) + self.store_arg_edges(h[0][1], new_node_id) h = [] elif segments_to_merge >= 2: for _, individual in h: @@ -1247,8 +1251,8 @@ def dtwf_climb_pedigree(self): for ploid in range(ind.ploidy): self.process_pedigree_common_ancestors(ind, ploid) - def store_arg_edges(self, segment, u=None): - if u is None: + def store_arg_edges(self, segment, u=-1): + if u == -1: u = len(self.tables.nodes) - 1 # Store edges pointing to current node to the left x = segment @@ -1724,6 +1728,7 @@ def merge_ancestors(self, H, pop_id, label, new_node_id=-1): pop = self.P[pop_id] defrag_required = False coalescence = False + store_edge = False alpha = None z = None merged_head = None @@ -1802,9 +1807,20 @@ def merge_ancestors(self, H, pop_id, label, new_node_id=-1): self.set_segment_mass(alpha) z = alpha if self.full_arg: + store_edge = True if not coalescence: - self.store_node(pop_id, flags=msprime.NODE_IS_CA_EVENT) - self.store_arg_edges(z) + new_node_id = self.store_node(pop_id, flags=msprime.NODE_IS_CA_EVENT) + + elif self.store_unary: + if new_node_id != -1: + store_edge = True + elif self.model == "dtwf": + # PASS_THROUGH_EVENT has been dealt with in merge_n_ancestors + new_node_id = self.store_node(pop_id, flags=msprime.NODE_IS_CA_EVENT) + store_edge = True + + if store_edge: + self.store_arg_edges(z, new_node_id) if defrag_required: self.defrag_segment_chain(z) if coalescence: @@ -1853,6 +1869,8 @@ def merge_two_ancestors(self, population_index, label, x, y): z = None coalescence = False defrag_required = False + store_edge = False + u = -1 while x is not None or y is not None: alpha = None if x is None or y is None: @@ -1939,10 +1957,17 @@ def merge_two_ancestors(self, population_index, label, x, y): z = alpha if self.full_arg: + store_edge = True if not coalescence: u = self.store_node(population_index, flags=msprime.NODE_IS_CA_EVENT) - self.store_arg_edges(z, u) - elif self.store_unary and coalescence: + elif self.store_unary: + if u != -1: + store_edge = True + elif self.model == "dtwf": + u = self.store_node(population_index) + store_edge = True + + if store_edge: self.store_arg_edges(z, u) if defrag_required: self.defrag_segment_chain(z) diff --git a/lib/msprime.c b/lib/msprime.c index 4734ca56a..aa4d9b8ed 100644 --- a/lib/msprime.c +++ b/lib/msprime.c @@ -2784,6 +2784,7 @@ msp_merge_two_ancestors(msp_t *self, population_id_t population_id, label_id_t l int ret = 0; bool coalescence = false; bool defrag_required = false; + bool store_edge = false; tsk_id_t v; double l, r, l_min, r_max; avl_node_t *node; @@ -2938,6 +2939,7 @@ msp_merge_two_ancestors(msp_t *self, population_id_t population_id, label_id_t l } } if (self->store_full_arg) { + store_edge = true; if (!coalescence) { ret = msp_store_node( self, MSP_NODE_IS_CA_EVENT, self->time, population_id, TSK_NULL); @@ -2945,19 +2947,31 @@ msp_merge_two_ancestors(msp_t *self, population_id_t population_id, label_id_t l goto out; } } - // is specified new_node_id impossible when recording full_arg? - ret = msp_store_arg_edges(self, z, TSK_NULL); - if (ret != 0) { - goto out; - } + } else { - if (self->store_unary && coalescence) { - ret = msp_store_arg_edges(self, z, new_node_id); - if (ret < 0) { - goto out; + if (self->store_unary) { + if (new_node_id != TSK_NULL) { + store_edge = true; + } else { + if (self->model.type == MSP_MODEL_DTWF) { + new_node_id = msp_store_node( + self, MSP_NODE_IS_CA_EVENT, self->time, population_id, TSK_NULL); + ret = (int) new_node_id; + if (ret < 0) { + goto out; + } + store_edge = true; + } } } } + + if (store_edge) { + ret = msp_store_arg_edges(self, z, new_node_id); + if (ret != 0) { + goto out; + } + } if (defrag_required) { ret = msp_defrag_segment_chain(self, z); if (ret != 0) { @@ -3020,6 +3034,7 @@ msp_merge_ancestors(msp_t *self, avl_tree_t *Q, population_id_t population_id, int ret = MSP_ERR_GENERIC; bool coalescence = false; bool defrag_required = false; + bool store_edge = false; uint32_t j, h; double l, r, r_max, next_l, l_min; avl_node_t *node; @@ -3180,6 +3195,7 @@ msp_merge_ancestors(msp_t *self, avl_tree_t *Q, population_id_t population_id, } } if (self->store_full_arg) { + store_edge = true; if (!coalescence) { ret = msp_store_node( self, MSP_NODE_IS_CA_EVENT, self->time, population_id, individual); @@ -3187,7 +3203,27 @@ msp_merge_ancestors(msp_t *self, avl_tree_t *Q, population_id_t population_id, goto out; } } - ret = msp_store_arg_edges(self, z, TSK_NULL); + } else { + if (self->store_unary) { + if (new_node_id != TSK_NULL) { + store_edge = true; + } else { + if (self->model.type == MSP_MODEL_DTWF) { + // PASS_THROUGH_EVENT have been dealth with in merge_n_ancestors + new_node_id = msp_store_node(self, MSP_NODE_IS_CA_EVENT, self->time, + population_id, individual); + ret = (int) new_node_id; + if (ret < 0) { + goto out; + } + store_edge = true; + } + } + } + } + + if (store_edge) { + ret = msp_store_arg_edges(self, z, new_node_id); if (ret != 0) { goto out; } @@ -3246,6 +3282,21 @@ msp_merge_n_ancestors(msp_t *self, avl_tree_t *Q, population_id_t population_id, if (num_common_ancestors == 1) { merged_head = msp_priority_queue_pop(self, Q); + if (self->store_unary) { + if (self->model.type == MSP_MODEL_DTWF) { + new_node_id = msp_store_node( + /*FIXME: not a CA_EVENT but PASS_THROUGH_EVENT*/ + self, MSP_NODE_IS_CA_EVENT, self->time, population_id, TSK_NULL); + ret = (int) new_node_id; + if (ret < 0) { + goto out; + } + } + ret = msp_store_arg_edges(self, merged_head, new_node_id); + if (ret != 0) { + goto out; + } + } } else if (num_common_ancestors >= 2) { msp_remove_individuals_from_population(self, Q); if (num_common_ancestors == 2) { diff --git a/lib/tests/test_ancestry.c b/lib/tests/test_ancestry.c index aee5f75c7..b372c3078 100644 --- a/lib/tests/test_ancestry.c +++ b/lib/tests/test_ancestry.c @@ -824,6 +824,49 @@ test_dtwf_multi_locus_simulation(void) tsk_table_collection_free(&tables); } +static void +test_dtwf_multi_locus_simulation_unary(void) +{ + int ret; + const char *model_name; + uint32_t n = 100; + msp_t msp; + tsk_size_t num_edges, num_edges_simple; + gsl_rng *rng = safe_rng_alloc(); + tsk_table_collection_t tables; + + ret = build_sim(&msp, &tables, rng, 100, 1, NULL, n); + CU_ASSERT_EQUAL(ret, 0); + ret = msp_initialise(&msp); + CU_ASSERT_EQUAL(ret, 0); + ret = msp_set_recombination_rate(&msp, 1); + CU_ASSERT_EQUAL_FATAL(ret, 0); + ret = msp_set_simulation_model_dtwf(&msp); + CU_ASSERT_EQUAL(ret, 0); + ret = msp_set_population_configuration(&msp, 0, n, 0, true); + CU_ASSERT_EQUAL(ret, 0); + ret = msp_set_store_unary(&msp, true); + CU_ASSERT_EQUAL(ret, 0); + model_name = msp_get_model_name(&msp); + CU_ASSERT_STRING_EQUAL(model_name, "dtwf"); + + ret = msp_run(&msp, DBL_MAX, UINT32_MAX); + CU_ASSERT_EQUAL(ret, 0); + msp_verify(&msp, 0); + + // verify whether there is at least one unary node + num_edges = tables.edges.num_rows; + ret = tsk_table_collection_simplify(&tables, NULL, 0, 0, NULL); + CU_ASSERT_EQUAL_FATAL(ret, 0); + num_edges_simple = tables.edges.num_rows; + CU_ASSERT_TRUE(num_edges_simple < num_edges); + + ret = msp_free(&msp); + CU_ASSERT_EQUAL(ret, 0); + gsl_rng_free(rng); + tsk_table_collection_free(&tables); +} + static void test_dtwf_deterministic(void) { @@ -3573,6 +3616,8 @@ main(int argc, char **argv) { "test_dtwf_single_locus_simulation", test_dtwf_single_locus_simulation }, { "test_dtwf_multi_locus_simulation", test_dtwf_multi_locus_simulation }, + { "test_dtwf_multi_locus_simulation_unary", + test_dtwf_multi_locus_simulation_unary }, { "test_dtwf_deterministic", test_dtwf_deterministic }, { "test_dtwf_simultaneous_historical_samples", test_dtwf_simultaneous_historical_samples }, diff --git a/lib/tests/test_pedigrees.c b/lib/tests/test_pedigrees.c index d188c77e3..ba9949d5e 100644 --- a/lib/tests/test_pedigrees.c +++ b/lib/tests/test_pedigrees.c @@ -103,10 +103,32 @@ verify_complete_pedigree_simulation( gsl_rng_free(rng); } +static void +verify_store_unary(tsk_table_collection_t *tables) +{ + tsk_size_t num_edges, num_edges_simple; + int ret; + tsk_treeseq_t ts; + + /* verify whether there is at least one unary node */ + num_edges = tables->edges.num_rows; + ret = tsk_table_collection_simplify( + tables, NULL, 0, TSK_SIMPLIFY_KEEP_INPUT_ROOTS, NULL); + CU_ASSERT_EQUAL_FATAL(ret, 0); + num_edges_simple = tables->edges.num_rows; + CU_ASSERT_TRUE(num_edges_simple < num_edges); + /* Is this a valid tree sequence? */ + ret = tsk_treeseq_init(&ts, tables, TSK_TS_INIT_BUILD_INDEXES); + CU_ASSERT_EQUAL_FATAL(ret, 0); + /* tsk_table_collection_print_state(tables, stdout); */ + + tsk_treeseq_free(&ts); +} + static void verify_pedigree(double recombination_rate, unsigned long seed, tsk_size_t num_individuals, tsk_id_t *parents, double *time, tsk_flags_t *is_sample, - tsk_id_t *population) + tsk_id_t *population, bool store_unary) { int ret; int ploidy = 2; @@ -125,6 +147,8 @@ verify_pedigree(double recombination_rate, unsigned long seed, num_nodes = msp.tables->nodes.num_rows; ret = msp_set_recombination_rate(&msp, recombination_rate); CU_ASSERT_EQUAL(ret, 0); + ret = msp_set_store_unary(&msp, store_unary); + CU_ASSERT_EQUAL_FATAL(ret, 0); ret = msp_initialise(&msp); CU_ASSERT_EQUAL_FATAL(ret, 0); @@ -163,6 +187,10 @@ verify_pedigree(double recombination_rate, unsigned long seed, } CU_ASSERT_EQUAL_FATAL(ret, 0); + if (store_unary) { + verify_store_unary(&tables); + } + tsk_tree_free(&tree); tsk_treeseq_free(&ts); gsl_rng_free(rng); @@ -177,7 +205,7 @@ test_trio(void) tsk_id_t parents[] = { -1, -1, -1, -1, 0, 1 }; double time[] = { 1, 1, 0 }; - verify_pedigree(0, 1, 3, parents, time, NULL, NULL); + verify_pedigree(0, 1, 3, parents, time, NULL, NULL, false); } static void @@ -186,8 +214,8 @@ test_three_generations(void) tsk_id_t parents[] = { -1, -1, -1, -1, -1, -1, -1, -1, 0, 1, 2, 3, 4, 5 }; double time[] = { 2, 2, 2, 2, 1, 1, 0 }; - verify_pedigree(0, 1234, 7, parents, time, NULL, NULL); - verify_pedigree(0.1, 1234, 7, parents, time, NULL, NULL); + verify_pedigree(0, 1234, 7, parents, time, NULL, NULL, false); + verify_pedigree(0.1, 1234, 7, parents, time, NULL, NULL, false); } static void @@ -198,8 +226,8 @@ test_sibs(void) int seed; for (seed = 1; seed < 10; seed++) { - verify_pedigree(1, seed, 4, parents, time, NULL, NULL); - verify_pedigree(0, seed, 4, parents, time, NULL, NULL); + verify_pedigree(1, seed, 4, parents, time, NULL, NULL, false); + verify_pedigree(0, seed, 4, parents, time, NULL, NULL, false); } } @@ -212,8 +240,8 @@ test_ancient_sample(void) int seed; for (seed = 1; seed < 10; seed++) { - verify_pedigree(0.1, seed, 4, parents, time, is_sample, NULL); - verify_pedigree(0, seed, 4, parents, time, is_sample, NULL); + verify_pedigree(0.1, seed, 4, parents, time, is_sample, NULL, false); + verify_pedigree(0, seed, 4, parents, time, is_sample, NULL, false); } } @@ -224,8 +252,10 @@ test_large_family(void) size_t n = sizeof(large_family_time) / sizeof(*large_family_time); for (seed = 1; seed < 10; seed++) { - verify_pedigree(1, seed, n, large_family_parents, large_family_time, NULL, NULL); - verify_pedigree(0, seed, n, large_family_parents, large_family_time, NULL, NULL); + verify_pedigree( + 1, seed, n, large_family_parents, large_family_time, NULL, NULL, false); + verify_pedigree( + 0, seed, n, large_family_parents, large_family_time, NULL, NULL, false); } } @@ -235,8 +265,8 @@ test_unrelated_n3(void) tsk_id_t parents[] = { -1, -1, -1, -1, -1, -1 }; double time[] = { 0.0, 0.0, 0.0 }; - verify_pedigree(0, 1, 3, parents, time, NULL, NULL); - verify_pedigree(0.1, 1, 3, parents, time, NULL, NULL); + verify_pedigree(0, 1, 3, parents, time, NULL, NULL, false); + verify_pedigree(0.1, 1, 3, parents, time, NULL, NULL, false); } static void @@ -244,8 +274,11 @@ test_very_deep_n2(void) { size_t n = sizeof(very_deep_n2_time) / sizeof(*very_deep_n2_time); - verify_pedigree(0, 1, n, very_deep_n2_parents, very_deep_n2_time, NULL, NULL); - verify_pedigree(0.1, 1, n, very_deep_n2_parents, very_deep_n2_time, NULL, NULL); + verify_pedigree(0, 1, n, very_deep_n2_parents, very_deep_n2_time, NULL, NULL, false); + verify_pedigree( + 0.1, 1, n, very_deep_n2_parents, very_deep_n2_time, NULL, NULL, false); + verify_pedigree( + 0.1, 1, n, very_deep_n2_parents, very_deep_n2_time, NULL, NULL, true); } static void @@ -253,8 +286,10 @@ test_two_pedigrees(void) { size_t n = sizeof(two_pedigrees_time) / sizeof(*two_pedigrees_time); - verify_pedigree(0, 1, n, two_pedigrees_parents, two_pedigrees_time, NULL, NULL); - verify_pedigree(0.1, 1, n, two_pedigrees_parents, two_pedigrees_time, NULL, NULL); + verify_pedigree( + 0, 1, n, two_pedigrees_parents, two_pedigrees_time, NULL, NULL, false); + verify_pedigree( + 0.1, 1, n, two_pedigrees_parents, two_pedigrees_time, NULL, NULL, false); } static void @@ -264,8 +299,9 @@ test_deep_n2(void) = { -1, -1, -1, -1, 0, 1, 0, 1, 2, 3, 2, 3, 4, 5, 4, 5, 6, 7, 6, 7, 8, 9, 8, 9 }; double time[] = { 5.0, 5.0, 4.0, 4.0, 3.0, 3.0, 2.0, 2.0, 1.0, 1.0, 0.0, 0.0 }; - verify_pedigree(0, 1, 12, parents, time, NULL, NULL); - verify_pedigree(0.1, 1, 12, parents, time, NULL, NULL); + verify_pedigree(0, 1, 12, parents, time, NULL, NULL, false); + verify_pedigree(0.1, 1, 12, parents, time, NULL, NULL, false); + verify_pedigree(0.1, 1, 12, parents, time, NULL, NULL, true); } static void @@ -273,8 +309,9 @@ test_deep_n10(void) { size_t n = sizeof(deep_n10_time) / sizeof(*deep_n10_time); - verify_pedigree(0, 1, n, deep_n10_parents, deep_n10_time, NULL, NULL); - verify_pedigree(0.1, 1, n, deep_n10_parents, deep_n10_time, NULL, NULL); + verify_pedigree(0, 1, n, deep_n10_parents, deep_n10_time, NULL, NULL, false); + verify_pedigree(0.1, 1, n, deep_n10_parents, deep_n10_time, NULL, NULL, false); + verify_pedigree(1, 1, n, deep_n10_parents, deep_n10_time, NULL, NULL, true); } static void @@ -317,8 +354,8 @@ test_shallow_n100(void) size_t n = sizeof(time) / sizeof(*time); for (int seed = 1; seed < 5; seed++) { - verify_pedigree(0, seed, n, parents, time, NULL, NULL); - verify_pedigree(0.1, seed, n, parents, time, NULL, NULL); + verify_pedigree(0, seed, n, parents, time, NULL, NULL, false); + verify_pedigree(0.1, seed, n, parents, time, NULL, NULL, false); } } @@ -860,7 +897,7 @@ test_trio_same_pop(void) double time[] = { 1, 1, 0 }; tsk_id_t population[] = { 2, 2, 2 }; - verify_pedigree(0, 1, 3, parents, time, NULL, population); + verify_pedigree(0, 1, 3, parents, time, NULL, population, false); } static void @@ -871,7 +908,7 @@ test_trio_child_different_pop(void) double time[] = { 1, 1, 0 }; tsk_id_t population[] = { 2, 2, 1 }; - verify_pedigree(0, 1, 3, parents, time, NULL, population); + verify_pedigree(0, 1, 3, parents, time, NULL, population, false); } int diff --git a/tests/test_algorithms.py b/tests/test_algorithms.py index 3d57cc9f5..a7bf0ddef 100644 --- a/tests/test_algorithms.py +++ b/tests/test_algorithms.py @@ -54,6 +54,43 @@ def verify_unary(ts): assert maxc >= 2 +def verify_pedigree_unary(ts): + direct_ancestors = [set() for _ in range(ts.num_nodes)] + + # collect all direct ancestors for each node + for node in ts.nodes(): + individual = ts.individual(node.individual) + for parent_id in individual.parents: + if parent_id != tskit.NULL: + parent = ts.individual(parent_id) + for parent_node in parent.nodes: + direct_ancestors[node.id].add(parent_node) + + # assert that in each marginal tree all nodes a parent from + # the constructed set + for tree in ts.trees(): + for node_id in tree.nodes(): + if tree.parent(node_id) != tskit.NULL: + direct_ancestor = direct_ancestors[node_id] + assert tree.parent(node_id) in direct_ancestor + + +def verify_dtwf_unary(ts): + # assert all nodes in each tree and subtree only differ one generation + for tree in ts.trees(): + queue = [tree.root] + while queue: + parent = queue.pop(0) + time = tree.time(parent) - 1 + num_children = 0 + for child in tree.children(parent): + assert tree.time(child) == time + queue.append(child) + num_children += 1 + if num_children == 0: + assert tree.time(parent) == 0 + + @pytest.mark.skipif(IS_WINDOWS, reason="Bintrees isn't availble on windows") class TestAlgorithms: def run_script(self, cmd): @@ -94,6 +131,13 @@ def test_dtwf_discrete(self): assert ts.num_trees > 1 assert has_discrete_genome(ts) + def test_dtwf_store_unary(self): + ts = self.run_script("10 --model=dtwf --store-unary") + assert ts.num_trees > 1 + assert not has_discrete_genome(ts) + assert ts.sequence_length == 100 + verify_dtwf_unary(ts) + def test_full_arg(self): ts = self.run_script("30 -L 200 --full-arg") assert ts.num_trees > 1 @@ -277,3 +321,13 @@ def test_one_gen_pedigree(self, num_founders): tables.dump(ts_path) ts = self.run_script(f"0 --from-ts {ts_path} -r 1 --model=fixed_pedigree") assert len(ts.dump_tables().edges) == 0 + + def test_store_unary_pedigree(self): + tables = simulate_pedigree(num_founders=4, num_generations=10) + with tempfile.TemporaryDirectory() as tmpdir: + ts_path = pathlib.Path(tmpdir) / "pedigree.trees" + tables.dump(ts_path) + ts = self.run_script( + f"0 --from-ts {ts_path} -r 1 --model=fixed_pedigree --store-unary" + ) + verify_pedigree_unary(ts) diff --git a/tests/test_ancestry.py b/tests/test_ancestry.py index 2f5abfdd2..7d4b554a1 100644 --- a/tests/test_ancestry.py +++ b/tests/test_ancestry.py @@ -408,6 +408,34 @@ def test_predefined_scenario(self): tree.next() assert tree.parent_array[left_root] == tree.root + def verify_dtwf_unary(self, ts): + # assert all nodes in each tree and subtree only differ one generation + for tree in ts.trees(): + queue = [tree.root] + while queue: + parent = queue.pop(0) + time = tree.time(parent) - 1 + num_children = 0 + for child in tree.children(parent): + assert tree.time(child) == time + queue.append(child) + num_children += 1 + if num_children == 0: + assert tree.time(parent) == 0 + + def test_dtwf_store_unary(self): + ts = msprime.sim_ancestry( + 10, + population_size=100, + model="dtwf", + ploidy=2, + random_seed=1234, + recombination_rate=1, + sequence_length=100, + record_unary=True, + ) + self.verify_dtwf_unary(ts) + class TestSimulator: """ diff --git a/tests/test_pedigree.py b/tests/test_pedigree.py index 2d2d7f96f..5a0c2f15e 100644 --- a/tests/test_pedigree.py +++ b/tests/test_pedigree.py @@ -1535,6 +1535,61 @@ def test_mismatched_node_populations(self): pedigrees.write_pedigree(tables.tree_sequence(), io.StringIO()) +class TestPedigreeUnary: + def test_pedigree_unary(self): + initial_state = simulate_pedigree( + num_founders=3, num_generations=5, sequence_length=100 + ) + ts = msprime.sim_ancestry( + recombination_rate=0.1, + model="fixed_pedigree", + record_unary=True, + initial_state=initial_state, + ) + self.verify_pedigree_unary(ts) + + def test_pedigree_unary_simple(self): + sequence_length = 100 + pb = msprime.PedigreeBuilder() + mmid = pb.add_individual(time=2) + mdid = pb.add_individual(time=2) + dmid = pb.add_individual(time=2) + ddid = pb.add_individual(time=2) + mid = pb.add_individual(time=1, parents=[mmid, mdid]) + did = pb.add_individual(time=1, parents=[dmid, ddid]) + pb.add_individual(time=0, parents=[mid, did], is_sample=True) + pedigree = pb.finalise(sequence_length) + ts = msprime.sim_ancestry( + recombination_rate=0.1, + model="fixed_pedigree", + record_unary=True, + initial_state=pedigree, + ) + self.verify_pedigree_unary(ts) + + def verify_pedigree_unary(self, ts): + direct_ancestors = [set() for _ in range(ts.num_nodes)] + + # collect all direct ancestors for each node + # assuming the initial tables are not modified by the + # pedigree simulation + for node in ts.nodes(): + individual = ts.individual(node.individual) + for parent_id in individual.parents: + if parent_id != tskit.NULL: + parent = ts.individual(parent_id) + for parent_node in parent.nodes: + direct_ancestors[node.id].add(parent_node) + + # assert that in each marginal tree all nodes a parent from + # the constructed set + for tree in ts.trees(): + for node_id in tree.nodes(): + if tree.parent(node_id) != tskit.NULL: + direct_ancestor = direct_ancestors[node_id] + assert tree.parent(node_id) in direct_ancestor + + if __name__ == "__main__": # Easy way to output a simulated pedigree for command line testing. # Run with python3 tests/test_pedigree.py NUM_FOUNDERS NUM_GENERATIONS > out.trees