Change Ne to population size in API

Author: nspope

.gitignore

@@ -9,3 +9,5 @@ tsdate.egg-info
 data/prior_1000df.txt
 *.ipynb
 .ipynb*
+build/
+.*.swp

tsdate/cli.py

@@ -80,8 +80,11 @@ def tsdate_cli_parser():
         help="The path and name of output file where the dated tree \
                         sequence will saved.",
     )
+    # TODO array specification from file?
     parser.add_argument(
-        "Ne", type=float, help="estimated effective (diploid) population size."
+        "population_size",
+        type=float,
+        help="estimated effective (diploid) population size.",
     )
     parser.add_argument(
         "-m",
@@ -190,7 +193,7 @@ def run_date(args):
     dated_ts = tsdate.date(
         ts,
         args.mutation_rate,
-        args.Ne,
+        args.population_size,
         recombination_rate=args.recombination_rate,
         probability_space=args.probability_space,
         method=args.method,

tsdate/core.py

@@ -932,7 +932,7 @@ def constrain_ages_topo(ts, post_mn, eps, nodes_to_date=None, progress=False):
 def date(
     tree_sequence,
     mutation_rate,
-    Ne=None,
+    population_size=None,
     recombination_rate=None,
     time_units=None,
     priors=None,
@@ -966,10 +966,12 @@ def date(
 
     :param TreeSequence tree_sequence: The input :class:`tskit.TreeSequence`, treated as
         one whose non-sample nodes are undated.
-    :param float Ne: The estimated (diploid) effective population size used to construct
-        the (default) conditional coalescent prior. This is what is used when ``priors``
-        is ``None``: a positive ``Ne`` value is therefore required in this case.
-        Conversely, if ``priors`` is not ``None``, no ``Ne`` value should be given.
+    :param float population_size: The estimated (diploid) effective population
+        size used to construct the (default) conditional coalescent prior. This may
+        be a single value, or a two-column array of epoch breakpoints and effective
+        population sizes within epochs. These are used when ``priors`` is ``None``.
+        Conversely, if ``priors`` is not ``None``, no ``population_size`` value
+        should be given.
     :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``
@@ -979,12 +981,12 @@ def date(
     :param str time_units: The time units used by the ``mutation_rate`` and
         ``recombination_rate`` values, and stored in the ``time_units`` attribute of the
         output tree sequence. If the conditional coalescent prior is used,
-        then this is also applies to the value of ``Ne``, which in standard coalescent
-        theory is measured in generations. Therefore if you wish to use mutation and
-        recombination rates measured in (say) years, and are using the conditional
-        coalescent prior, the ``Ne`` value which you provide must be scaled by
-        multiplying by the number of years per generation. If ``None`` (default), assume
-        ``"generations"``.
+        then this is also applies to the value of ``population_size``, which in
+        standard coalescent theory is measured in generations. Therefore if you
+        wish to use mutation and recombination rates measured in (say) years,
+        and are using the conditional coalescent prior, the ``population_size``
+        value which you provide must be scaled by multiplying by the number of
+        years per generation. If ``None`` (default), assume ``"generations"``.
     :param NodeGridValues priors: NodeGridValue object containing the prior probabilities
         for each node at a set of discrete time points. If ``None`` (default), use the
         conditional coalescent prior with a standard set of time points as given by
@@ -1016,7 +1018,7 @@ def date(
         time_units = "generations"
     tree_sequence, dates, posteriors, timepoints, eps, nds = get_dates(
         tree_sequence,
-        Ne=Ne,
+        population_size=population_size,
         mutation_rate=mutation_rate,
         recombination_rate=recombination_rate,
         priors=priors,
@@ -1034,7 +1036,7 @@ def date(
         tables,
         "date",
         mutation_rate=mutation_rate,
-        Ne=Ne,
+        population_size=population_size,
         recombination_rate=recombination_rate,
         progress=progress,
         **kwargs,
@@ -1051,7 +1053,7 @@ def date(
 def get_dates(
     tree_sequence,
     mutation_rate,
-    Ne=None,
+    population_size=None,
     recombination_rate=None,
     priors=None,
     *,
@@ -1087,24 +1089,24 @@ def get_dates(
         approx_priors = True
 
     if priors is None:
-        if Ne is None:
+        if population_size is None:
             raise ValueError(
-                "Must specify Ne if priors are not already built using \
-                        tsdate.build_prior_grid()"
+                "Must specify population_size if priors are not already built \
+                        using tsdate.build_prior_grid()"
             )
         priors = prior.build_grid(
             tree_sequence,
-            Ne=Ne,
+            population_size=population_size,
             eps=eps,
             progress=progress,
             approximate_priors=approx_priors,
         )
     else:
         logging.info("Using user-specified priors")
-        if Ne is not None:
+        if population_size is not None:
             raise ValueError(
-                "Cannot specify Ne in tsdate.date() or tsdate.get_dates() if \
-                        specifying priors from tsdate.build_prior_grid()"
+                "Cannot specify population_size in tsdate.date() or tsdate.get_dates() \
+                        if specifying priors from tsdate.build_prior_grid()"
             )
         priors = priors
 

tsdate/prior.py

@@ -70,7 +70,11 @@ class ConditionalCoalescentTimes:
     """
 
     def __init__(
-        self, precalc_approximation_n, Ne, prior_distr="lognorm", progress=False
+        self,
+        precalc_approximation_n,
+        population_size,
+        prior_distr="lognorm",
+        progress=False,
     ):
         """
         :param bool precalc_approximation_n: the size of tree used for
@@ -79,7 +83,7 @@ def __init__(
             and therefore do not allow approximate priors to be used
         """
         self.n_approx = precalc_approximation_n
-        self.Ne = Ne
+        self.population_size = population_size
         self.prior_store = {}
         self.progress = progress
 
@@ -177,8 +181,8 @@ def add(self, total_tips, approximate=None):
             priors[1] = PriorParams(alpha=0, beta=1, mean=0, var=0)
         for var, tips in zip(variances, all_tips):
             # NB: it should be possible to vectorize this in numpy
-            var = var * ((2 * self.Ne) ** 2)
-            expectation = self.tau_expect(tips, total_tips) * 2 * self.Ne
+            var = var * ((2 * self.Ne) ** 2)  # TODO
+            expectation = self.tau_expect(tips, total_tips) * 2 * self.Ne  # TODO
             alpha, beta = self.func_approx(expectation, var)
             priors[tips] = PriorParams(
                 alpha=alpha, beta=beta, mean=expectation, var=var
@@ -947,7 +951,9 @@ def gamma_cdf(t_set, alpha, beta):
     return np.insert(t_set, 0, 0)
 
 
-def fill_priors(node_parameters, timepoints, ts, Ne, *, prior_distr, progress=False):
+def fill_priors(
+    node_parameters, timepoints, ts, population_size, *, prior_distr, progress=False
+):
     """
     Take the alpha and beta values from the node_parameters array, which contains
     one row for each node in the TS (including fixed nodes)
@@ -993,7 +999,7 @@ def fill_priors(node_parameters, timepoints, ts, Ne, *, prior_distr, progress=Fa
 
 def build_grid(
     tree_sequence,
-    Ne,
+    population_size,
     timepoints=20,
     *,
     approximate_priors=False,
@@ -1011,9 +1017,11 @@ def build_grid(
 
     :param TreeSequence tree_sequence: The input :class:`tskit.TreeSequence`, treated as
         undated.
-    :param float Ne: The estimated (diploid) effective population size: must be
-        specified. Using standard (unscaled) values for ``Ne`` results in a prior where
-        times are measures in generations.
+    :param float population_size: The estimated (diploid) effective population
+        size: must be specified. May be a single value, or a two-column array with
+        epoch breakpoints and effective population sizes. Using standard (unscaled)
+        values for ``population_size`` results in a prior where times are measures
+        in generations.
     :param int_or_array_like timepoints: The number of quantiles used to create the
         time slices, or manually-specified time slices as a numpy array. Default: 20
     :param bool approximate_priors: Whether to use a precalculated approximate prior or
@@ -1033,8 +1041,9 @@ def build_grid(
         inference and a discretised time grid
     :rtype:  base.NodeGridValues Object
     """
-    if Ne <= 0:
-        raise ValueError("Parameter 'Ne' must be greater than 0")
+    # TODO
+    if population_size <= 0:
+        raise ValueError("Parameter 'population_size' must be greater than 0")
     if approximate_priors:
         if not approx_prior_size:
             approx_prior_size = 1000
@@ -1053,7 +1062,7 @@ def build_grid(
     span_data = SpansBySamples(contmpr_ts, progress=progress, allow_unary=allow_unary)
 
     base_priors = ConditionalCoalescentTimes(
-        approx_prior_size, Ne, prior_distribution, progress=progress
+        approx_prior_size, population_size, prior_distribution, progress=progress
     )
 
     base_priors.add(contmpr_ts.num_samples, approximate_priors)
@@ -1088,7 +1097,7 @@ def build_grid(
         prior_params,
         timepoints,
         tree_sequence,
-        Ne,
+        population_size,
         prior_distr=prior_distribution,
         progress=progress,
     )