Retain unary nodes with simplify()

c/CHANGELOG.rst

@@ -1,3 +1,14 @@
+---------------------
+[0.99.3] - 2019-XX-XX
+---------------------
+
+In development.
+
+**New features**
+
+- Add the ``TSK_KEEP_UNARY`` option to simplify (:user:`gtsambos`). See :issue:`1`
+  and :pr:`143`.
+
 ---------------------
 [0.99.2] - 2019-03-27
 ---------------------

c/tests/test_trees.c

@@ -872,14 +872,27 @@ test_simplest_records(void)
         "0  1   0";
     const char *edges =
         "0  1   2   0,1\n";
-    tsk_treeseq_t ts;
+    tsk_treeseq_t ts, simplified;
+    tsk_id_t sample_ids[] = {0, 1};
+    int ret;
 
     tsk_treeseq_from_text(&ts, 1, nodes, edges, NULL, NULL, NULL, NULL, NULL);
     CU_ASSERT_EQUAL(tsk_treeseq_get_num_samples(&ts), 2);
     CU_ASSERT_EQUAL(tsk_treeseq_get_sequence_length(&ts), 1.0);
     CU_ASSERT_EQUAL(tsk_treeseq_get_num_nodes(&ts), 3);
     CU_ASSERT_EQUAL(tsk_treeseq_get_num_mutations(&ts), 0);
     CU_ASSERT_EQUAL(tsk_treeseq_get_num_trees(&ts), 1);
+
+    ret = tsk_treeseq_simplify(&ts, sample_ids, 2, 0, &simplified, NULL);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    CU_ASSERT_TRUE(tsk_table_collection_equals(ts.tables, simplified.tables));
+    tsk_treeseq_free(&simplified);
+
+    ret = tsk_treeseq_simplify(&ts, sample_ids, 2, TSK_KEEP_UNARY, &simplified, NULL);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    CU_ASSERT_TRUE(tsk_table_collection_equals(ts.tables, simplified.tables));
+    tsk_treeseq_free(&simplified);
+
     tsk_treeseq_free(&ts);
 }
 
@@ -894,14 +907,27 @@ test_simplest_nonbinary_records(void)
         "0  1   0";
     const char *edges =
         "0  1   4   0,1,2,3\n";
-    tsk_treeseq_t ts;
+    tsk_treeseq_t ts, simplified;
+    tsk_id_t sample_ids[] = {0, 1, 2, 3};
+    int ret;
 
     tsk_treeseq_from_text(&ts, 1, nodes, edges, NULL, NULL, NULL, NULL, NULL);
     CU_ASSERT_EQUAL(tsk_treeseq_get_num_samples(&ts), 4);
     CU_ASSERT_EQUAL(tsk_treeseq_get_sequence_length(&ts), 1.0);
     CU_ASSERT_EQUAL(tsk_treeseq_get_num_nodes(&ts), 5);
     CU_ASSERT_EQUAL(tsk_treeseq_get_num_mutations(&ts), 0);
     CU_ASSERT_EQUAL(tsk_treeseq_get_num_trees(&ts), 1);
+
+    ret = tsk_treeseq_simplify(&ts, sample_ids, 4, 0, &simplified, NULL);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    CU_ASSERT_TRUE(tsk_table_collection_equals(ts.tables, simplified.tables));
+    tsk_treeseq_free(&simplified);
+
+    ret = tsk_treeseq_simplify(&ts, sample_ids, 4, TSK_KEEP_UNARY, &simplified, NULL);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    CU_ASSERT_TRUE(tsk_table_collection_equals(ts.tables, simplified.tables));
+    tsk_treeseq_free(&simplified);
+
     tsk_treeseq_free(&ts);
 }
 
@@ -935,11 +961,17 @@ test_simplest_unary_records(void)
     CU_ASSERT_EQUAL(tsk_treeseq_get_num_samples(&simplified), 2);
     CU_ASSERT_EQUAL(tsk_treeseq_get_sequence_length(&simplified), 1.0);
     CU_ASSERT_EQUAL(tsk_treeseq_get_num_nodes(&simplified), 3);
