Allow priors to be truncated or not

And require a flag to use historical samples

Author: hyanwong

tsdate/prior.py

@@ -991,7 +991,7 @@ def fill_priors(node_parameters, timepoints, ts, Ne, *, prior_distr, progress=Fa
     return prior_times
 
 
-def truncate_priors(ts, priors, nodes_to_date=None, progress=False):
+def _truncate_priors(ts, priors, nodes_to_date=None, progress=False):
     """
     Truncate priors so they conform to the age of nodes in the tree sequence
     """
@@ -1043,7 +1043,7 @@ def truncate_priors(ts, priors, nodes_to_date=None, progress=False):
         grid_data = grid_data - rowmax[:, np.newaxis]
 
     priors.grid_data[:] = grid_data
-    return constrained_min_times, constrained_max_times, priors
+    return priors
 
 
 def build_grid(
@@ -1054,6 +1054,8 @@ def build_grid(
     approximate_priors=False,
     approx_prior_size=None,
     prior_distribution="lognorm",
+    allow_historical_samples=None,
+    truncate_priors=None,
     eps=1e-6,
     progress=False,
 ):
@@ -1075,18 +1077,25 @@ def build_grid(
     :param int approx_prior_size: Number of samples from which to precalculate prior.
         Should only enter value if approximate_priors=True. If approximate_priors=True
         and no value specified, defaults to 1000. Default: None
-    :param string prior_distr: What distribution to use to approximate the conditional
-        coalescent prior. Can be "lognorm" for the lognormal distribution (generally a
-        better fit, but slightly slower to calculate) or "gamma" for the gamma
-        distribution (slightly faster, but a poorer fit for recent nodes). Default:
-        "lognorm"
+    :param string prior_distribution: What distribution to use to approximate the
+        conditional coalescent prior. Can be "lognorm" for the lognormal distribution
+        (generally a better fit, but slightly slower to calculate) or "gamma" for the
+        gamma distribution (slightly faster, but a poorer fit for recent nodes).
+        Default: "lognorm"
+    :param bool allow_historical_samples: should we allow historical samples (i.e. at
+        times > 0. This invalidates the assumptions of the conditional coalescent, but
+        may be acceptable if the historical samples are recent or if there are many
+        contemporaneous samples. Default: `False`
+    :param bool truncate_priors: If there are historical samples, should we truncate the
+        priors of their direct ancestor nodes so that the probability of being younger
+        than the oldest descendant sample is zero. If the tree sequence is trustworthy
+        this should give better restults. Default: `True`
     :param float eps: Specify minimum distance separating points in the time grid. Also
         specifies the error factor in time difference calculations. Default: 1e-6
     :return: A prior object to pass to tsdate.date() containing prior values for
         inference and a discretised time grid
     :rtype:  base.NodeGridValues Object
     """
-    # tree_sequence = tree_sequence.simplify(tree_sequence.samples())
 
     if Ne <= 0:
         raise ValueError("Parameter 'Ne' must be greater than 0")
@@ -1098,6 +1107,10 @@ def build_grid(
             raise ValueError(
                 "Can't set approx_prior_size if approximate_prior is False"
             )
+    if truncate_priors is None:
+        truncate_priors = True
+    if allow_historical_samples is None:
+        allow_historical_samples = False
 
     span_data = SpansBySamples(tree_sequence, progress=progress)
 
@@ -1140,5 +1153,13 @@ def build_grid(
         progress=progress,
     )
     if np.any(tree_sequence.tables.nodes.time[tree_sequence.samples()] != 0):
-        _, _, priors = truncate_priors(tree_sequence, priors, progress=progress)
+        if not allow_historical_samples:
+            raise ValueError(
+                "There are samples at non-zero times, invalidating the conditional "
+                "coalescent prior. You can set allow_historical_samples=True to carry "
+                "on regardless, calculating a prior as if all samples were "
+                "contemporaneous (reasonable if you only have a few ancient samples)"
+            )
+        if truncate_priors:
+            priors = _truncate_priors(tree_sequence, priors, progress=progress)
     return priors