Check for valid input in ARG likelihood

msprime/likelihood.py

@@ -21,6 +21,8 @@
 """
 import math
 
+import numpy as np
+
 from msprime import _msprime
 
 
@@ -136,6 +138,14 @@ def log_arg_likelihood(ts, recombination_rate, Ne=1):
         sequence contains at least one recombination event, then returns
         `-DBL_MAX`.
     """
+    for tree in ts.trees():
+        if np.any(tree.num_children_array > 2):
+            raise ValueError(
+                "ARG likelihood encountered a polytomy."
+                " Tree sequences must contain binary mergers only for"
+                " valid likelihood evaluation."
+            )
+
     # Get the tables into the format we need to interchange with the low-level code.
     lw_tables = _msprime.LightweightTableCollection()
     lw_tables.fromdict(ts.tables.asdict())

tests/test_likelihood.py

@@ -169,6 +169,33 @@ def test_log_likelihoods(self):
                 msprime.log_mutation_likelihood(arg, t),
             )
 
+    def test_arg_likelihood_error_handling(self):
+        tables = tskit.TableCollection(sequence_length=1)
+        tables.nodes.add_row(
+            flags=tskit.NODE_IS_SAMPLE, population=0, individual=-1, time=0
+        )
+        tables.nodes.add_row(
+            flags=tskit.NODE_IS_SAMPLE, population=0, individual=-1, time=0
+        )
+        tables.nodes.add_row(
+            flags=tskit.NODE_IS_SAMPLE, population=0, individual=-1, time=0
+        )
+        tables.nodes.add_row(flags=0, population=0, individual=-1, time=1)
+
+        tables.edges.add_row(left=0, right=1, parent=3, child=0)
+        tables.edges.add_row(left=0, right=1, parent=3, child=1)
+        tables.edges.add_row(left=0, right=1, parent=3, child=2)
+        tables.populations.add_row()
+        arg = tables.tree_sequence()
+
+        with pytest.raises(
+            ValueError,
+            match="ARG likelihood encountered a polytomy."
+            " Tree sequences must contain binary mergers only for"
+            " valid likelihood evaluation.",
+        ):
+            msprime.log_arg_likelihood(arg, 1)
+
     def test_multiple_mrcas(self):
         tables = tskit.TableCollection(sequence_length=1)
         tables.nodes.add_row(