Merge pull request #599 from hyanwong/change-default-ts

Change time defaults for from_tree_sequence

Author: mergify[bot]

CHANGELOG.rst

@@ -16,6 +16,10 @@
 - Oldest nodes in a standard inferred tree sequence are no longer set to frequencies ~2
   and ~3 (i.e. 2 or 3 times as old as all the other nodes), but are spaced above the
   others by the mean time between unique ancestor ages (:pr:`485`, :user:`hyanwong`)
+  
+- The ``tsinfer.SampleData.from_tree_sequence()`` function now defaults to setting
+  ``use_sites_time`` and ``use_individuals_time`` to ``False`` rather than ``True``
+  (:pr:`599`, :user:`hyanwong`)
 
 ********************
 [0.2.1] - 2021-05-26

tests/test_formats.py

@@ -366,7 +366,7 @@ def test_from_tree_sequence_simple(self):
         ts = tsutil.get_example_ts(10, 10, 1)
         sd1 = formats.SampleData(sequence_length=ts.sequence_length)
         self.verify_data_round_trip(ts, sd1)
-        sd2 = formats.SampleData.from_tree_sequence(ts)
+        sd2 = formats.SampleData.from_tree_sequence(ts, use_sites_time=True)
         assert sd1.data_equal(sd2)
 
     def test_from_tree_sequence_variable_allele_number(self):
@@ -398,7 +398,7 @@ def test_from_tree_sequence_variable_allele_number(self):
         num_alleles = sd1.num_alleles()
         for var in ts.variants():
             assert len(var.alleles) == num_alleles[var.site.id]
-        sd2 = formats.SampleData.from_tree_sequence(ts)
+        sd2 = formats.SampleData.from_tree_sequence(ts, use_sites_time=True)
         assert sd1.data_equal(sd2)
 
     def test_from_tree_sequence_with_metadata(self):
@@ -412,34 +412,63 @@ def test_from_tree_sequence_with_metadata(self):
         ts_no_individuals = tables.tree_sequence()
         sd1 = formats.SampleData(sequence_length=ts.sequence_length)
         self.verify_data_round_trip(ts_no_individuals, sd1)
-        sd2 = formats.SampleData.from_tree_sequence(ts_no_individuals)
+        sd2 = formats.SampleData.from_tree_sequence(
+            ts_no_individuals, use_sites_time=True
+        )
         assert sd1.data_equal(sd2)
 
     def test_from_tree_sequence_with_metadata_and_individuals(self):
         ts = tsutil.get_example_individuals_ts_with_metadata(5, 3, 10)
         sd1 = formats.SampleData(sequence_length=ts.sequence_length)
         self.verify_data_round_trip(ts, sd1)
-        sd2 = formats.SampleData.from_tree_sequence(ts)
+        sd2 = formats.SampleData.from_tree_sequence(ts, use_sites_time=True)
         sd1.assert_data_equal(sd2)
 
-    def test_from_historical_tree_sequence(self):
+    def test_from_historical_tree_sequence_with_times(self):
+        n_indiv = 5
         ploidy = 2
-        individual_times = np.arange(5)
+        individual_times = np.arange(n_indiv)
         ts = tsutil.get_example_historical_sampled_ts(individual_times, ploidy, 10)
+        # Test on a tree seq containing an individual with no nodes
+        keep_samples = [u for i in ts.individuals() for u in i.nodes if i.id < n_indiv]
+        ts = ts.simplify(samples=keep_samples, filter_individuals=False)
         sd1 = formats.SampleData(sequence_length=ts.sequence_length)
         self.verify_data_round_trip(ts, sd1)
-        sd2 = formats.SampleData.from_tree_sequence(ts)
+        sd2 = formats.SampleData.from_tree_sequence(
+            ts, use_sites_time=True, use_individuals_time=True
+        )
         assert sd1.data_equal(sd2)
+        # Fails if use_individuals_time is not set
+        sd2 = formats.SampleData.from_tree_sequence(ts, use_sites_time=True)
+        assert not sd1.data_equal(sd2)
 
-    def test_from_tree_sequence_use_time(self):
+    def test_from_tree_sequence_no_times(self):
+        n_indiv = 5
         ploidy = 2
-        individual_times = np.arange(5)
+        individual_times = np.arange(n_indiv + 1)
         ts = tsutil.get_example_historical_sampled_ts(individual_times, ploidy, 10)
