Add method to impute unknown mutation times

Author: duncanMR

python/CHANGELOG.rst

@@ -1,3 +1,11 @@
+--------------------
+[0.5.X] - 2023-XX-XX
+--------------------
+
+**Features**
+
+- Add ``TreeSequence.impute_unknown_mutations_time`` method to return an array of mutation times based on the times of associated nodes (:user:`duncanMR`, :pr:`2760`, :issue:`2758`)
+
 --------------------
 [0.5.5] - 2023-05-17
 --------------------

python/tests/test_highlevel.py

@@ -2596,6 +2596,35 @@ def test_arrays_equal_to_tables(self, ts_fixture):
             ts.indexes_edge_removal_order, tables.indexes.edge_removal_order
         )
 
+    @pytest.mark.parametrize("ts", get_example_tree_sequences())
+    def test_impute_unknown_mutations_time(self, ts):
+        # Tests for method='min'
+        imputed_time = ts.impute_unknown_mutations_time(method="min")
+        mutations = ts.tables.mutations
+        nodes_time = ts.nodes_time
+        table_time = np.zeros(len(mutations))
+
+        for mut_idx, mut in enumerate(mutations):
+            if tskit.is_unknown_time(mut.time):
+                node_time = nodes_time[mut.node]
+                table_time[mut_idx] = node_time
+            else:
+                table_time[mut_idx] = mut.time
+
+        assert np.allclose(imputed_time, table_time, rtol=1e-10, atol=1e-10)
+
+        # Check we have valid times
+        tables = ts.dump_tables()
+        tables.mutations.time = imputed_time
+        tables.sort()
+        tables.tree_sequence()
+
+        # Test for unallowed methods
+        with pytest.raises(
+            ValueError, match="Mutations time imputation method must be chosen"
+        ):
+            ts.impute_unknown_mutations_time(method="foobar")
+
 
 class TestSimplify:
     # This class was factored out of the old TestHighlevel class 2022-12-13,

python/tskit/trees.py

@@ -8270,6 +8270,7 @@ def Tajimas_D(self, sample_sets=None, windows=None, mode="site"):
         :return: A ndarray with shape equal to (num windows, num statistics).
             If there is one sample set and windows=None, a numpy scalar is returned.
         """
+
         # TODO this should be done in C as we'll want to support this method there.
         def tjd_func(sample_set_sizes, flattened, **kwargs):
             n = sample_set_sizes
@@ -8882,6 +8883,36 @@ def coalescence_time_distribution(
             span_normalise=span_normalise,
         )
 
+    def impute_unknown_mutations_time(
+        self,
+        method=None,
+    ):
+        """
+        Returns an array of mutation times, where any unknown times are
+        imputed from the times of associated nodes. Not to be confused with
+        :meth:`TableCollection.compute_mutation_times`, which modifies the
+        ``time`` column of the mutations table in place.
+
+        :param str method: The method used to impute the unknown mutation times.
+            Currently only "min" is supported, which uses the time of the node
+            below the mutation as the mutation time. The "min" method can also
+            be specified by ``method=None`` (Default: ``None``).
+        :return: An array of length equal to the number of mutations in the
+            tree sequence.
+        """
+        allowed_methods = ["min"]
+        if method is None:
+            method = "min"
+        if method not in allowed_methods:
+            raise ValueError(
+                f"Mutations time imputation method must be chosen from {allowed_methods}"
+            )
+        if method == "min":
+            mutations_time = self.mutations_time.copy()
+            unknown = tskit.is_unknown_time(mutations_time)
+            mutations_time[unknown] = self.nodes_time[self.mutations_node[unknown]]
+            return mutations_time
+
     ############################################
     #
     # Deprecated APIs. These are either already unsupported, or will be unsupported in a