Delete topology in flanking regions

Fixes #483

Author: hyanwong

evaluation.py

@@ -1692,11 +1692,13 @@ def run_ancestor_quality(args):
 def get_node_degree_by_depth(ts):
     """
     Returns a tuple (degree, depth) for each node in each tree in the
-    specified tree sequence.
+    specified tree sequence (empty flanking regions are omitted)
     """
     degree = []
     depth = []
     for tree in ts.trees():
+        if tree.num_edges == 0 and (tree.index == 0 or tree.index == ts.num_trees - 1):
+            continue
         stack = [(tree.root, 0)]
         while len(stack) > 0:
             u, d = stack.pop()

tests/test_inference.py

@@ -753,7 +753,8 @@ def verify(self, genotypes):
                 exclude_positions=exclude_positions,
             )
             for tree in output_ts.trees():
-                assert tree.num_roots == 1
+                if tree.num_edges > 0 or 0 < tree.index < output_ts.num_trees - 1:
+                    assert tree.num_roots == 1
             for site in output_ts.sites():
                 inf_type = json.loads(site.metadata)["inference_type"]
                 if len(site.mutations) == 0:
@@ -2682,6 +2683,23 @@ def test_single_tree(self, small_ts_fixture):
     def test_many_trees(self, medium_ts_fixture):
         self.verify(medium_ts_fixture)
 
+    def test_flanking_regions_deleted(self, small_sd_fixture):
+        ts1 = tsinfer.infer(small_sd_fixture)
+        assert ts1.site(-1).position + 1 < ts1.sequence_length
+        assert ts1.first().num_edges == 0
+        assert ts1.last().num_edges == 0
+        assert ts1.first().interval.right == small_sd_fixture.sites_position[0]
+        assert ts1.last().interval.left == small_sd_fixture.sites_position[-1] + 1
+
+        # If seq length is less than the last pos + 1, right flank is not deleted
+        sd = small_sd_fixture.subset(sequence_length=ts1.site(-1).position + 0.1)
+        ts2 = tsinfer.infer(sd)
+        assert ts2.first().num_edges == 0
+        assert ts2.last().num_edges != 0
+        assert ts2.first().interval.right == sd.sites_position[0]
+
+        assert ts2.num_trees == ts1.num_trees - 1
+
     def test_standalone_post_process(self, medium_sd_fixture):
         # test we can post process separately, e.g. omitting the MRCA splitting step
         ts_unsimplified = tsinfer.infer(medium_sd_fixture, post_process=False)
@@ -2700,7 +2718,7 @@ def test_standalone_post_process(self, medium_sd_fixture):
         md = json.loads(md.decode())  # At the moment node metadata has no schema
         assert md["ancestor_data_id"] == 1
 
-        ts = tsinfer.post_process(ts_unsimplified, split_mrca=True)
+        ts = tsinfer.post_process(ts_unsimplified, split_mrca=True, erase_flanks=False)
         oldest_parent_id = ts.edge(-1).parent
         assert np.sum(ts.tables.nodes.time == ts.node(oldest_parent_id).time) > 1
         roots = set()
@@ -2719,11 +2737,13 @@ def test_post_process_non_tsinfer(self, small_ts_fixture, caplog):
             small_ts_fixture.samples() == np.arange(small_ts_fixture.num_samples)
         )
         with caplog.at_level(logging.WARNING):
-            ts_postprocessed = tsinfer.post_process(small_ts_fixture)
+            ts_postprocessed = tsinfer.post_process(
+                small_ts_fixture, erase_flanks=False
+            )
             assert caplog.text.count("virtual-root-like") == 0
         with caplog.at_level(logging.WARNING):
             ts_postprocessed = tsinfer.post_process(
-                small_ts_fixture, warn_if_unexpected_format=True
+                small_ts_fixture, warn_if_unexpected_format=True, erase_flanks=False
             )
             assert caplog.text.count("virtual-root-like") == 1
 
@@ -2747,8 +2767,8 @@ def test_virtual_like_root_removed(self, medium_sd_fixture):
 
     def test_split_edges_one_tree(self, small_sd_fixture):
         ts = tsinfer.infer(small_sd_fixture, post_process=False)
-        ts = tsinfer.post_process(ts, split_mrca=False)
         assert ts.num_trees == 1
+        ts = tsinfer.post_process(ts, split_mrca=False)
         # Check that we don't delete and recreate the oldest node if there's only 1 tree
         tables = ts.dump_tables()
         oldest_node_in_topology = tables.edges[-1].parent
@@ -2758,7 +2778,7 @@ def test_split_edges_one_tree(self, small_sd_fixture):
 
     def test_dont_split_edges_twice(self, medium_sd_fixture, caplog):
         ts = tsinfer.infer(medium_sd_fixture, post_process=False)
-        ts = tsinfer.post_process(ts, split_mrca=False)
+        ts = tsinfer.post_process(ts, split_mrca=False, erase_flanks=False)
         assert ts.num_trees > 1
         assert tsinfer.has_same_root_everywhere(ts)
         # Once the mrca has been split, it can't be split again
@@ -2801,6 +2821,18 @@ def test_sample_order(self, medium_sd_fixture):
             inferred_individual = ts.individual(ts.node(inferred_node_id).individual)
             assert sd_individual.metadata["id"] == inferred_individual.metadata["id"]
 