-        sd1 = formats.SampleData.from_tree_sequence(ts, use_individuals_time=False)
+        # Test on a tree seq containing an individual with no nodes
+        keep_samples = [u for i in ts.individuals() for u in i.nodes if i.id < n_indiv]
+        ts = ts.simplify(samples=keep_samples, filter_individuals=False)
+        sd1 = formats.SampleData.from_tree_sequence(ts)
+        assert sd1.num_individuals == n_indiv
         assert np.all(sd1.individuals_time[:] == 0)
-        sd2 = formats.SampleData.from_tree_sequence(ts, use_sites_time=False)
-        assert np.all(tskit.is_unknown_time(sd2.sites_time[:]))
-        assert np.array_equal(sd2.individuals_time[:], individual_times)
+
+    def test_from_tree_sequence_time_incompatibilities(self):
+        ploidy = 2
+        individual_times = np.arange(5)
+        ts = tsutil.get_example_historical_sampled_ts(individual_times, ploidy, 10)
+        with pytest.raises(ValueError, match="Incompatible timescales"):
+            _ = formats.SampleData.from_tree_sequence(ts, use_individuals_time=True)
+        # Similar error if no individuals in the TS
+        tables = ts.dump_tables()
+        tables.individuals.clear()
+        tables.nodes.individual = np.full(
+            tables.nodes.num_rows, tskit.NULL, dtype=tables.nodes.individual.dtype
+        )
+        ts = tables.tree_sequence()
+        with pytest.raises(ValueError, match="Incompatible timescales"):
+            _ = formats.SampleData.from_tree_sequence(ts, use_individuals_time=True)
 
     def test_chunk_size(self):
         ts = tsutil.get_example_ts(4, 2)
@@ -802,7 +831,7 @@ def test_sites(self):
     def test_sites_subset(self):
         ts = tsutil.get_example_ts(11, 15)
         assert ts.num_sites > 1
-        input_file = formats.SampleData.from_tree_sequence(ts)
+        input_file = formats.SampleData.from_tree_sequence(ts, use_sites_time=True)
         assert list(input_file.sites([])) == []
         index = np.arange(input_file.num_sites)
         site_list = list(input_file.sites())
@@ -1587,7 +1616,7 @@ def verify(self, sd1, sd2):
     def test_merge_identical(self):
         n = 10
         ts = tsutil.get_example_ts(n, 10, 1)
-        sd1 = formats.SampleData.from_tree_sequence(ts)
+        sd1 = formats.SampleData.from_tree_sequence(ts, use_sites_time=True)
         sd2 = sd1.merge(sd1)
         assert sd2.num_sites == sd1.num_sites
         assert sd2.num_samples == 2 * sd1.num_samples
@@ -1721,7 +1750,7 @@ class TestMinSiteTimes:
 
     def test_no_historical(self):
         ts = tsutil.get_example_ts(10, 10, 1)
-        sd1 = formats.SampleData.from_tree_sequence(ts)
+        sd1 = formats.SampleData.from_tree_sequence(ts, use_sites_time=True)
         # No arguments and individuals_only=True should give array of zeros
         bounds_individuals_only = sd1.min_site_times(individuals_only=True)
         assert np.array_equal(bounds_individuals_only, np.zeros(sd1.num_sites))
@@ -1731,7 +1760,9 @@ def test_no_historical(self):
     def test_simple_case(self):
         individual_times = [0, 0, 0.5, 1]
         ts = tsutil.get_example_historical_sampled_ts(individual_times, ploidy=1)
-        sd1 = formats.SampleData.from_tree_sequence(ts)
+        sd1 = formats.SampleData.from_tree_sequence(
+            ts, use_sites_time=True, use_individuals_time=True
+        )
         time_bound_individuals_only = sd1.min_site_times(individuals_only=True)
         # Because this is a haploid tree sequence we can use the
         # individual and sample IDs interchangably.

tests/test_inference.py

@@ -1116,7 +1116,11 @@ def test_from_historical_tree_sequence(self):
         seq_len = 10
         individual_times = np.arange(n_indiv)
         ts = tsutil.get_example_historical_sampled_ts(individual_times, ploidy, seq_len)
