Implement TSK_SIMPLIFY_NO_UPDATE_SAMPLE_FLAGS

Closes tskit-dev#2662

Author: jeromekelleher

c/tests/test_trees.c

@@ -3323,6 +3323,36 @@ test_simplest_no_node_filter(void)
     tsk_treeseq_free(&ts);
 }
 
+static void
+test_simplest_no_update_flags(void)
+{
+    const char *nodes = "0  0   0\n"
+                        "1  0   0\n"
+                        "0  1   0\n";
+    const char *edges = "0  1   2   0,1\n";
+    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, 0);
+
+    /* We have a mixture of sample and non-samples in the input tables */
+    ret = tsk_treeseq_simplify(
+        &ts, sample_ids, 2, TSK_SIMPLIFY_NO_UPDATE_SAMPLE_FLAGS, &simplified, NULL);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    CU_ASSERT_TRUE(tsk_table_collection_equals(ts.tables, simplified.tables, 0));
+    tsk_treeseq_free(&simplified);
+
+    ret = tsk_treeseq_simplify(&ts, sample_ids, 2,
+        TSK_SIMPLIFY_NO_UPDATE_SAMPLE_FLAGS | TSK_SIMPLIFY_NO_FILTER_NODES, &simplified,
+        NULL);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    CU_ASSERT_TRUE(tsk_table_collection_equals(ts.tables, simplified.tables, 0));
+    tsk_treeseq_free(&simplified);
+
+    tsk_treeseq_free(&ts);
+}
+
 static void
 test_simplest_map_mutations(void)
 {
@@ -8093,6 +8123,7 @@ main(int argc, char **argv)
         { "test_simplest_population_filter", test_simplest_population_filter },
         { "test_simplest_individual_filter", test_simplest_individual_filter },
         { "test_simplest_no_node_filter", test_simplest_no_node_filter },
+        { "test_simplest_no_update_flags", test_simplest_no_update_flags },
         { "test_simplest_map_mutations", test_simplest_map_mutations },
         { "test_simplest_nonbinary_map_mutations",
             test_simplest_nonbinary_map_mutations },

c/tskit/tables.c

@@ -8808,16 +8808,18 @@ static tsk_id_t TSK_WARN_UNUSED
 simplifier_record_node(simplifier_t *self, tsk_id_t input_id)
 {
     tsk_node_t node;
-    tsk_flags_t flags;
+    bool update_flags = !(self->options & TSK_SIMPLIFY_NO_UPDATE_SAMPLE_FLAGS);
 
     tsk_node_table_get_row_unsafe(&self->input_tables.nodes, (tsk_id_t) input_id, &node);
-    /* Zero out the sample bit */
-    flags = node.flags & (tsk_flags_t) ~TSK_NODE_IS_SAMPLE;
-    if (self->is_sample[input_id]) {
-        flags |= TSK_NODE_IS_SAMPLE;
+    if (update_flags) {
+        /* Zero out the sample bit */
+        node.flags &= (tsk_flags_t) ~TSK_NODE_IS_SAMPLE;
+        if (self->is_sample[input_id]) {
+            node.flags |= TSK_NODE_IS_SAMPLE;
+        }
     }
     self->node_id_map[input_id] = (tsk_id_t) self->tables->nodes.num_rows;
-    return tsk_node_table_add_row(&self->tables->nodes, flags, node.time,
+    return tsk_node_table_add_row(&self->tables->nodes, node.flags, node.time,
         node.population, node.individual, node.metadata, node.metadata_length);
 }
 
@@ -9108,6 +9110,7 @@ simplifier_init_nodes(simplifier_t *self, const tsk_id_t *samples)
     tsk_size_t j;
     const tsk_size_t num_nodes = self->input_tables.nodes.num_rows;
     bool filter_nodes = !(self->options & TSK_SIMPLIFY_NO_FILTER_NODES);
+    bool update_flags = !(self->options & TSK_SIMPLIFY_NO_UPDATE_SAMPLE_FLAGS);
     tsk_flags_t *node_flags = self->tables->nodes.flags;
     tsk_id_t *node_id_map = self->node_id_map;
 
@@ -9123,13 +9126,17 @@ simplifier_init_nodes(simplifier_t *self, const tsk_id_t *samples)
         }
     } else {
         tsk_bug_assert(self->tables->nodes.num_rows == num_nodes);
-        /* The node table has not been changed */
-        for (j = 0; j < num_nodes; j++) {
-            /* Reset the sample flags */
-            node_flags[j] &= (tsk_flags_t) ~TSK_NODE_IS_SAMPLE;
-            if (self->is_sample[j]) {
-                node_flags[j] |= TSK_NODE_IS_SAMPLE;
+        if (update_flags) {
+            for (j = 0; j < num_nodes; j++) {
+                /* Reset the sample flags */
+                node_flags[j] &= (tsk_flags_t) ~TSK_NODE_IS_SAMPLE;
+                if (self->is_sample[j]) {
+                    node_flags[j] |= TSK_NODE_IS_SAMPLE;
+                }
             }
+        }
+
+        for (j = 0; j < num_nodes; j++) {
             node_id_map[j] = (tsk_id_t) j;
         }
     }

c/tskit/tables.h

@@ -694,6 +694,10 @@ first.
 */
 #define TSK_SIMPLIFY_NO_FILTER_NODES (1 << 7)
 /**
+Do not update the sample status of nodes as a result of simplification.
+*/
+#define TSK_SIMPLIFY_NO_UPDATE_SAMPLE_FLAGS (1 << 8)
+/**
 Reduce the topological information in the tables to the minimum necessary to
 represent the trees that contain sites. If there are zero sites this will
 result in an zero output edges. When the number of sites is greater than zero,
@@ -3919,9 +3923,10 @@ or :c:macro:`TSK_NULL` if the node has been removed. Thus, ``node_map`` must be
 of at least ``self->nodes.num_rows`` :c:type:`tsk_id_t` values.
 
 If the `TSK_SIMPLIFY_NO_FILTER_NODES` option is specified, the node table will be
-unaltered except for changing the sample status of nodes that were samples in the
-input tables, but not in the specified list of sample IDs (if provided). The
-``node_map`` (if specified) will always be the identity mapping, such that
+unaltered except for changing the sample status of nodes (but see the
+`TSK_SIMPLIFY_NO_UPDATE_SAMPLE_FLAGS` option below) and to update references
+to other tables that may have changed as a result of filtering (see below).
+The ``node_map`` (if specified) will always be the identity mapping, such that
 ``node_map[u] == u`` for all nodes. Note also that the order of the list of
 samples is not important in this case.
 
@@ -3941,6 +3946,17 @@ sample status flag of nodes.
    may be entirely unreferenced entities in the input tables, which
    are not affected by whether we filter nodes or not.
 
+By default, the node sample flags are updated by unsetting
+:c:macro:`TSK_NODE_IS_SAMPLE` flag for all nodes and subsequently
+setting it for the nodes provided as input to this function.
+The `TSK_SIMPLIFY_NO_UPDATE_SAMPLE_FLAGS` option will prevent this
+from occuring, making it the responsibility of calling code to
+keep track of the ultimate sample status of nodes. Using
+this option in conjunction with `TSK_SIMPLIFY_NO_FILTER_NODES`
+(and without the `TSK_SIMPLIFY_FILTER_POPULATIONS` and
+`TSK_SIMPLIFY_FILTER_INDIVIDUALS`) guarantees that the node table
+will not be written to during the lifetime of this function.
+
 The table collection will always be unindexed after simplify successfully completes.
 
 .. note:: Migrations are currently not supported by simplify, and an error will
@@ -3956,6 +3972,7 @@ Options can be specified by providing one or more of the following bitwise
 - :c:macro:`TSK_SIMPLIFY_FILTER_POPULATIONS`
 - :c:macro:`TSK_SIMPLIFY_FILTER_INDIVIDUALS`
 - :c:macro:`TSK_SIMPLIFY_NO_FILTER_NODES`
+- :c:macro:`TSK_SIMPLIFY_NO_UPDATE_SAMPLE_FLAGS`
 - :c:macro:`TSK_SIMPLIFY_REDUCE_TO_SITE_TOPOLOGY`
 - :c:macro:`TSK_SIMPLIFY_KEEP_UNARY`
 - :c:macro:`TSK_SIMPLIFY_KEEP_INPUT_ROOTS`

python/_tskitmodule.c

@@ -6589,21 +6589,23 @@ TableCollection_simplify(TableCollection *self, PyObject *args, PyObject *kwds)
     int filter_individuals = false;
     int filter_populations = false;
     int filter_nodes = true;
+    int update_sample_flags = true;
     int keep_unary = false;
     int keep_unary_in_individuals = false;
     int keep_input_roots = false;
     int reduce_to_site_topology = false;
-    static char *kwlist[] = { "samples", "filter_sites", "filter_populations",
-        "filter_individuals", "filter_nodes", "reduce_to_site_topology", "keep_unary",
-        "keep_unary_in_individuals", "keep_input_roots", NULL };
+    static char *kwlist[]
+        = { "samples", "filter_sites", "filter_populations", "filter_individuals",
+              "filter_nodes", "update_sample_flags", "reduce_to_site_topology",
+              "keep_unary", "keep_unary_in_individuals", "keep_input_roots", NULL };
 
     if (TableCollection_check_state(self) != 0) {
         goto out;
     }
-    if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|iiiiiiii", kwlist, &samples,
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|iiiiiiiii", kwlist, &samples,
             &filter_sites, &filter_populations, &filter_individuals, &filter_nodes,
-            &reduce_to_site_topology, &keep_unary, &keep_unary_in_individuals,
-            &keep_input_roots)) {
+            &update_sample_flags, &reduce_to_site_topology, &keep_unary,
+            &keep_unary_in_individuals, &keep_input_roots)) {
         goto out;
     }
     samples_array = (PyArrayObject *) PyArray_FROMANY(
@@ -6625,6 +6627,9 @@ TableCollection_simplify(TableCollection *self, PyObject *args, PyObject *kwds)
     if (!filter_nodes) {
         options |= TSK_SIMPLIFY_NO_FILTER_NODES;
     }
+    if (!update_sample_flags) {
+        options |= TSK_SIMPLIFY_NO_UPDATE_SAMPLE_FLAGS;
+    }
     if (reduce_to_site_topology) {
         options |= TSK_SIMPLIFY_REDUCE_TO_SITE_TOPOLOGY;
     }

python/tests/simplify.py

@@ -111,7 +111,8 @@ def __init__(
         keep_unary=False,
         keep_unary_in_individuals=False,
         keep_input_roots=False,
-        filter_nodes=True,  # If this is False, the order in `sample` is ignored
+        filter_nodes=True,
+        update_sample_flags=True,
     ):
         self.ts = ts
         self.n = len(sample)
@@ -121,6 +122,7 @@ def __init__(
         self.filter_populations = filter_populations
         self.filter_individuals = filter_individuals
         self.filter_nodes = filter_nodes
+        self.update_sample_flags = update_sample_flags
         self.keep_unary = keep_unary
         self.keep_unary_in_individuals = keep_unary_in_individuals
         self.keep_input_roots = keep_input_roots
@@ -152,14 +154,14 @@ def __init__(
             # NOTE In the C implementation we would really just not touch the
             # original tables.
             self.tables.nodes.replace_with(self.ts.tables.nodes)
-            # TODO make this optional somehow
-            flags = self.tables.nodes.flags
-            # Zero out other sample flags
-            flags = np.bitwise_and(flags, ~tskit.NODE_IS_SAMPLE)
-            flags[sample] |= tskit.NODE_IS_SAMPLE
-            self.tables.nodes.flags = flags.astype(np.uint32)
-            self.node_id_map[:] = np.arange(ts.num_nodes)
+            if self.update_sample_flags:
+                flags = self.tables.nodes.flags
+                # Zero out other sample flags
+                flags = np.bitwise_and(flags, ~tskit.NODE_IS_SAMPLE)
+                flags[sample] |= tskit.NODE_IS_SAMPLE
+                self.tables.nodes.flags = flags.astype(np.uint32)
 
+            self.node_id_map[:] = np.arange(ts.num_nodes)
             for sample_id in sample:
                 self.add_ancestry(sample_id, 0, self.sequence_length, sample_id)
         else:
@@ -178,10 +180,11 @@ def record_node(self, input_id):
         """
         node = self.ts.node(input_id)
         flags = node.flags
-        # Need to zero out the sample flag
-        flags &= ~tskit.NODE_IS_SAMPLE
-        if self.is_sample[input_id]:
-            flags |= tskit.NODE_IS_SAMPLE
+        if self.update_sample_flags:
+            # Need to zero out the sample flag
+            flags &= ~tskit.NODE_IS_SAMPLE
+            if self.is_sample[input_id]:
+                flags |= tskit.NODE_IS_SAMPLE
         output_id = self.tables.nodes.append(node.replace(flags=flags))
         self.node_id_map[input_id] = output_id
         return output_id

python/tests/test_lowlevel.py

@@ -349,6 +349,8 @@ def test_simplify_bad_args(self):
             tc.simplify([0, 1], filter_populations="x")
         with pytest.raises(TypeError):
             tc.simplify([0, 1], filter_nodes="x")
+        with pytest.raises(TypeError):
+            tc.simplify([0, 1], update_sample_flags="x")
         with pytest.raises(_tskit.LibraryError):
             tc.simplify([0, -1])
 
@@ -360,6 +362,7 @@ def test_simplify_bad_args(self):
             "filter_populations",
             "filter_individuals",
             "filter_nodes",
+            "update_sample_flags",
             "reduce_to_site_topology",
             "keep_unary",
             "keep_unary_in_individuals",

python/tests/test_topology.py

@@ -2686,7 +2686,7 @@ def verify_simplify(
             filter_sites=filter_sites,
             keep_input_roots=keep_input_roots,
             filter_nodes=filter_nodes,
-            compare_lib=True,  # TMP
+            compare_lib=True,
         )
         if debug:
             print("before")
@@ -4757,6 +4757,7 @@ def do_simplify(
     filter_nodes=True,
     keep_unary=False,
     keep_input_roots=False,
+    update_sample_flags=True,
 ):
     """
     Runs the Python test implementation of simplify.
@@ -4772,6 +4773,7 @@ def do_simplify(
         filter_nodes=filter_nodes,
         keep_unary=keep_unary,
         keep_input_roots=keep_input_roots,
+        update_sample_flags=update_sample_flags,
     )
     new_ts, node_map = s.simplify()
     if compare_lib:
@@ -4781,28 +4783,16 @@ def do_simplify(
             filter_individuals=filter_individuals,
             filter_populations=filter_populations,
             filter_nodes=filter_nodes,
+            update_sample_flags=update_sample_flags,
             keep_unary=keep_unary,
             keep_input_roots=keep_input_roots,
             map_nodes=True,
         )
         lib_tables1 = sts.dump_tables()
 
-        lib_tables2 = ts.dump_tables()
-        lib_node_map2 = lib_tables2.simplify(
-            samples,
-            filter_sites=filter_sites,
-            keep_unary=keep_unary,
-            keep_input_roots=keep_input_roots,
-            filter_individuals=filter_individuals,
-            filter_populations=filter_populations,
-            filter_nodes=filter_nodes,
-        )
-
         py_tables = new_ts.dump_tables()
         py_tables.assert_equals(lib_tables1, ignore_provenance=True)
-        py_tables.assert_equals(lib_tables2, ignore_provenance=True)
         assert all(node_map == lib_node_map1)
-        assert all(node_map == lib_node_map2)
     return new_ts, node_map
 
 
@@ -6091,6 +6081,64 @@ def test_mutations_on_removed_branches(self):
         assert ts2.num_mutations == 0
 
 
+class TestSimplifyNoUpdateSampleFlags:
+    """
+    Tests for simplify when we don't update the sample flags.
+    """
+
+    def test_simple_case_filter_nodes(self):
+        # 2.00┊    6    ┊
+        #     ┊  ┏━┻━┓  ┊
+        # 1.00┊  4   5  ┊
+        #     ┊ ┏┻┓ ┏┻┓ ┊
+        # 0.00┊ 0 1 2 3 ┊
+        #     0         1
+        ts1 = tskit.Tree.generate_balanced(4).tree_sequence
+        ts2, node_map = do_simplify(
+            ts1,
+            [0, 1, 6],
+            update_sample_flags=False,
+        )
+        # Because we don't retain 2 and 3 here, they don't stay as
+        # samples. But, we specified 6 as a sample, so it's coming
+        # through where it would ordinarily be dropped.
+
+        # 2.00┊  2  ┊
+        #     ┊  ┃  ┊
+        # 1.00┊  3  ┊
+        #     ┊ ┏┻┓ ┊
+        # 0.00┊ 0 1 ┊
+        #     0     1
+        assert list(ts2.nodes_flags) == [1, 1, 0, 0]
+        tree = ts2.first()
+        assert list(tree.parent_array) == [3, 3, -1, 2, -1]
+
+    def test_simple_case_no_filter_nodes(self):
+        # 2.00┊    6    ┊
+        #     ┊  ┏━┻━┓  ┊
+        # 1.00┊  4   5  ┊
+        #     ┊ ┏┻┓ ┏┻┓ ┊
+        # 0.00┊ 0 1 2 3 ┊
+        #     0         1
+        ts1 = tskit.Tree.generate_balanced(4).tree_sequence
+        ts2, node_map = do_simplify(
+            ts1,
+            [0, 1, 6],
+            update_sample_flags=False,
+            filter_nodes=False,
+        )
+
+        # 2.00┊  6      ┊
+        #     ┊  ┃      ┊
+        # 1.00┊  4      ┊
+        #     ┊ ┏┻┓     ┊
+        # 0.00┊ 0 1 2 3 ┊
+        #     0         1
+        assert list(ts2.nodes_flags) == list(ts1.nodes_flags)
+        tree = ts2.first()
+        assert list(tree.parent_array) == [4, 4, -1, -1, 6, -1, -1, -1]
+
+
 class TestMapToAncestors:
     """
     Tests the AncestorMap class.