Add keep_input_roots to simplify.

Also updates the internal algorithm to remove ancestry that has been
processed.

Author: jeromekelleher

appveyor.yml

@@ -29,6 +29,7 @@ build_script:
   - cmd: python -m pip install newick
   - cmd: python -m pip install python_jsonschema_objects
   - cmd: python -m pip install xmlunittest
+  - cmd: python -m pip install portion
   - cmd: python -m nose -vs --processes=%NUMBER_OF_PROCESSORS% --process-timeout=5000
 
 after_test:

c/tests/test_trees.c

@@ -3640,6 +3640,7 @@ test_single_tree_iter_depths(void)
     tsk_tree_free(&tree);
     tsk_treeseq_free(&ts);
 }
+
 static void
 test_single_tree_simplify(void)
 {
@@ -3697,6 +3698,55 @@ test_single_tree_simplify(void)
     tsk_table_collection_free(&tables);
 }
 
+static void
+test_single_tree_simplify_debug(void)
+{
+    tsk_treeseq_t ts, simplified;
+    tsk_id_t samples[] = { 0, 1 };
+    int ret;
+    FILE *save_stdout = stdout;
+    FILE *tmp = fopen(_tmp_file_name, "w");
+
+    CU_ASSERT_FATAL(tmp != NULL);
+    tsk_treeseq_from_text(&ts, 1, single_tree_ex_nodes, single_tree_ex_edges, NULL,
+        single_tree_ex_sites, single_tree_ex_mutations, NULL, NULL, 0);
+
+    stdout = tmp;
+    ret = tsk_treeseq_simplify(&ts, samples, 2, TSK_DEBUG, &simplified, NULL);
+    stdout = save_stdout;
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    CU_ASSERT_TRUE(ftell(tmp) > 0);
+
+    fclose(tmp);
+    tsk_treeseq_free(&ts);
+    tsk_treeseq_free(&simplified);
+}
+
+static void
+test_single_tree_simplify_keep_input_roots(void)
+{
+    tsk_treeseq_t ts;
+    tsk_table_collection_t tables;
+    tsk_id_t samples[] = { 0, 1 };
+    int ret;
+
+    tsk_treeseq_from_text(&ts, 1, single_tree_ex_nodes, single_tree_ex_edges, NULL,
+        single_tree_ex_sites, single_tree_ex_mutations, NULL, NULL, 0);
+    verify_simplify(&ts);
+    ret = tsk_treeseq_copy_tables(&ts, &tables, 0);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+
+    ret = tsk_table_collection_simplify(&tables, samples, 2, TSK_KEEP_INPUT_ROOTS, NULL);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    CU_ASSERT_EQUAL(tables.nodes.num_rows, 4);
+    CU_ASSERT_EQUAL(tables.edges.num_rows, 3);
+    CU_ASSERT_EQUAL(tables.sites.num_rows, 3);
+    CU_ASSERT_EQUAL(tables.mutations.num_rows, 4);
+
+    tsk_treeseq_free(&ts);
+    tsk_table_collection_free(&tables);
+}
+
 static void
 test_single_tree_simplify_no_sample_nodes(void)
 {
@@ -4283,6 +4333,60 @@ test_nonbinary_multi_tree(void)
     tsk_treeseq_free(&ts);
 }
 
+static void
+test_simplify_keep_input_roots_multi_tree(void)
+{
+
+    /*
+    0.25┊     8   ┊         ┊         ┊
+        ┊   ┏━┻━┓ ┊         ┊         ┊
+    0.20┊   ┃   ┃ ┊         ┊   7     ┊
+        ┊   ┃   ┃ ┊         ┊ ┏━┻━┓   ┊
+    0.17┊   6   ┃ ┊   6     ┊ ┃   ┃   ┊
+        ┊ ┏━┻┓  ┃ ┊ ┏━┻━┓   ┊ ┃   ┃   ┊
+    0.09┊ ┃  5  ┃ ┊ ┃   5   ┊ ┃   5   ┊
+        ┊ ┃ ┏┻┓ ┃ ┊ ┃ ┏━┻┓  ┊ ┃ ┏━┻┓  ┊
+    0.07┊ ┃ ┃ ┃ ┃ ┊ ┃ ┃  4  ┊ ┃ ┃  4  ┊
+        ┊ ┃ ┃ ┃ ┃ ┊ ┃ ┃ ┏┻┓ ┊ ┃ ┃ ┏┻┓ ┊
+    0.00┊ 0 1 3 2 ┊ 0 1 2 3 ┊ 0 1 2 3 ┊
+      0.00      2.00      7.00      10.00
+
+    Simplifies to
+
+    0.25┊  4  ┊     ┊     ┊
+        ┊  ┃  ┊     ┊     ┊
+    0.20┊  ┃  ┊     ┊  3  ┊
+        ┊  ┃  ┊     ┊ ┏┻┓ ┊
+    0.17┊  2  ┊  2  ┊ ┃ ┃ ┊
+        ┊ ┏┻┓ ┊ ┏┻┓ ┊ ┃ ┃ ┊
+    0.00┊ 0 1 ┊ 0 1 ┊ 0 1 ┊
+      0.00  2.00  7.00  10.00
+
+    */
+    int ret = 0;
+    // clang-format off
+    tsk_id_t parents[] = {
+        2, 2, 4, -1, -1,
+        2, 2, -1, -1, -1,
+        3, 3, -1, -1, -1,
+    };
+    // clang-format on
+    uint32_t num_trees = 3;
+
+    tsk_id_t samples[] = { 0, 3 };
+    tsk_treeseq_t ts, simplified;
+
+    tsk_treeseq_from_text(&ts, 10, paper_ex_nodes, paper_ex_edges, NULL, paper_ex_sites,
+        paper_ex_mutations, paper_ex_individuals, NULL, 0);
+    tsk_treeseq_dump(&ts, "tmp.trees", 0);
+    ret = tsk_treeseq_simplify(&ts, samples, 2, TSK_KEEP_INPUT_ROOTS, &simplified, NULL);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    verify_trees(&simplified, num_trees, parents);
+
+    tsk_treeseq_free(&ts);
+    tsk_treeseq_free(&simplified);
+}
+
 static void
 test_left_to_right_multi_tree(void)
 {
@@ -5839,6 +5943,9 @@ main(int argc, char **argv)
         { "test_single_tree_iter_times", test_single_tree_iter_times },
         { "test_single_tree_iter_depths", test_single_tree_iter_depths },
         { "test_single_tree_simplify", test_single_tree_simplify },
+        { "test_single_tree_simplify_debug", test_single_tree_simplify_debug },
+        { "test_single_tree_simplify_keep_input_roots",
+            test_single_tree_simplify_keep_input_roots },
         { "test_single_tree_simplify_no_sample_nodes",
             test_single_tree_simplify_no_sample_nodes },
         { "test_single_tree_simplify_null_samples",
@@ -5859,6 +5966,8 @@ main(int argc, char **argv)
         { "test_internal_sample_multi_tree", test_internal_sample_multi_tree },
         { "test_internal_sample_simplified_multi_tree",
             test_internal_sample_simplified_multi_tree },
+        { "test_simplify_keep_input_roots_multi_tree",
+            test_simplify_keep_input_roots_multi_tree },
         { "test_left_to_right_multi_tree", test_left_to_right_multi_tree },
         { "test_gappy_multi_tree", test_gappy_multi_tree },
         { "test_tsk_treeseq_bad_records", test_tsk_treeseq_bad_records },