Merge pull request #275 from hyanwong/provenance-saving

Provenance saving

Author: hyanwong

CHANGELOG.rst

@@ -33,6 +33,14 @@
   probabilities (scaled so that they sum to one) rather than standardized
   "probabilities" whose maximum value is one.
 
+- The population size is saved in provenance metadata when possible (e.g. if it is
+  a single number)
+
+- ``preprocess_ts`` always records provenance as being from the ``preprocess_ts``
+  command, even if no gaps are removed. The command now has a (rarely used)
+  `delete_intervals` parameter, which is normally filled out and saved in provenance
+  (as it was before). If no gap deletion is done, the param is saved as ``[]``
+
 --------------------
 [0.1.5] - 2022-06-07
 --------------------

tests/test_functions.py

@@ -1781,11 +1781,28 @@ def test_no_intervals(self):
         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.tables.edges.left
-        rights = trimmed.tables.edges.right
+        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))
 

tests/test_provenance.py

@@ -0,0 +1,107 @@
+# MIT License
+#
+# Copyright (c) 2021-23 Tskit Developers
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+"""
+Test cases for saving provenances in tsdate.
+"""
+import json
+
+import numpy as np
+import pytest
+import utility_functions
+
+import tsdate
+from tsdate import provenance
+
+
+class TestProvenance:
+    def test_bad_get_dict(self):
+        with pytest.raises(ValueError, match="cannot be None"):
+            provenance.get_provenance_dict(None)
+
+    def test_date_cmd_recorded(self):
+        ts = utility_functions.single_tree_ts_n2()
+        dated_ts = tsdate.date(ts, population_size=1, mutation_rate=None)
+        rec = json.loads(dated_ts.provenance(-1).record)
+        assert rec["software"]["name"] == "tsdate"
+        assert rec["parameters"]["command"] == "date"
+
+    def test_date_params_recorded(self):
+        ts = utility_functions.single_tree_ts_n2()
+        mu = 0.123
+        Ne = 9
+        dated_ts = tsdate.date(ts, population_size=Ne, mutation_rate=mu)
+        rec = json.loads(dated_ts.provenance(-1).record)
+        assert np.isclose(rec["parameters"]["mutation_rate"], mu)
+        assert np.isclose(rec["parameters"]["population_size"], Ne)
+
+    @pytest.mark.skip(
+        reason="Not implemented yet: https://github.com/tskit-dev/tsdate/issues/274"
+    )
+    def test_date_popsizehist_recorded(self):
+        ts = utility_functions.single_tree_ts_n2()
+        mu = 0.123
+        popsize = tsdate.demography.PopulationSizeHistory(
+            np.array([1, 1]), np.array([1])
+        )
+        dated_ts = tsdate.date(ts, population_size=popsize, mutation_rate=mu)
+        rec = json.loads(dated_ts.provenance(-1).record)
+        assert np.isclose(rec["parameters"]["mutation_rate"], mu)
+        assert "population_size" in rec["parameters"]
+        # TODO: check that the population size history is recorded correctly
+        # https://github.com/tskit-dev/tsdate/issues/274
+        assert np.isclose(
+            # This is wrong - left in for now to show the sort of thing we want
+            (np.array(rec["parameters"]["population_size"])),
+            popsize,
+        )
+
+    def test_preprocess_cmd_recorded(self):
+        ts = utility_functions.ts_w_data_desert(40, 60, 100)
+        preprocessed_ts = tsdate.preprocess_ts(ts)
+        rec = json.loads(preprocessed_ts.provenance(-1).record)
+        assert rec["software"]["name"] == "tsdate"
+        assert rec["parameters"]["command"] == "preprocess_ts"
+
+    def test_preprocess_defaults_recorded(self):
+        ts = utility_functions.ts_w_data_desert(40, 60, 100)
+        preprocessed_ts = tsdate.preprocess_ts(ts)
+        rec = json.loads(preprocessed_ts.provenance(-1).record)
+        assert rec["parameters"]["remove_telomeres"]
+        assert rec["parameters"]["minimum_gap"] == 1000000
+        assert rec["parameters"]["delete_intervals"] == []
+
+    def test_preprocess_interval_recorded(self):
+        ts = utility_functions.ts_w_data_desert(40, 60, 100)
+        preprocessed_ts = tsdate.preprocess_ts(
+            ts, minimum_gap=20, remove_telomeres=False
+        )
+        rec = json.loads(preprocessed_ts.provenance(-1).record)
+        print(rec)
+        assert rec["parameters"]["minimum_gap"] == 20
+        assert rec["parameters"]["remove_telomeres"] is not None
+        assert not rec["parameters"]["remove_telomeres"]
+        deleted_intervals = rec["parameters"]["delete_intervals"]
+        assert len(deleted_intervals) == 1
+        assert deleted_intervals[0][0] < deleted_intervals[0][1]
+        assert 40 < deleted_intervals[0][0] < 60
+        assert 40 < deleted_intervals[0][1] < 60

