Rename remove_telomeres to erase_flanks

And move unit tests to one place. Fixes #398

Author: hyanwong

CHANGELOG.md

@@ -12,6 +12,8 @@
 - An environment variable `TSDATE_ENABLE_NUMBA_CACHE` can be set to cache JIT
   compiled code, speeding up loading time (useful when testing).
 
+- The time taken for running _tsdate_ is now recorded in the provenance data
+
 **Documentation**
 
 - Various fixes in documentation, including documenting returned fits.
@@ -27,6 +29,10 @@
   of more than one tree, as tests have found that span-weighting the conditional coalescent
   causes substantial bias.
 
+- The `trim_telomeres` parameter in the `tsdate.preprocess_ts()` function has been renamed
+  to `erase_flanks`, to match `tinfer.preprocess()`. The previous name is kept as a
+  deprecated alias.
+
 ## [0.2.1] - 2024-07-31
 
 **Bugfixes**

docs/usage.md

@@ -133,7 +133,7 @@ print(
 ```
 :::{note}
 In simulated data you may not have missing data regions, and you may be able to
-pass `remove_telomeres=False` to the `preprocess_ts` function.
+pass `erase_flanks=False` to the `preprocess_ts` function.
 :::
 
 The inference in this case is much more noisy (as illustrated using the original

tests/test_cli.py

@@ -146,7 +146,7 @@ def test_default_values_preprocess(self):
         assert args.tree_sequence == self.infile
         assert args.output == self.output
         assert args.minimum_gap == 1000000
-        assert args.trim_telomeres
+        assert args.erase_flanks
         assert args.split_disjoint
 
 

tests/test_functions.py

@@ -1691,96 +1691,6 @@ def test_constrain_ages_leastsquare(self):
         assert r2_2 > r2_1
 
 
-class TestPreprocessTs(unittest.TestCase):
-    """
-    Test preprocess_ts works as expected
-    """
-
-    def verify(self, ts, minimum_gap=None, remove_telomeres=None, **kwargs):
-        with self.assertLogs("tsdate.util", level="INFO") as logs:
-            if minimum_gap is not None and remove_telomeres is not None:
-                ts = tsdate.preprocess_ts(
-                    ts, minimum_gap=minimum_gap, remove_telomeres=remove_telomeres
-                )
-            elif minimum_gap is not None and remove_telomeres is None:
-                ts = tsdate.preprocess_ts(ts, minimum_gap=minimum_gap)
-            elif remove_telomeres is not None and minimum_gap is None:
-                ts = tsdate.preprocess_ts(ts, remove_telomeres=remove_telomeres)
-            else:
-                ts = tsdate.preprocess_ts(ts, **kwargs)
-        messages = [record.msg for record in logs.records]
-        assert "Beginning preprocessing" in messages
-        return ts
-
-    def test_no_sites(self):
-        ts = utility_functions.two_tree_ts()
-        with pytest.raises(ValueError):
-            tsdate.preprocess_ts(ts)
-
-    def test_invariant_sites(self):
-        # Test that invariant sites are not removed by default
-        # (and simularly for unused individuals & populations)
-        ts = utility_functions.site_no_mutations()
-        assert ts.num_sites != 0
-        assert ts.num_individuals != 0
-        assert ts.num_populations != 0
-        removed = self.verify(ts)
-        assert removed.num_sites == ts.num_sites
-        assert removed.num_individuals == ts.num_individuals
-        assert removed.num_populations == ts.num_populations
-        assert tsdate.preprocess_ts(ts, **{"filter_sites": True}).num_sites == 0
-        assert (
-            tsdate.preprocess_ts(ts, **{"filter_populations": True}).num_populations == 0
-        )
-        assert (
-            tsdate.preprocess_ts(ts, **{"filter_individuals": True}).num_individuals == 0
-        )
-
-    def test_no_intervals(self):
-        ts = utility_functions.two_tree_mutation_ts()
-        assert ts.tables.edges == self.verify(ts, remove_telomeres=False).tables.edges
-        assert ts.tables.edges == self.verify(ts, minimum_gap=0.05).tables.edges
-
-    def test_passed_intervals(self):
-        # Mostly we would not pass in user-defined intervals: this is mainly for testing
-        ts = utility_functions.single_tree_ts_n3()  # No sites!
-        ts = tsdate.preprocess_ts(
-            ts, delete_intervals=[(0, 0.1), (0.5, ts.sequence_length)]
-        )
-        assert ts.num_edges > 1
-        assert np.allclose(ts.edges_left, 0.1)
-        assert np.allclose(ts.edges_right, 0.5)
-
-    def test_bad_delete_intervals(self):
-        ts = utility_functions.two_tree_mutation_ts()
-        with pytest.raises(ValueError, match="specify both"):
-            tsdate.preprocess_ts(ts, delete_intervals=[(0, 0.1)], minimum_gap=0.05)
-        with pytest.raises(ValueError, match="specify both"):
-            tsdate.preprocess_ts(ts, delete_intervals=[(0, 0.1)], remove_telomeres=True)
-
-    def test_delete_interval(self):
-        ts = utility_functions.ts_w_data_desert(40, 60, 100)
-        trimmed = self.verify(ts, minimum_gap=20, remove_telomeres=False)
-        lefts = trimmed.edges_left
-        rights = trimmed.edges_right
-        assert not np.any(np.logical_and(lefts > 41, lefts < 59))
-        assert not np.any(np.logical_and(rights > 41, rights < 59))
-
-    def test_remove_telomeres(self):
-        ts = utility_functions.ts_w_data_desert(0, 5, 100)
-        removed = self.verify(ts, minimum_gap=ts.get_sequence_length())
-        lefts = removed.tables.edges.left
-        rights = removed.tables.edges.right
-        assert not np.any(np.logical_and(lefts > 0, lefts < 4))
-        assert not np.any(np.logical_and(rights > 0, rights < 4))
-        ts = utility_functions.ts_w_data_desert(95, 100, 100)
-        removed = self.verify(ts, minimum_gap=ts.get_sequence_length())
-        lefts = removed.tables.edges.left
-        rights = removed.tables.edges.right
-        assert not np.any(np.logical_and(lefts > 96, lefts < 100))
-        assert not np.any(np.logical_and(rights > 96, rights < 100))
-
-
 class TestNodeTimes:
     """
     Test node_times works as expected.

tests/test_provenance.py

@@ -113,7 +113,7 @@ def test_preprocess_defaults_recorded(self):
         preprocessed_ts = tsdate.preprocess_ts(ts)
         assert preprocessed_ts.num_provenances == num_provenances + 1
         rec = json.loads(preprocessed_ts.provenance(-1).record)
-        assert rec["parameters"]["remove_telomeres"]
+        assert rec["parameters"]["erase_flanks"]
         assert rec["parameters"]["minimum_gap"] == 1000000
         assert rec["parameters"]["delete_intervals"] == []
 
@@ -124,8 +124,8 @@ def test_preprocess_interval_recorded(self):
         assert preprocessed_ts.num_provenances == num_provenances + 1
         rec = json.loads(preprocessed_ts.provenance(-1).record)
         assert rec["parameters"]["minimum_gap"] == 20
-        assert rec["parameters"]["remove_telomeres"] is not None
-        assert not rec["parameters"]["remove_telomeres"]
+        assert rec["parameters"]["erase_flanks"] is not None
+        assert not rec["parameters"]["erase_flanks"]
         deleted_intervals = rec["parameters"]["delete_intervals"]
         assert len(deleted_intervals) == 1
         assert deleted_intervals[0][0] < deleted_intervals[0][1]

tests/test_util.py

@@ -31,6 +31,7 @@
 import pytest
 import tsinfer
 import tskit
+import utility_functions
 
 import tsdate
 
@@ -203,6 +204,85 @@ def test_inferred(self):
 
 
 class TestPreprocessTs:
+    def verify(self, ts, caplog, minimum_gap=None, erase_flanks=None, **kwargs):
+        with caplog.at_level(logging.INFO):
+            if minimum_gap is not None and erase_flanks is not None:
+                ts = tsdate.preprocess_ts(
+                    ts, minimum_gap=minimum_gap, erase_flanks=erase_flanks
+                )
+            elif minimum_gap is not None and erase_flanks is None:
+                ts = tsdate.preprocess_ts(ts, minimum_gap=minimum_gap)
+            elif erase_flanks is not None and minimum_gap is None:
+                ts = tsdate.preprocess_ts(ts, erase_flanks=erase_flanks)
+            else:
+                ts = tsdate.preprocess_ts(ts, **kwargs)
+
+            assert "Beginning preprocessing" in caplog.text
+        return ts
+
+    def test_invariant_sites(self, caplog):
+        # Test that invariant sites are not removed by default
+        # (and simularly for unused individuals & populations)
+        ts = utility_functions.site_no_mutations()
+        assert ts.num_sites != 0
+        assert ts.num_individuals != 0
+        assert ts.num_populations != 0
+        removed = self.verify(ts, caplog)
+        assert removed.num_sites == ts.num_sites
+        assert removed.num_individuals == ts.num_individuals
+        assert removed.num_populations == ts.num_populations
+        assert tsdate.preprocess_ts(ts, **{"filter_sites": True}).num_sites == 0
+        assert (
+            tsdate.preprocess_ts(ts, **{"filter_populations": True}).num_populations == 0
+        )
+        assert (
+            tsdate.preprocess_ts(ts, **{"filter_individuals": True}).num_individuals == 0
+        )
+
+    def test_no_intervals(self, caplog):
+        ts = utility_functions.two_tree_mutation_ts()
+        assert ts.tables.edges == self.verify(ts, caplog, erase_flanks=False).tables.edges
+        assert ts.tables.edges == self.verify(ts, caplog, minimum_gap=0.05).tables.edges
+
+    def test_passed_intervals(self):
+        # Mostly we would not pass in user-defined intervals: this is mainly for testing
+        ts = utility_functions.single_tree_ts_n3()  # No sites!
+        ts = tsdate.preprocess_ts(
+            ts, delete_intervals=[(0, 0.1), (0.5, ts.sequence_length)]
+        )
+        assert ts.num_edges > 1
+        assert np.allclose(ts.edges_left, 0.1)
+        assert np.allclose(ts.edges_right, 0.5)
+
+    def test_bad_delete_intervals(self):
+        ts = utility_functions.two_tree_mutation_ts()
+        with pytest.raises(ValueError, match="specify both"):
+            tsdate.preprocess_ts(ts, delete_intervals=[(0, 0.1)], minimum_gap=0.05)
+        with pytest.raises(ValueError, match="specify both"):
+            tsdate.preprocess_ts(ts, delete_intervals=[(0, 0.1)], erase_flanks=True)
+
+    def test_delete_interval(self, caplog):
+        ts = utility_functions.ts_w_data_desert(40, 60, 100)
+        trimmed = self.verify(ts, caplog, minimum_gap=20, erase_flanks=False)
+        lefts = trimmed.edges_left
+        rights = trimmed.edges_right
+        assert not np.any(np.logical_and(lefts > 41, lefts < 59))
+        assert not np.any(np.logical_and(rights > 41, rights < 59))
+
+    def test_erase_flanks(self, caplog):
+        ts = utility_functions.ts_w_data_desert(0, 5, 100)
+        removed = self.verify(ts, caplog, minimum_gap=ts.get_sequence_length())
+        lefts = removed.tables.edges.left
+        rights = removed.tables.edges.right
+        assert not np.any(np.logical_and(lefts > 0, lefts < 4))
+        assert not np.any(np.logical_and(rights > 0, rights < 4))
+        ts = utility_functions.ts_w_data_desert(95, 100, 100)
+        removed = self.verify(ts, caplog, minimum_gap=ts.get_sequence_length())
+        lefts = removed.tables.edges.left
+        rights = removed.tables.edges.right
+        assert not np.any(np.logical_and(lefts > 96, lefts < 100))
+        assert not np.any(np.logical_and(rights > 96, rights < 100))
+
     def test_no_sites(self):
         ts = tskit.Tree.generate_comb(3).tree_sequence
         with pytest.raises(ValueError, match="no sites"):
@@ -273,19 +353,25 @@ def test_record_provenance(self):
         ts = tsdate.preprocess_ts(ts, record_provenance=False)
         assert ts.num_provenances == num_provenances + 1
 
-    def test_trim_flanks(self):
+    def test_no_erase_flanks(self):
         tables = tskit.Tree.generate_comb(3, span=100).tree_sequence.dump_tables()
         tables.sites.add_row(10, "A")
         tables.sites.add_row(90, "A")
         ts = tables.tree_sequence()
         assert ts.sequence_length == 100
         assert ts.num_trees == 1
-        ts = tsdate.preprocess_ts(ts)
-        assert ts.num_trees == 3
-        assert ts.first().num_edges == 0
-        assert ts.first().interval.right == 10 - 1
-        assert ts.last().num_edges == 0
-        assert ts.last().interval.left == 90 + 1
+        for param_name in ("erase_flanks", "remove_telomeres"):
+            params = {param_name: False}
+            new_ts = tsdate.preprocess_ts(ts, **params)
+            assert new_ts.num_trees == 1
+            assert new_ts.sequence_length == 100
+
+    @pytest.mark.parametrize("bool1", [True, False])
+    @pytest.mark.parametrize("bool2", [True, False])
+    def test_erase_flanks_telomeres_combo(self, bool1, bool2):
+        ts = tskit.Tree.generate_comb(3, span=100).tree_sequence
+        with pytest.raises(ValueError, match="specify both"):
+            tsdate.preprocess_ts(ts, erase_flanks=bool1, remove_telomeres=bool2)
 
     def test_sim_example(self):
         # Test a larger example
@@ -310,8 +396,6 @@ def test_sim_example(self):
         # Next assumes no breakpoints before first site or after last
         assert ts.num_trees == num_trees + first_empty + last_empty
 
-    # TODO - test minimum_gap param
-
 
 class TestUnaryNodeCheck:
     def test_inferred(self):

tsdate/cli.py

@@ -225,7 +225,8 @@ def tsdate_cli_parser():
         default=1000000,
     )
     parser.add_argument(
-        "--trim-telomeres",
+        "--erase-flanks",
+        "--trim_telomeres",
         type=bool,
         help=(
             "Should all material before the first site and after the "
@@ -313,7 +314,7 @@ def run_preprocess(args):
     except tskit.FileFormatError as ffe:
         error_exit(f"FileFormatError loading '{args.tree_sequence}: {ffe}")
     snipped_ts = tsdate.preprocess_ts(
-        ts, minimum_gap=args.minimum_gap, remove_telomeres=args.trim_telomeres
+        ts, minimum_gap=args.minimum_gap, erase_flanks=args.erase_flanks
     )
     snipped_ts.dump(args.output)