Save PopulationSizeHistory params

hyanwong · hyanwong · commit 868e177a3643 · 2023-06-24T09:41:21.000+01:00
Fixes #274
diff --git a/tests/test_functions.py b/tests/test_functions.py
@@ -2147,16 +2147,12 @@ def test_change_time_measure_numerically(self):
             assert np.isclose(x, gens[i])
 
     def test_to_coalescent_timescale(self):
-        demography = PopulationSizeHistory(
-            np.array([1000, 2000, 3000]), np.array([500, 2500])
-        )
+        demography = PopulationSizeHistory([1000, 2000, 3000], [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])
-        )
+        demography = PopulationSizeHistory([1000, 2000, 3000], [500, 2500])
         time = demography.to_natural_timescale(np.array([0.125, 0.5]))
         assert np.allclose(time, [250, 1500])
 
@@ -2169,9 +2165,7 @@ def test_single_epoch(self):
     def test_moments_numerically(self):
         alpha = 2.8
         beta = 1.7
-        demography = PopulationSizeHistory(
-            np.array([1000, 2000, 3000]), np.array([500, 2500])
-        )
+        demography = PopulationSizeHistory([1000, 2000, 3000], [500, 2500])
         numer_mn, _ = scipy.integrate.quad(
             lambda t: demography.to_natural_timescale(np.array([t]))
             * scipy.stats.gamma.pdf(t, alpha, scale=1 / beta),
@@ -2192,10 +2186,16 @@ def test_moments_numerically(self):
         assert np.isclose(numer_va, analy_va)
 
     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="greater than 0"):
+            PopulationSizeHistory([None])
+        with pytest.raises(ValueError, match="finite"):
+            PopulationSizeHistory(np.inf)
+        with pytest.raises(ValueError, match="a numpy float array"):
+            PopulationSizeHistory(["a"])
+        with pytest.raises(TypeError, match="a numpy float array"):
+            PopulationSizeHistory([{}])
+        with pytest.raises(TypeError, match="a numpy float array"):
+            PopulationSizeHistory(np.array([1, 1]), [{}])
         with pytest.raises(ValueError, match="must be greater than 0"):
             PopulationSizeHistory(0)
         with pytest.raises(ValueError, match="must be greater than 0"):
diff --git a/tests/test_provenance.py b/tests/test_provenance.py
@@ -54,26 +54,29 @@ def test_date_params_recorded(self):
         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"
+    @pytest.mark.parametrize(
+        "popdict",
+        [
+            {"population_size": [1, 2, 3], "time_breaks": [1, 1.2]},
+            {"population_size": [123]},
+        ],
     )
-    def test_date_popsizehist_recorded(self):
+    def test_date_popsizehist_recorded(self, popdict):
         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,
-        )
+        for use_class in (False, True):
+            if use_class:
+                popsize = tsdate.demography.PopulationSizeHistory(**popdict)
+            else:
+                popsize = popdict
+            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"]
+            popsz = rec["parameters"]["population_size"]
+            assert len(popsz) == len(popdict)
+            for param, val in popdict.items():
+                assert np.all(np.isclose(val, popsz[param]))
 
     def test_preprocess_cmd_recorded(self):
         ts = utility_functions.ts_w_data_desert(40, 60, 100)
diff --git a/tsdate/core.py b/tsdate/core.py
@@ -39,6 +39,7 @@
 
 from . import approx
 from . import base
+from . import demography
 from . import prior
 from . import provenance
 
@@ -1169,10 +1170,12 @@ def date(
         one whose non-sample nodes are undated.
     :param PopulationSizeHistory population_size: The estimated (diploid) effective
         population size used to construct the (default) conditional coalescent
-        prior. This may be a single value (for a population with constant size), or
-        a :class:`PopulationSizeHistory` object (for a population with time-varying
-        size). This is used when ``priors`` is ``None``.  Conversely, if ``priors``
-        is not ``None``, no ``population_size`` value should be given.
+        prior. For a population with constant size, this can be given as a single
+        value. For a population with time-varying size, this can be given directly as
+        a :class:`PopulationSizeHistory` object or a parameter dictionary passed
+        to initialise a class:`PopulationSizeHistory` object. This is used when
+        ``priors`` is ``None``.  Conversely, if ``priors`` is not ``None``, no
+        ``population_size`` value should be specified.
     :param float mutation_rate: The estimated mutation rate per unit of genome per
         unit time. If provided, the dating algorithm will use a mutation rate clock to
         help estimate node dates. Default: ``None``
@@ -1226,6 +1229,10 @@ def date(
             )
         else:
             population_size = Ne
+
+    if isinstance(population_size, dict):
+        population_size = demography.PopulationSizeHistory(**population_size)
+
     if method == "variational_gamma":
         tree_sequence, dates, posteriors, timepoints, eps, nds = variational_dates(
             tree_sequence,
@@ -1261,8 +1268,8 @@ def date(
     )
     if isinstance(population_size, (int, float)):
         params["population_size"] = population_size
-    else:
-        params["population_size"] = "TODO: PopulationSizeHistory object"
+    elif isinstance(population_size, demography.PopulationSizeHistory):
+        params["population_size"] = population_size.as_dict()
     provenance.record_provenance(
         tables,
         "date",
diff --git a/tsdate/demography.py b/tsdate/demography.py
@@ -69,39 +69,48 @@ def _change_time_measure(time_ago, breakpoints, time_measure):
 
     def __init__(self, population_size, time_breaks=None):
         """
-        :param np.ndarray population_size: A numpy array containing diploid
+        :param array_like population_size: An array containing diploid
             population sizes per epoch
-        :param np.ndarray time_breaks: A sorted numpy array containing time
+        :param array_like time_breaks: A sorted 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)
+            time_breaks = []
 
         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:
+            try:
+                population_size = np.array(population_size, dtype=float)
+            except (ValueError, TypeError) as e:
+                raise e.__class__(
+                    "Population sizes must be convertable to a numpy float array"
+                ) from e
+        if not np.all(population_size > 0.0):
+            raise ValueError("Population sizes must be greater than 0")
+        if not np.all(np.isfinite(population_size)):
+            raise ValueError("Population sizes must be finite")
+
+        try:
+            time_breaks = np.array(time_breaks, dtype=float)
+        except (ValueError, TypeError) as e:
+            raise e.__class__(
+                "Time breaks must be convertable to a numpy float array"
+            ) from e
+        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(
-                    "The length of the population size array must be one less "
-                    "than the number of epoch time breaks"
+                    "Epoch time breaks must be unique and in increasing order"
                 )
-            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()
@@ -111,6 +120,17 @@ def __init__(self, population_size, time_breaks=None):
         self.coalescent_breaks = coalescent_breaks
         self.coalescent_rate = coalescent_rate
 
+    def as_dict(self):
+        """
+        Return the population size history as a dictionary of parameters
+        that can be used to initialise a new object
+        """
+        ret_val = {"population_size": list(self.population_size / 2)}
+        assert self.time_breaks[0] == 0.0
+        if len(self.time_breaks) > 1:
+            ret_val["time_breaks"] = list(self.time_breaks[1:])
+        return ret_val
+
     def to_natural_timescale(self, coalescent_time_ago):
         """
         Convert a vector of times from coalescent units to generations