+    def test_erase_flanks(self, small_sd_fixture):
+        ts1 = tsinfer.infer(small_sd_fixture, post_process=False)
+        ts2 = tsinfer.post_process(ts1, erase_flanks=False)
+        assert ts2.first().num_edges > 0
+        assert ts2.last().num_edges > 0
+        assert ts1.num_trees == ts2.num_trees
+
+        ts2 = tsinfer.post_process(ts1, erase_flanks=True)
+        assert ts2.first().num_edges == 0
+        assert ts2.last().num_edges == 0
+        assert ts1.num_trees == ts2.num_trees - 2
+
 
 def get_default_inference_sites(sample_data):
     """
@@ -3990,7 +4022,8 @@ def test_nan_sites(self):
             sample_data.add_site(0.4, [1, 1, 0], time=np.nan)
             sample_data.add_site(0.6, [1, 1, 0])
         ts = tsinfer.infer(sample_data)
-        assert ts.num_trees == 1
+        num_nonempty_trees = sum(1 for tree in ts.trees() if tree.num_edges > 0)
+        assert num_nonempty_trees == 1
         inf_type = [json.loads(site.metadata)["inference_type"] for site in ts.sites()]
         assert inf_type[0] == tsinfer.INFERENCE_FULL
         assert inf_type[1] == tsinfer.INFERENCE_PARSIMONY

tsinfer/eval_util.py

@@ -776,9 +776,14 @@ def node_span(ts):
             for u in [edge.parent, edge.child]:
                 if start[u] == -1 and tree.num_samples(u) > 0:
                     start[u] = left
+    # add intervals for the last tree
     for u in tree.nodes():
         if tree.num_samples(u) > 0:
             S[u] += ts.sequence_length - start[u]
+    # isolated sample nodes (e.g. where topology was deleted) will have been missed when
+    # iterating over edges, but by definition they span the entire ts, so override them
+    for u in ts.samples():
+        S[u] = ts.sequence_length
     return S
 
 

tsinfer/inference.py

@@ -1950,29 +1950,35 @@ def post_process(
     ts,
     *,
     split_mrca=None,
+    erase_flanks=None,
     # Parameters below deliberately undocumented
     warn_if_unexpected_format=None,
     simplify_only=None,
 ):
     """
-    post_process(ts, *, split_mrca=None)
+    post_process(ts, *, split_mrca=None, erase_flanks=None)
 
     Post-process a tsinferred tree sequence into a more conventional form. This is
     the function run by default on the final tree sequence output by
-    :func:`match_samples`. It involves the following 3 steps:
+    :func:`match_samples`. It involves the following 4 steps:
 
     #. If the oldest node is connected to a single child via an edge that spans the
        entire tree sequence, this oldest node is removed, so that its child becomes
        the new root (this step is undertaken to remove the "virtual-root-like node"
        which is added to ancestor tree sequences to enable matching).
     #. If the oldest node is removed in the first step and the new root spans the
-       entire genome, it is treated as the "grand MRCA" and (if split_mrca is ``True``)
-       the node is split into multiple coexisiting nodes with the splits
+       entire genome, it is treated as the "grand MRCA" and (unless split_mrca is
+       ``False``) the node is split into multiple coexisiting nodes with the splits
        occurring whenever the children of the grand MRCA change. The rationale
        is that tsinfer creates a grand MRCA consisting of a single ancestral haplotype
        with all inference sites in the ancestral state: this is, however, unlikely
        to represent a single ancestor in the past. If nodes in the tree sequence are
        then dated, these MRCA nodes can be pushed to different times.
+    #. Often, extensive regions of genome exist before the first defined site and after
+       the last defined site. Since no data exists in these sections of the genome, post
+       processing by default erases the inferred topology in these regions. However,
+       if ``erase_flanks`` is False, the flanking regions at the start and end will be
+       assigned the same topology as inferred at the first and last site respectively.
     #. The sample nodes are reordered such that they are the first nodes listed in the
        node table,  removing tree nodes and edges that are not on a path between the
        root and any of the samples (by applying the :meth:`~tskit.TreeSequence.simplify`
@@ -1982,11 +1988,17 @@ def post_process(
     :param bool split_mrca: If ``True`` (default) and the oldest node is the only
         parent to a single "grand MRCA", split the grand MRCA into separate nodes
         (see above). If ``False`` do not attempt to identify or split a grand MRCA.
+    :param bool erase_flanks: If ``True`` (default), keep only the
+        inferred topology between the first and last sites. If ``False``,
+        output regions of topology inferred before the first site and after
+        the last site.
     :return: The post-processed tree sequence.
     :rtype: tskit.TreeSequence
     """
     if split_mrca is None:
         split_mrca = True
+    if erase_flanks is None:
+        erase_flanks = True
 
     tables = ts.dump_tables()
 
@@ -2011,10 +2023,21 @@ def post_process(
                 "Cannot find a virtual-root-like ancestor during preprocessing"
             )
 
+        if erase_flanks and ts.num_sites > 0:
+            logger.info("Removing topology in flanking regions with keep_intervals")
+            # So that the last site falls within a tree, we must add one to the
+            # site position (or simply extend to the end of the ts)
+            upper_cutoff = min(ts.sites_position[-1] + 1, ts.sequence_length)
+            tables.keep_intervals(
+                [[ts.sites_position[0], upper_cutoff]],
+                simplify=False,
+                record_provenance=False,
+            )
+
     logger.info(
         "Simplifying with filter_sites=False, filter_populations=False, "
         "filter_individuals=False, and keep_unary=True on "
-        f"{tables.nodes.num_rows} nodes and {ts.num_edges} edges"
+        f"{tables.nodes.num_rows} nodes and {tables.edges.num_rows} edges"
     )
     # NB: if this is an inferred TS, match_samples is guaranteed to produce samples
     # in the same order as passed in to sample_indexes, and simplification will
@@ -2024,6 +2047,7 @@ def post_process(
         filter_populations=False,
         filter_individuals=False,
         keep_unary=True,
+        record_provenance=False,
     )
     logger.info(
         "Finished simplify; now have {} nodes and {} edges".format(