-        ts_inferred = tsinfer.infer(tsinfer.SampleData.from_tree_sequence(ts))
+        ts_inferred = tsinfer.infer(
+            tsinfer.SampleData.from_tree_sequence(
+                ts, use_sites_time=True, use_individuals_time=True
+            )
+        )
         assert ts.sequence_length == ts_inferred.sequence_length
         assert ts.metadata_schema == ts_inferred.metadata_schema
         assert ts.metadata == ts_inferred.metadata

tsinfer/formats.py

@@ -1419,8 +1419,8 @@ def min_site_times(self, individuals_only=False):
     def from_tree_sequence(
         cls,
         ts,
-        use_sites_time=True,
-        use_individuals_time=True,
+        use_sites_time=None,
+        use_individuals_time=None,
         **kwargs,
     ):
         """
@@ -1444,16 +1444,18 @@ def from_tree_sequence(
 
         :param TreeSequence ts: The :class:`tskit.TreeSequence` from which to generate
             samples.
-        :param bool use_sites_time: If True (default), the times of nodes in the tree
+        :param bool use_sites_time: If ``True``, the times of nodes in the tree
             sequence are used to set a time for each site (which affects the relative
             temporal order of ancestors during inference). Times for a site are only
             used if there is a single mutation at that site, in which case the node
             immediately below the mutation is taken as the origination time for the
-            variant. If False, the frequency of the variant is used as a proxy for the
-            relative variant time (see :meth:`.add_site`).
-        :param bool use_individuals_time: If True (default), set a time for individuals
-            that contain historical sample nodes. If False, all individuals are
-            set at time 0.
+            variant. If ``False``, the frequency of the variant is used as a proxy for
+            the relative variant time (see :meth:`.add_site`). Defaults to ``False``.
+        :param bool use_individuals_time: If ``True``, use the time of the sample nodes
+            in the tree sequence as the time of the individuals associated with
+            those nodes in the sample data file. This is likely only to be meaningful if
+            ``use_sites_time`` is also ``True``. If ``False``, all individuals are set
+            to time 0. Defaults to ``False``.
         :param \\**kwargs: Further arguments passed to the :class:`SampleData`
             constructor.
         :return: A :class:`.SampleData` object.
@@ -1468,6 +1470,11 @@ def encode_metadata(metadata, schema):
                     metadata = None
             return metadata
 
+        if use_sites_time is None:
+            use_sites_time = False
+        if use_individuals_time is None:
+            use_individuals_time = False
+
         tables = ts.tables
         self = cls(sequence_length=ts.sequence_length, **kwargs)
 
@@ -1488,6 +1495,7 @@ def encode_metadata(metadata, schema):
         for individual in ts.individuals():
             nodes = individual.nodes
             if len(nodes) > 0:
+                time = 0
                 first_node = ts.node(nodes[0])
                 for u in nodes[1:]:
                     if ts.node(u).time != first_node.time:
@@ -1502,18 +1510,37 @@ def encode_metadata(metadata, schema):
                             "population".format(individual.id)
                         )
                 metadata = encode_metadata(individual.metadata, schema)
+                if use_individuals_time:
+                    time = first_node.time
+                    if time != 0 and not use_sites_time:
+                        raise ValueError(
+                            "Incompatible timescales: site frequencies used for times "
+                            f"(use_sites_time=False), but node {first_node.id} in "
+                            f"individual {individual.id} has a nonzero time and "
+                            "use_individuals_time=True. Please set site times manually."
+                        )
                 self.add_individual(
                     location=individual.location,
                     metadata=metadata,
                     population=first_node.population,
                     flags=individual.flags,
-                    time=first_node.time if use_individuals_time else 0,
+                    time=time,
                     ploidy=len(nodes),
                 )
         for u in ts.samples():
             node = ts.node(u)
             if node.individual == tskit.NULL:
                 # The sample node has no individual: create a haploid individual for it
+                time = 0
+                if use_individuals_time:
+                    time = node.time
+                    if time != 0 and not use_sites_time:
+                        raise ValueError(
+                            "Incompatible timescales: site frequencies used for times "
+                            f"(use_sites_time=False), but node {node.id} "
+                            "has a nonzero time and use_individuals_time=True. "
+                            "Please set site times manually."
+                        )
                 self.add_individual(
                     population=node.population,
                     time=node.time if use_individuals_time else 0,