Truncate on the basis of whether the nodes are fixed

Author: hyanwong

tsdate/base.py

@@ -95,6 +95,12 @@ def __init__(
         ] = (-np.arange(num_nodes - self.num_nonfixed) - 1)
         self.probability_space = LIN
 
+    def fixed_node_ids(self):
+        return np.where(self.row_lookup < 0)[0]
+
+    def nonfixed_node_ids(self):
+        return np.where(self.row_lookup >= 0)[0]
+
     def force_probability_space(self, probability_space):
         """
         probability_space can be "logarithmic" or "linear": this function will force

tsdate/prior.py

@@ -996,38 +996,42 @@ 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, progress=False):
     """
-    Truncate priors so they conform to the age of nodes in the tree sequence
+    Truncate priors for the nodes listed in truncate_nodes (or all nonfixed nodes
+    if truncate_nodes in None) so they conform to the age of fixed 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])]
+    truncate_nodes = priors.nonfixed_node_ids()
+    # ensure truncate_nodes is ordered by node time
+    truncate_nodes = truncate_nodes[np.argsort(tables.nodes.time[truncate_nodes])]
+
+    fixed_nodes = priors.fixed_node_ids()
+    fixed_times = tables.nodes.time[fixed_nodes]
 
     grid_data = np.copy(priors.grid_data[:])
     timepoints = priors.timepoints
-    if np.max(tables.nodes.time[ts.samples()]) >= np.max(timepoints):
-        raise ValueError("Sample times cannot be larger than the oldest timepoint")
+    if np.max(fixed_times) >= np.max(timepoints):
+        raise ValueError("Fixed node times cannot be older than the oldest timepoint")
     if priors.probability_space == "linear":
         zero_value = 0
     elif priors.probability_space == "logarithmic":
         zero_value = -np.inf
     constrained_min_times = np.zeros_like(tables.nodes.time)
-    constrained_min_times[ts.samples()] = tables.nodes.time[ts.samples()]
+    # Set the min times of fixed nodes to those in the tree sequence
+    constrained_min_times[fixed_nodes] = fixed_times
     constrained_max_times = np.full_like(constrained_min_times, np.inf)
 
     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)]
+    parent_indices = parents_unique[1][np.isin(parents_unique[0], truncate_nodes)]
     for index, nd in tqdm(
-        enumerate(nodes_to_date), desc="Constrain Ages", disable=not progress
+        enumerate(truncate_nodes), desc="Constrain Ages", disable=not progress
     ):
-        if index + 1 != len(nodes_to_date):
+        if index + 1 != len(truncate_nodes):
             children_index = np.arange(parent_indices[index], parent_indices[index + 1])
         else:
             children_index = np.arange(parent_indices[index], ts.num_edges)