From cc6025a66c6b25687fd3cca4b557b502be4fe513 Mon Sep 17 00:00:00 2001
From: Jere Koskela <J.Koskela@warwick.ac.uk>
Date: Sat, 15 Oct 2022 11:52:20 +0100
Subject: [PATCH] Check for valid input in ARG likelihood

---
 msprime/likelihood.py    | 10 ++++++++++
 tests/test_likelihood.py | 27 +++++++++++++++++++++++++++
 2 files changed, 37 insertions(+)

diff --git a/msprime/likelihood.py b/msprime/likelihood.py
index e1aac7a2f..6da02f791 100644
--- a/msprime/likelihood.py
+++ b/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())
diff --git a/tests/test_likelihood.py b/tests/test_likelihood.py
index 7874d3e31..3d0830ed5 100644
--- a/tests/test_likelihood.py
+++ b/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(