Refactor into PopulationSizeHistory class

Author: nspope

tests/test_accuracy.py

@@ -170,8 +170,7 @@ def test_piecewise_scaling(self, bkwd_rate, trio_tmrca):
         time = np.linspace(0, 10, 100)
         ne = 0.5 * np.exp(bkwd_rate * time)
         ts = tskit.Tree.generate_comb(3).tree_sequence
-        dts = tsdate.date(
-            ts, population_size=np.column_stack([time, ne]), mutation_rate=None
-        )
+        demo = tsdate.demography.PopulationSizeHistory(ne, time[1:])
+        dts = tsdate.date(ts, population_size=demo, mutation_rate=None)
         # Check the date is within 10% of the expected
         assert 0.9 < dts.node(dts.first().root).time / trio_tmrca < 1.1

tests/test_functions.py

@@ -45,12 +45,12 @@
 from tsdate.core import Likelihoods
 from tsdate.core import LogLikelihoods
 from tsdate.core import posterior_mean_var
+from tsdate.demography import PopulationSizeHistory
 from tsdate.prior import ConditionalCoalescentTimes
 from tsdate.prior import fill_priors
 from tsdate.prior import gamma_approx
 from tsdate.prior import PriorParams
 from tsdate.prior import SpansBySamples
-from tsdate.util import change_time_measure
 from tsdate.util import nodes_time_unconstrained
 
 
@@ -512,7 +512,7 @@ def test_two_tree_mutation_ts_intervals(self):
 class TestPriorVals:
     def verify_prior_vals(self, ts, prior_distr, **kwargs):
         span_data = SpansBySamples(ts, **kwargs)
-        Ne = np.array([[0, 0.5]])
+        Ne = PopulationSizeHistory(0.5)
         priors = ConditionalCoalescentTimes(None, prior_distr=prior_distr)
         priors.add(ts.num_samples, approximate=False)
         grid = np.linspace(0, 3, 3)
@@ -1102,7 +1102,7 @@ def run_outside_algorithm(
         self, ts, prior_distr="lognorm", standardize=False, ignore_oldest_root=False
     ):
         span_data = SpansBySamples(ts)
-        Ne = np.array([[0, 0.5]])
+        Ne = PopulationSizeHistory(0.5)
         priors = ConditionalCoalescentTimes(None, prior_distr)
         priors.add(ts.num_samples, approximate=False)
         grid = np.array([0, 1.2, 2])
@@ -1205,7 +1205,7 @@ class TestTotalFunctionalValueTree:
     def find_posterior(self, ts, prior_distr):
         grid = np.array([0, 1.2, 2])
         span_data = SpansBySamples(ts)
-        Ne = np.array([[0, 0.5]])
+        Ne = PopulationSizeHistory(0.5)
         priors = ConditionalCoalescentTimes(None, prior_distr=prior_distr)
         priors.add(ts.num_samples, approximate=False)
         mixture_priors = priors.get_mixture_prior_params(span_data)
@@ -1269,7 +1269,7 @@ def test_gil_tree(self):
             ts = utility_functions.gils_example_tree()
             span_data = SpansBySamples(ts)
             prior_distr = "lognorm"
-            Ne = np.array([[0, 0.5]])
+            Ne = PopulationSizeHistory(0.5)
             priors = ConditionalCoalescentTimes(None, prior_distr=prior_distr)
             priors.add(ts.num_samples, approximate=False)
             grid = np.array([0, 0.1, 0.2, 0.5, 1, 2, 5])
@@ -1978,45 +1978,95 @@ def test_historical_samples(self):
         )
 
 
-class TestRescaleTime:
-    def test_rescale(self):
+class TestPopulationSizeHistory:
+    def test_change_time_measure_scalar(self):
         Ne = 10000
         coaltime = np.array([0, 1, 2, 3])
         coalstart = np.array([0])
         coalrate = np.array([1 / (2 * Ne)])