+    CU_ASSERT_EQUAL(tsk_treeseq_get_num_edges(&simplified), 2);
     CU_ASSERT_EQUAL(tsk_treeseq_get_num_mutations(&simplified), 0);
     CU_ASSERT_EQUAL(tsk_treeseq_get_num_trees(&simplified), 1);
+    tsk_treeseq_free(&simplified);
 
-    tsk_treeseq_free(&ts);
+    ret = tsk_treeseq_simplify(&ts, sample_ids, 2, TSK_KEEP_UNARY, &simplified, NULL);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    CU_ASSERT_TRUE(tsk_table_collection_equals(ts.tables, simplified.tables));
     tsk_treeseq_free(&simplified);
+
+    tsk_treeseq_free(&ts);
 }
 
 static void
@@ -1072,8 +1104,13 @@ test_simplest_degenerate_multiple_root_records(void)
     CU_ASSERT_EQUAL(tsk_treeseq_get_num_samples(&simplified), 2);
     CU_ASSERT_EQUAL(tsk_treeseq_get_sequence_length(&simplified), 1.0);
     CU_ASSERT_EQUAL(tsk_treeseq_get_num_nodes(&simplified), 2);
+    tsk_treeseq_free(&simplified);
 
+    ret = tsk_treeseq_simplify(&ts, sample_ids, 2, TSK_KEEP_UNARY, &simplified, NULL);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    CU_ASSERT_TRUE(tsk_table_collection_equals(ts.tables, simplified.tables));
     tsk_treeseq_free(&simplified);
+
     tsk_treeseq_free(&ts);
     tsk_tree_free(&t);
 }
@@ -1370,7 +1407,8 @@ test_simplest_holey_tree_sequence(void)
     char *haplotype;
     unsigned int j;
     int ret;
-    tsk_treeseq_t ts;
+    tsk_treeseq_t ts, simplified;
+    tsk_id_t sample_ids[] = {0, 1};
     tsk_hapgen_t hapgen;
 
     tsk_treeseq_from_text(&ts, 3, nodes_txt, edges_txt, NULL, sites_txt,
@@ -1390,8 +1428,18 @@ test_simplest_holey_tree_sequence(void)
         CU_ASSERT_EQUAL(ret, 0);
         CU_ASSERT_STRING_EQUAL(haplotype, haplotypes[j]);
     }
-
     tsk_hapgen_free(&hapgen);
+
+    ret = tsk_treeseq_simplify(&ts, sample_ids, 2, 0, &simplified, NULL);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    CU_ASSERT_TRUE(tsk_table_collection_equals(ts.tables, simplified.tables));
+    tsk_treeseq_free(&simplified);
+
+    ret = tsk_treeseq_simplify(&ts, sample_ids, 2, TSK_KEEP_UNARY, &simplified, NULL);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    CU_ASSERT_TRUE(tsk_table_collection_equals(ts.tables, simplified.tables));
+    tsk_treeseq_free(&simplified);
+
     tsk_treeseq_free(&ts);
 }
 
@@ -1462,14 +1510,15 @@ test_simplest_initial_gap_tree_sequence(void)
     char *haplotype;
     unsigned int j;
     int ret;
-    tsk_treeseq_t ts;
+    tsk_treeseq_t ts, simplified;
     tsk_hapgen_t hapgen;
     const tsk_id_t z = TSK_NULL;
     tsk_id_t parents[] = {
         z, z, z,
         2, 2, z,
     };
-    uint32_t num_trees = 2;
+    tsk_size_t num_trees = 2;
+    tsk_id_t sample_ids[] = {0, 1};
 
     tsk_treeseq_from_text(&ts, 3, nodes, edges, NULL, sites, mutations, NULL, NULL);
     CU_ASSERT_EQUAL(tsk_treeseq_get_num_samples(&ts), 2);
@@ -1490,6 +1539,17 @@ test_simplest_initial_gap_tree_sequence(void)
         CU_ASSERT_STRING_EQUAL(haplotype, haplotypes[j]);
     }
     tsk_hapgen_free(&hapgen);
