Allow truncate_prior to actually function!

Author: hyanwong

tsdate/prior.py

@@ -419,7 +419,7 @@ def __init__(self, tree_sequence, progress=False):
 
         self.ts = tree_sequence
         self.sample_node_set = set(self.ts.samples())
-        #if np.any(self.ts.tables.nodes.time[self.ts.samples()] != 0):
+        # if np.any(self.ts.tables.nodes.time[self.ts.samples()] != 0):
         #    raise ValueError(
         #        "The SpansBySamples class needs a tree seq with all samples at time 0"
         #    )
@@ -991,34 +991,36 @@ def fill_priors(node_parameters, timepoints, ts, Ne, *, prior_distr, progress=Fa
     return prior_times
 
 
-def truncate_priors(ts, sample_times, 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
     """
+    tables = ts.tables
+    if nodes_to_date is None:
+        nodes_to_date = np.arange(ts.num_nodes, dtype=np.uint64)
+        nodes_to_date = nodes_to_date[~np.isin(nodes_to_date, ts.samples())]
+    # ensure nodes_to_date is ordered by node time
+    nodes_to_date = nodes_to_date[np.argsort(tables.nodes.time[nodes_to_date])]
+
     grid_data = np.copy(priors.grid_data[:])
     timepoints = priors.timepoints
-    if np.max(sample_times) >= np.max(timepoints):
+    if np.max(tables.nodes.time[ts.samples()]) >= np.max(timepoints):
         raise ValueError("Sample times cannot be larger than the oldest timepoint")
     if priors.probability_space == "linear":
         zero_value = 0
-        one_value = 1
     elif priors.probability_space == "logarithmic":
         zero_value = -np.inf
-        one_value = 0
-    constrained_min_times = np.copy(sample_times)
-    constrained_max_times = np.full(sample_times.shape[0], np.inf)
-    if nodes_to_date is None:
-        nodes_to_date = np.arange(ts.num_nodes, dtype=np.uint64)
-        nodes_to_date = nodes_to_date[~np.isin(nodes_to_date, ts.samples())]
+    constrained_min_times = np.copy(tables.nodes.time)
+    constrained_min_times[ts.samples()] = tables.nodes.time[ts.samples()]
+    constrained_max_times = np.full_like(constrained_min_times, np.inf)
 
-    tables = ts.tables
     parents = tables.edges.parent
     nd_children = tables.edges.child[np.argsort(parents)]
     parents = sorted(parents)
     parents_unique = np.unique(parents, return_index=True)
     parent_indices = parents_unique[1][np.isin(parents_unique[0], nodes_to_date)]
     for index, nd in tqdm(
-        enumerate(sorted(nodes_to_date)), desc="Constrain Ages", disable=not progress
+        enumerate(nodes_to_date), desc="Constrain Ages", disable=not progress
     ):
         if index + 1 != len(nodes_to_date):
             children_index = np.arange(parent_indices[index], parent_indices[index + 1])
@@ -1033,17 +1035,17 @@ def truncate_priors(ts, sample_times, priors, nodes_to_date=None, progress=False
         lookup_index = priors.row_lookup[int(nd)]
         grid_data[lookup_index][:nearest_time] = zero_value
     assert np.all(constrained_min_times < constrained_max_times)
-    all_zeros = np.where(np.all(grid_data == zero_value, axis=1))[0]
 
     rowmax = grid_data[:, 1:].max(axis=1)
     if priors.probability_space == "linear":
         grid_data = grid_data / rowmax[:, np.newaxis]
     elif priors.probability_space == "logarithmic":
         grid_data = grid_data - rowmax[:, np.newaxis]
-                    
+
     priors.grid_data[:] = grid_data
     return constrained_min_times, constrained_max_times, priors
 
+
 def build_grid(
     tree_sequence,
     Ne,
@@ -1054,7 +1056,6 @@ def build_grid(
     prior_distribution="lognorm",
     eps=1e-6,
     progress=False,
-    sample_times=None
 ):
     """
     Using the conditional coalescent, calculate the prior distribution for the age of
@@ -1085,7 +1086,7 @@ def build_grid(
         inference and a discretised time grid
     :rtype:  base.NodeGridValues Object
     """
-    #tree_sequence = tree_sequence.simplify(tree_sequence.samples())
+    # tree_sequence = tree_sequence.simplify(tree_sequence.samples())
 
     if Ne <= 0:
         raise ValueError("Parameter 'Ne' must be greater than 0")
@@ -1139,6 +1140,5 @@ def build_grid(
         progress=progress,
     )
     if np.any(tree_sequence.tables.nodes.time[tree_sequence.samples()] != 0):
-        if False:
-            priors = truncate_priors(tree_sequence, sample_times, priors, eps, progress=progress)
+        _, _, priors = truncate_priors(tree_sequence, priors, progress=progress)
     return priors