-        gens, _, _ = change_time_measure(coaltime, coalstart, coalrate)
+        gens, _, _ = PopulationSizeHistory._change_time_measure(
+            coaltime, coalstart, coalrate
+        )
         assert np.allclose(gens, 2 * coaltime * Ne)
 
-    def test_piecewise(self):
+    def test_change_time_measure_piecewise(self):
         Ne = np.array([2000, 3000, 5000])
         start = np.array([0, 4000, 10000])
         gens = np.array([2000, 7000, 15000])
-        coaltime, coalstart, coalrate = change_time_measure(gens, start, 2 * Ne)
+        coaltime, coalstart, coalrate = PopulationSizeHistory._change_time_measure(
+            gens, start, 2 * Ne
+        )
         assert np.allclose(coalstart, np.array([0, 1, 2]))
         assert np.allclose(coaltime, np.array([0.5, 1.5, 2.5]))
         assert np.allclose(coalrate, 1 / (2 * Ne))
 
-    def test_piecewise_bijection(self):
+    def test_change_time_measure_bijection(self):
         hapNe = np.array([2000, 3000, 5000])
         start = np.array([0, 4000, 10000])
         gens = np.array([500, 7000, 15000])
-        coaltime, coalstart, coalrate = change_time_measure(gens, start, hapNe)
-        gens_back, start_back, hapNe_back = change_time_measure(
+        coaltime, coalstart, coalrate = PopulationSizeHistory._change_time_measure(
+            gens, start, hapNe
+        )
+        gens_back, start_back, hapNe_back = PopulationSizeHistory._change_time_measure(
             coaltime, coalstart, coalrate
         )
         assert np.allclose(gens, gens_back)
         assert np.allclose(start, start_back)
         assert np.allclose(hapNe, hapNe_back)
 
-    def test_piecewise_numerically(self):
+    def test_change_time_measure_numerically(self):
         coalrate = np.array([0.001, 0.01, 0.1])
         coalstart = np.array([0, 1, 2])
         coaltime = np.linspace(0, 3, 10)
-        gens, _, _ = change_time_measure(coaltime, coalstart, coalrate)
+        gens, _, _ = PopulationSizeHistory._change_time_measure(
+            coaltime, coalstart, coalrate
+        )
         for i in range(gens.size):
             x, _ = scipy.integrate.quad(
                 lambda t: 1 / coalrate[np.digitize(t, coalstart) - 1],
                 a=0,
                 b=coaltime[i],
             )
             assert np.isclose(x, gens[i])
+
+    def test_to_coalescent_timescale(self):
+        demography = PopulationSizeHistory(
+            np.array([1000, 2000, 3000]), np.array([500, 2500])
+        )
+        coaltime = demography.to_coalescent_timescale(np.array([250, 1500]))
+        assert np.allclose(coaltime, [0.125, 0.5])
+
+    def test_to_natural_timescale(self):
+        demography = PopulationSizeHistory(
+            np.array([1000, 2000, 3000]), np.array([500, 2500])
+        )
+        time = demography.to_natural_timescale(np.array([0.125, 0.5]))
+        assert np.allclose(time, [250, 1500])
+
+    def test_single_epoch(self):
+        for Ne in [10000, np.array([10000])]:
+            demography = PopulationSizeHistory(Ne)
+            time = demography.to_natural_timescale(np.array([0, 1, 2, 3]))
+            assert np.allclose(time, [0.0, 20000, 40000, 60000])
+
+    def test_bad_arguments(self):
+        with pytest.raises(ValueError, match="a numpy array"):
+            PopulationSizeHistory([1])
+        with pytest.raises(ValueError, match="a numpy array"):
+            PopulationSizeHistory(np.array([1, 1]), [1])
+        with pytest.raises(ValueError, match="must be greater than 0"):
+            PopulationSizeHistory(0)
+        with pytest.raises(ValueError, match="must be greater than 0"):
+            PopulationSizeHistory(np.array([0, 0]), np.array([1]))
+        with pytest.raises(ValueError, match="must be greater than 0"):
+            PopulationSizeHistory(np.array([1, 1]), np.array([0]))
+        with pytest.raises(ValueError, match="one less than the number"):
+            PopulationSizeHistory(np.array([1]), np.array([1]))
+        with pytest.raises(ValueError, match="increasing order"):
+            PopulationSizeHistory(np.array([1, 1, 1]), np.array([2, 1]))
+        demography = PopulationSizeHistory(1)
+        for time in [1, [1]]:
+            with pytest.raises(ValueError, match="a numpy array"):
+                demography.to_natural_timescale(time)
+            with pytest.raises(ValueError, match="a numpy array"):
+                demography.to_coalescent_timescale(time)

tests/test_inference.py

@@ -55,24 +55,9 @@ def test_not_needed_population_size(self):
 
     def test_bad_population_size(self):
         ts = utility_functions.two_tree_mutation_ts()
-        with pytest.raises(ValueError, match="greater than 0"):
-            tsdate.date(ts, mutation_rate=None, population_size=0)
-        with pytest.raises(ValueError, match="greater than 0"):
-            tsdate.date(ts, mutation_rate=None, population_size=-1)
-        with pytest.raises(ValueError, match="two-dimensional"):
-            tsdate.date(ts, mutation_rate=None, population_size=np.array([[[1]]]))
-        with pytest.raises(ValueError, match="two columns"):
-            tsdate.date(ts, mutation_rate=None, population_size=np.array([[1]]))
-        with pytest.raises(ValueError, match="nonnegative"):
-            tsdate.date(ts, mutation_rate=None, population_size=np.array([[-1, 1]]))
-        with pytest.raises(ValueError, match="positive"):
-            tsdate.date(ts, mutation_rate=None, population_size=np.array([[0, 0]]))
-        with pytest.raises(ValueError, match="start at time 0"):
-            tsdate.date(ts, mutation_rate=None, population_size=np.array([[1, 1]]))
-        with pytest.raises(ValueError, match="unique and increasing"):
-            tsdate.date(
-                ts, mutation_rate=None, population_size=np.array([[0, 1], [0, 1]])
-            )
+        for Ne in [0, -1]:
+            with pytest.raises(ValueError, match="greater than 0"):
+                tsdate.date(ts, mutation_rate=None, population_size=Ne)
 
     def test_dangling_failure(self):
         ts = utility_functions.single_tree_ts_n2_dangling()

tsdate/core.py

@@ -1032,16 +1032,11 @@ def date(
     # Remove any times associated with mutations
     tables.mutations.time = np.full(tree_sequence.num_mutations, tskit.UNKNOWN_TIME)
     tables.sort()
-    population_size_provenance = (
-        population_size.tolist()
-        if isinstance(population_size, np.ndarray)
-        else population_size
-    )
+    # TODO: record population_size provenance, or record that it is omitted
     provenance.record_provenance(
         tables,
         "date",
         mutation_rate=mutation_rate,
-        population_size=population_size_provenance,
         recombination_rate=recombination_rate,
         progress=progress,
         **kwargs,

tsdate/demography.py

@@ -0,0 +1,152 @@
+# MIT License
+#
+# Copyright (c) 2020 University of Oxford
+# Copyright (c) 2021-2023 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.
+"""
+Routines and classes for manipulating demographic histories in tsdate
+"""
+import numpy as np
+
+
+class PopulationSizeHistory:
+    """
+    Stores a piecewise constant population size history and tranforms time from
+    a natural (generational) scale to a coalescent one
+    """
+
+    @staticmethod
+    def _change_time_measure(time_ago, breakpoints, time_measure):
+        """
+        Rescales time given a piecewise-constant time measure. To convert from
+        generations to coalescent units, the time measure per epoch should be 2 *
+        effective population size.  To convert from coalescent units to
+        generations, the time measure should be the coalescent rate ``1/(2 * Ne)``.
+
+        :param np.ndarray time_ago: An increasing vector of time points
+        :param np.ndarray breakpoints: Start times of pieces
+        :param np.ndarray time_measure: Time measure within pieces
+
+        :return: Inputs in new time measure
+        """
+
+        assert np.all(np.diff(breakpoints) > 0.0)
+        assert np.min(breakpoints) == 0.0
+        assert np.all(time_ago >= 0.0)
+        assert np.all(time_measure > 0.0)
+        assert breakpoints.size == time_measure.size
+        index = np.searchsorted(breakpoints, time_ago, side="right") - 1
+        step = np.concatenate(
+            [
+                [0.0],
+                np.cumsum(
+                    breakpoints[1:] * (1.0 / time_measure[:-1] - 1.0 / time_measure[1:])
+                ),
+            ]
+        )
+        new_time_ago = time_ago * 1.0 / time_measure[index] + step[index]
+        new_breakpoints = breakpoints * 1.0 / time_measure + step
+        new_time_measure = 1.0 / time_measure
+        return new_time_ago, new_breakpoints, new_time_measure
+
+    def __init__(self, population_size, time_breaks=None):
+        """
+        :param np.ndarray population_size: A numpy array containing diploid
+            population sizes per epoch
+        :param np.ndarray time_breaks: A sorted numpy array containing time
+            breaks that divide epochs, measured in units of generations in the
+            past
+        """
+
+        if time_breaks is None:
+            time_breaks = np.array([], dtype=float)
+
+        if isinstance(population_size, (int, float)):
+            if not population_size > 0:
+                raise ValueError("Population size must be greater than 0")
+            population_size = np.array([population_size], dtype=float)
+        else:
+            if not isinstance(population_size, np.ndarray):
+                raise ValueError("Population sizes must be in a numpy array")
+            if not np.all(population_size > 0.0):
+                raise ValueError("Population sizes must be greater than 0")
+            if not isinstance(time_breaks, np.ndarray):
+                raise ValueError("Epoch time breaks must be in a numpy array")
+            if not time_breaks.size == population_size.size - 1:
+                raise ValueError(
+                    "The length of the population size array must be one less "
+                    "than the number of epoch time breaks"
+                )
+            if time_breaks.size > 0:
+                if not np.all(time_breaks > 0.0):
+                    raise ValueError("Epoch time breaks must be greater than 0")
+                if not np.all(np.diff(time_breaks) > 0.0):
+                    raise ValueError(
+                        "Epoch time breaks must be unique and in increasing order"
+                    )
+
+        self.time_breaks = np.append([0.0], time_breaks.flatten())
+        self.population_size = 2 * population_size.flatten()
+        _, coalescent_breaks, coalescent_rate = self._change_time_measure(
+            self.time_breaks, self.time_breaks, self.population_size
+        )
+        self.coalescent_breaks = coalescent_breaks
+        self.coalescent_rate = coalescent_rate
+
+    def to_natural_timescale(self, coalescent_time_ago):
+        """
+        Convert a vector of times from coalescent units to generations
+
+        :param np.ndarray coalescent_time_ago: Times in the past, in coalescent units
+        :return: Times in the past, in generations
+        """
+
+        if not isinstance(coalescent_time_ago, np.ndarray):
+            raise ValueError("Times must be in a numpy array")
+        time_ago, _, _ = self._change_time_measure(
+            coalescent_time_ago,
+            self.coalescent_breaks,
+            self.coalescent_rate,
+        )
+        return time_ago
+
+    def to_coalescent_timescale(self, time_ago):
+        """
+        Convert a vector of times from generations to coalescent units
+
+        :param np.ndarray time_ago: Times in the past, in generations
+        :return: Times in the past, in coalescent units
+        """
+
+        if not isinstance(time_ago, np.ndarray):
+            raise ValueError("Times must be in a numpy array")
+        coalescent_time_ago, _, _ = self._change_time_measure(
+            time_ago,
+            self.time_breaks,
+            self.population_size,
+        )
+        return coalescent_time_ago
+
+    # TODO:
+    # @staticmethod
+    # def from_demes(filename):
+    #    """
+    #    Create a `PopulationSizeHistory` instance from a `demes` format YAML
+    #    """