tsdate/core.py

@@ -1215,11 +1215,7 @@ def date(
         Ignoring outside root provides greater stability when dating tree sequences
         inferred from real data. Default: False
     :param bool progress: Whether to display a progress bar. Default: False
-    :param float Ne: the estimated (diploid) effective population size used to
-        construct the (default) conditional coalescent prior. This is used when
-        ``priors`` is ``None``.  Conversely, if ``priors`` is not ``None``, no
-        ``population_size`` value should be given.  (Deprecated, use the
-        ``population_size`` argument instead).
+    :param float Ne: Deprecated, use the``population_size`` argument instead.
     :return: A copy of the input tree sequence but with altered node times, or (if
         ``return_posteriors`` is True) a tuple of that tree sequence plus a dictionary
         of posterior probabilities from the "inside_outside" estimation ``method``.
@@ -1262,13 +1258,19 @@ def date(
     # Remove any times associated with mutations
     tables.mutations.time = np.full(tree_sequence.num_mutations, tskit.UNKNOWN_TIME)
     tables.sort()
-    # TODO: record population_size provenance, or record that it is omitted
-    provenance.record_provenance(
-        tables,
-        "date",
+    params = dict(
         mutation_rate=mutation_rate,
         recombination_rate=recombination_rate,
         progress=progress,
+    )
+    if isinstance(population_size, (int, float)):
+        params["population_size"] = population_size
+    else:
+        params["population_size"] = "TODO: PopulationSizeHistory object"
+    provenance.record_provenance(
+        tables,
+        "date",
+        **params,
         **kwargs,
     )
     if return_posteriors:

tsdate/demography.py

@@ -245,3 +245,7 @@ def to_gamma(self, shape=1, rate=1):
     #    """
     #    Create a `PopulationSizeHistory` instance from a `demes` format YAML
     #    """
+
+    # TODO:
+    # output a PopulationSizeHistory as a json object, so we can store it
+    # in the provenance. See https://github.com/tskit-dev/tsdate/issues/274

tsdate/provenance.py

@@ -63,13 +63,13 @@ def get_environment():
     return env
 
 
-def get_provenance_dict(command=None, **kwargs):
+def get_provenance_dict(command, **kwargs):
     """
     Returns a dictionary encoding an execution of tsdate conforming to the
     tskit provenance schema.
     """
     if command is None:
-        raise ValueError("Command must be provided")
+        raise ValueError("`command` cannot be None")
     parameters = dict(kwargs)
     parameters["command"] = command
     document = {

tsdate/util.py

@@ -64,11 +64,12 @@ def reduce_to_contemporaneous(ts):
 def preprocess_ts(
     tree_sequence,
     *,
-    minimum_gap=1000000,
-    remove_telomeres=True,
+    minimum_gap=None,
+    remove_telomeres=None,
     filter_populations=False,
     filter_individuals=False,
     filter_sites=False,
+    delete_intervals=None,
     **kwargs,
 ):
     """
@@ -81,9 +82,10 @@ def preprocess_ts(
     :param TreeSequence tree_sequence: The input :class`tskit.TreeSequence`
         to be preprocessed.
     :param float minimum_gap: The minimum gap between sites to remove from the tree
-        sequence. Default: "1000000"
+        sequence. Default: ``None`` treated as ``1000000``
     :param bool remove_telomeres: Should all material before the first site and after the
-        last site be removed, regardless of the length. Default: "True"
+        last site be removed, regardless of the length. Default: ``None`` treated as
+        ``True``
     :param bool filter_populations: parameter passed to the ``tskit.simplify``
         command. Unlike calling that command directly, this defaults to ``False``, such
         that all populations in the tree sequence are kept.
@@ -93,71 +95,96 @@ def preprocess_ts(
     :param bool filter_sites: parameter passed to the ``tskit.simplify``
         command. Unlike calling that command directly, this defaults to ``False``, such
         that all sites in the tree sequence are kept
+    :param array_like delete_intervals: A list (start, end) pairs describing the
+        genomic intervals (gaps) to delete. This is usually left as ``None``
+        (the default) in which case ``minimum_gap`` and ``remove_telomeres`` are used
+        to determine the gaps to remove, and the calculated intervals are recorded in
+        the provenance of the resulting tree sequence.
     :param \\**kwargs: All further keyword arguments are passed to the ``tskit.simplify``
         command.
 
     :return: A tree sequence with gaps removed.
     :rtype: tskit.TreeSequence
     """
+
     logger.info("Beginning preprocessing")
     logger.info(f"Minimum_gap: {minimum_gap} and remove_telomeres: {remove_telomeres}")
-    if tree_sequence.num_sites < 1:
-        raise ValueError("Invalid tree sequence: no sites present")
+    if delete_intervals is not None and (
+        minimum_gap is not None or remove_telomeres is not None
+    ):
+        raise ValueError(
+            "Cannot specify both delete_intervals and minimum_gap/remove_telomeres"
+        )
 
     tables = tree_sequence.dump_tables()
     sites = tables.sites.position[:]
-    delete_intervals = []
-    if remove_telomeres:
-        first_site = sites[0] - 1
-        if first_site > 0:
-            delete_intervals.append([0, first_site])
-            logger.info(
-                "REMOVING TELOMERE: Snip topology "
-                "from 0 to first site at {}.".format(first_site)
-            )
-        last_site = sites[-1] + 1
-        sequence_length = tables.sequence_length
-        if last_site < sequence_length:
-            delete_intervals.append([last_site, sequence_length])
-            logger.info(
-                "REMOVING TELOMERE: Snip topology "
-                "from {} to end of sequence at {}.".format(last_site, sequence_length)
-            )
-    gaps = sites[1:] - sites[:-1]
-    threshold_gaps = np.where(gaps >= minimum_gap)[0]
-    for gap in threshold_gaps:
-        gap_start = sites[gap] + 1
-        gap_end = sites[gap + 1] - 1
-        if gap_end > gap_start:
-            logger.info(
-                "Gap Size is {}. Snip topology "
-                "from {} to {}.".format(gap_end - gap_start, gap_start, gap_end)
-            )
-            delete_intervals.append([gap_start, gap_end])
-    delete_intervals = sorted(delete_intervals, key=lambda x: x[0])
+    if delete_intervals is None:
+        if minimum_gap is None:
+            minimum_gap = 1000000
+        if remove_telomeres is None:
+            remove_telomeres = True
+
+        if tree_sequence.num_sites < 1:
+            raise ValueError("Invalid tree sequence: no sites present")
+        delete_intervals = []
+        if remove_telomeres:
+            first_site = sites[0] - 1
+            if first_site > 0:
+                delete_intervals.append([0, first_site])
+                logger.info(
+                    "REMOVING TELOMERE: Snip topology "
+                    "from 0 to first site at {}.".format(first_site)
+                )
+            last_site = sites[-1] + 1
+            sequence_length = tables.sequence_length
+            if last_site < sequence_length:
+                delete_intervals.append([last_site, sequence_length])
+                logger.info(
+                    "REMOVING TELOMERE: Snip topology "
+                    "from {} to end of sequence at {}.".format(
+                        last_site, sequence_length
+                    )
+                )
+        gaps = sites[1:] - sites[:-1]
+        threshold_gaps = np.where(gaps >= minimum_gap)[0]
+        for gap in threshold_gaps:
+            gap_start = sites[gap] + 1
+            gap_end = sites[gap + 1] - 1
+            if gap_end > gap_start:
+                logger.info(
+                    "Gap Size is {}. Snip topology "
+                    "from {} to {}.".format(gap_end - gap_start, gap_start, gap_end)
+                )
+                delete_intervals.append([gap_start, gap_end])
+        delete_intervals = sorted(delete_intervals, key=lambda x: x[0])
     if len(delete_intervals) > 0:
         tables.delete_intervals(delete_intervals, simplify=False)
         tables.simplify(
             filter_populations=filter_populations,
             filter_individuals=filter_individuals,
             filter_sites=filter_sites,
+            record_provenance=False,
             **kwargs,
         )
-        provenance.record_provenance(
-            tables,
-            "preprocess_ts",
-            minimum_gap=minimum_gap,
-            remove_telomeres=remove_telomeres,
-            delete_intervals=delete_intervals,
-        )
     else:
         logger.info("No gaps to remove")
         tables.simplify(
             filter_populations=filter_populations,
             filter_individuals=filter_individuals,
             filter_sites=filter_sites,
+            record_provenance=False,
             **kwargs,
         )
+    provenance.record_provenance(
+        tables,
+        "preprocess_ts",
+        minimum_gap=minimum_gap,
+        remove_telomeres=remove_telomeres,
+        filter_populations=filter_populations,
+        filter_individuals=filter_individuals,
+        filter_sites=filter_sites,
+        delete_intervals=delete_intervals,
+    )
     return tables.tree_sequence()