+
+    ret = tsk_treeseq_simplify(&ts, sample_ids, 2, 0, &simplified, NULL);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    CU_ASSERT_TRUE(tsk_table_collection_equals(ts.tables, simplified.tables));
+    tsk_treeseq_free(&simplified);
+
+    ret = tsk_treeseq_simplify(&ts, sample_ids, 2, TSK_KEEP_UNARY, &simplified, NULL);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    CU_ASSERT_TRUE(tsk_table_collection_equals(ts.tables, simplified.tables));
+    tsk_treeseq_free(&simplified);
+
     tsk_treeseq_free(&ts);
 }
 

c/tskit/tables.c

@@ -4006,6 +4006,8 @@ simplifier_print_state(simplifier_t *self, FILE *out)
             !!(self->options & TSK_FILTER_SITES));
     fprintf(out, "\treduce_to_site_topology  : %d\n",
             !!(self->options & TSK_REDUCE_TO_SITE_TOPOLOGY));
+    fprintf(out, "\tkeep_unary  : %d\n",
+            !!(self->options & TSK_KEEP_UNARY));
 
     fprintf(out, "===\nInput tables\n==\n");
     tsk_table_collection_print_state(&self->input_tables, out);
@@ -4531,6 +4533,7 @@ simplifier_merge_ancestors(simplifier_t *self, tsk_id_t input_id)
     tsk_id_t ancestry_node;
     tsk_id_t output_id = self->node_id_map[input_id];
     bool is_sample = output_id != TSK_NULL;
+    bool keep_unary = !!(self->options & TSK_KEEP_UNARY);
 
     if (is_sample) {
         /* Free up the existing ancestry mapping. */
@@ -4558,6 +4561,14 @@ simplifier_merge_ancestors(simplifier_t *self, tsk_id_t input_id)
                     goto out;
                 }
                 ancestry_node = output_id;
+            } else if (keep_unary) {
+                if (output_id == TSK_NULL) {
+                    output_id = simplifier_record_node(self, input_id, false);
+                }
+                ret = simplifier_record_edge(self, left, right, ancestry_node);
+                if (ret != 0) {
+                    goto out;
+                }
             }
         } else {
             if (output_id == TSK_NULL) {
@@ -4583,6 +4594,9 @@ simplifier_merge_ancestors(simplifier_t *self, tsk_id_t input_id)
                 goto out;
             }
         }
