use root_threshold to check for multi-root trees

Author: hyanwong

tests/test_priors.py

@@ -73,6 +73,11 @@ def test_repr(self):
             for u in t.leaves():
                 assert rep.count(f"{u}: {{}}") == 1
 
+    def test_error_on_multiroot(self):
+        ts = utility_functions.multiroot()
+        with pytest.raises(ValueError, match="multiple roots"):
+            SpansBySamples(ts)
+
 
 class TestTimepoints:
     def test_create_timepoints(self):

tests/utility_functions.py

@@ -669,6 +669,10 @@ def two_tree_two_mrcas():
     )
 
 
+def multiroot():
+    return two_tree_two_mrcas().decapitate(1.5)
+
+
 def loopy_tree():
     r"""
     Testing a tree with a loop.

tsdate/core.py

@@ -41,7 +41,6 @@
 from . import base
 from . import prior
 from . import provenance
-from . import util
 
 FORMAT_NAME = "tsdate"
 
@@ -618,12 +617,9 @@ def __init__(self, priors, lik, *, progress=False):
         self.spans = np.pad(self.spans, (0, self.ts.num_nodes - len(self.spans)))
 
         self.root_spans = defaultdict(float)
-        for tree in self.ts.trees():
-            root = util.get_single_root(tree)
-            if root is not None:
-                self.root_spans[
-                    root
-                ] += tree.span  # Count span if we have a single tree
+        for tree in self.ts.trees(root_threshold=2):
+            if tree.has_single_root:
+                self.root_spans[tree.root] += tree.span
         # Add on the spans when this is a root
         for root, span_when_root in self.root_spans.items():
             self.spans[root] += span_when_root

tsdate/prior.py

@@ -130,7 +130,7 @@ def prior_with_max_total_tips(self):
 
     def add(self, total_tips, approximate=None):
         """
-        Create and store a numpy array used to lookup prior paramsn and mean + variance
+        Create and store a numpy array used to lookup prior params and mean + variance
         of ages for nodes with descendant sample tips range from 2..``total_tips``
         given that the total number of tips in the coalescent tree is
         ``total_tips``. The array is indexed by (num_tips / total_tips).
@@ -604,12 +604,13 @@ def save_to_spans(prev_tree, node, num_fixed_at_0_treenodes):
         # Iterate over trees and remaining edge diffs
         focal_tips = list(self.sample_node_set)
         for prev_tree in tqdm(
-            self.ts.trees(tracked_samples=focal_tips),
+            self.ts.trees(tracked_samples=focal_tips, root_threshold=2),
             desc="Find Node Spans",
             total=self.ts.num_trees,
             disable=not self.progress,
         ):
-            util.get_single_root(prev_tree)  # Check only one root
+            if prev_tree.has_multiple_roots:
+                raise ValueError(f"Tree {prev_tree.index} has multiple roots")
             try:
                 # Get the edge diffs from the prev tree to the new tree
                 _, e_out, e_in = next(edge_diff_iter)

tsdate/util.py

@@ -34,16 +34,6 @@
 logger = logging.getLogger(__name__)
 
 
-def get_single_root(tree):
-    # TODO - use new 'root_threshold=2' to avoid having to check isolated nodes
-    topological_roots = [r for r in tree.roots if tree.num_children(r) != 0]
-    if len(topological_roots) > 1:
-        raise ValueError(f"Invalid tree sequence: tree {tree.index} has >1 root")
-    if len(topological_roots) == 0:
-        return None  # Empty tree
-    return topological_roots[0]
-
-
 def reduce_to_contemporaneous(ts):
     """
     Simplify the ts to only the contemporaneous samples, and return the new ts + node map