+        if (keep_unary) {
+            ancestry_node = output_id;
+        }
         ret = simplifier_add_ancestry(self, input_id, left, right, ancestry_node);
         if (ret != 0) {
             goto out;

c/tskit/tables.h

@@ -586,6 +586,7 @@ typedef struct {
 #define TSK_FILTER_POPULATIONS           (1 << 1)
 #define TSK_FILTER_INDIVIDUALS           (1 << 2)
 #define TSK_REDUCE_TO_SITE_TOPOLOGY      (1 << 3)
+#define TSK_KEEP_UNARY                   (1 << 4)
 
 /* Flags for check_integrity */
 #define TSK_CHECK_OFFSETS                (1 << 0)
@@ -2134,6 +2135,10 @@ TSK_REDUCE_TO_SITE_TOPOLOGY
     For a given site, the topology of the tree containing that site will be 
     identical (up to node ID remapping) to the topology of the corresponding tree 
     in the input.
+TSK_KEEP_UNARY
+    By default simplify removes unary nodes (i.e., nodes with exactly one child) 
+    along the path from samples to root. If this option is specified such unary
+    nodes will be preserved in the output.
 
 .. note:: Migrations are currently not supported by simplify, and an error will
     be raised if we attempt call simplify on a table collection with greater

python/CHANGELOG.rst

@@ -1,3 +1,14 @@
+--------------------
+[0.1.6] - 2019-XX-XX
+--------------------
+
+In development.
+
+**New features**
+
+- Add the ``keep_unary`` option to simplify (:user:`gtsambos`). See :issue:`1`
+  and :pr:`143`.
+
 --------------------
 [0.1.5] - 2019-03-27
 --------------------

python/_tskitmodule.c

@@ -5410,14 +5410,15 @@ TableCollection_simplify(TableCollection *self, PyObject *args, PyObject *kwds)
     int filter_sites = true;
     int filter_individuals = false;
     int filter_populations = false;
+    int keep_unary = false;
     int reduce_to_site_topology = false;
     static char *kwlist[] = {
         "samples", "filter_sites", "filter_populations", "filter_individuals",
-        "reduce_to_site_topology", NULL};
+        "reduce_to_site_topology", "keep_unary", NULL};
 
-    if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|iiii", kwlist,
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|iiiii", kwlist,
             &samples, &filter_sites, &filter_populations, &filter_individuals,
-            &reduce_to_site_topology)) {
+            &reduce_to_site_topology, &keep_unary)) {
         goto out;
     }
     samples_array = (PyArrayObject *) PyArray_FROMANY(samples, NPY_INT32, 1, 1,
@@ -5439,6 +5440,9 @@ TableCollection_simplify(TableCollection *self, PyObject *args, PyObject *kwds)
     if (reduce_to_site_topology) {
         options |= TSK_REDUCE_TO_SITE_TOPOLOGY;
     }
+    if (keep_unary) {
+        options |= TSK_KEEP_UNARY;
+    }
 
     /* Allocate a new array to hold the node map. */
     dims = self->tables->nodes.num_rows;

python/tests/simplify.py

@@ -98,14 +98,15 @@ class Simplifier(object):
     """
     def __init__(
             self, ts, sample, reduce_to_site_topology=False, filter_sites=True,
-            filter_populations=True, filter_individuals=True):
+            filter_populations=True, filter_individuals=True, keep_unary=False):
         self.ts = ts
         self.n = len(sample)
         self.reduce_to_site_topology = reduce_to_site_topology
         self.sequence_length = ts.sequence_length
         self.filter_sites = filter_sites
         self.filter_populations = filter_populations
         self.filter_individuals = filter_individuals
+        self.keep_unary = keep_unary
         self.num_mutations = ts.num_mutations
         self.input_sites = list(ts.sites())
         self.A_head = [None for _ in range(ts.num_nodes)]
@@ -260,6 +261,10 @@ def merge_labeled_ancestors(self, S, input_id):
                 if is_sample:
                     self.record_edge(left, right, output_id, ancestry_node)
                     ancestry_node = output_id
+                elif self.keep_unary:
+                    if output_id == -1:
+                        output_id = self.record_node(input_id)
+                    self.record_edge(left, right, output_id, ancestry_node)
             else:
                 if output_id == -1:
                     output_id = self.record_node(input_id)
@@ -269,6 +274,8 @@ def merge_labeled_ancestors(self, S, input_id):
             if is_sample and left != prev_right:
                 # Fill in any gaps in the ancestry for the sample
                 self.add_ancestry(input_id, prev_right, left, output_id)
+            if self.keep_unary:
+                ancestry_node = output_id
             self.add_ancestry(input_id, left, right, ancestry_node)
             prev_right = right
 
@@ -446,7 +453,22 @@ def check_state(self):
     # Simple CLI for running simplifier above.
     ts = tskit.load(sys.argv[1])
     samples = list(map(int, sys.argv[2:]))
-    s = Simplifier(ts, samples)
+
+    # When keep_unary = True
+    print('When keep_unary = True:')
+    s = Simplifier(ts, samples, keep_unary=True)
+    # s.print_state()
+    tss, _ = s.simplify()
+    tables = tss.dump_tables()
+    print("Output:")
+    print(tables.nodes)
+    print(tables.edges)
+    print(tables.sites)
+    print(tables.mutations)
+
+    # When keep_unary = False
+    print('\nWhen keep_unary = False:')
+    s = Simplifier(ts, samples, keep_unary=False)
     # s.print_state()
     tss, _ = s.simplify()
     tables = tss.dump_tables()