Improve HK genedrop test coverage

sgkit/tests/test_genedrop.py

@@ -65,6 +65,99 @@ def test_random_inheritance_Hamilton_Kerr():
     np.testing.assert_allclose(observed_lambda, 0.25, atol=0.001)  # lambda working
 
 
+@pytest.mark.parametrize(
+    "genotypes",
+    [
+        [
+            [0, 1, 2, 3],
+            [4, 5, -2, -2],  # padding after alleles
+            [-1, -1, -1, -1],
+        ],
+        [
+            [0, 1, 2, 3],
+            [-2, -2, 4, 5],  # padding before alleles
+            [-1, -1, -1, -1],
+        ],
+    ],
+)
+def test_random_inheritance_Hamilton_Kerr__mixed_ploidy(genotypes):
+    seed_numba(0)
+    genotypes = np.array(genotypes)
+    parent = np.array(
+        [
+            [-1, -1],
+            [-1, -1],
+            [0, 1],
+        ]
+    )
+    tau = np.array(
+        [
+            [2, 2],
+            [1, 1],
+            [2, 2],
+        ],
+        dtype=np.uint64,
+    )
+    lambda_ = np.array(
+        [
+            [0.0, 0.0],
+            [0.0, 0.0],
+            [0.0, 0.25],
+        ]
+    )
+    marked = np.zeros(4, bool)
+    marked = np.zeros(4, bool)
+    n_reps = 1_000_000
+    results = np.zeros((n_reps, 4), int)
+    for i in range(n_reps):
+        _random_inheritance_Hamilton_Kerr(genotypes, parent, tau, lambda_, marked, 2)
+        results[i] = genotypes[2]
+    # check for allelic bias in first parent
+    unique, counts = np.unique(results[:, 0], return_counts=True)
+    assert set(unique) == set(genotypes[0])  # all alleles present
+    np.testing.assert_allclose(1 / 4, counts / counts.sum(), atol=0.001)  # no bias
+    assert np.all(results[:, 0] != results[:, 1])  # no duplicates
+    # check for allelic bias in second parent (lambda > 0)
+    unique, counts = np.unique(results[:, 2], return_counts=True)
+    assert set(unique) == {4, 5}  # all alleles present
+    np.testing.assert_allclose(1 / 2, counts / counts.sum(), atol=0.001)  # no bias
+    observed_lambda = np.mean(results[:, 2] == results[:, 3])
+    np.testing.assert_allclose(observed_lambda, 0.25, atol=0.001)  # lambda working
+
+
+def test_random_inheritance_Hamilton_Kerr__padding():
+    seed_numba(0)
+    genotypes = np.array(
+        [
+            [0, 1, 2, 3],
+            [-1, -1, -1, -1],
+        ]
+    )
+    parent = np.array(
+        [
+            [-1, -1],
+            [-1, 0],  # half-clone
+        ]
+    )
+    tau = np.array(
+        [
+            [2, 2],
+            [0, 2],  # half-clone
+        ],
+        dtype=np.uint64,
+    )
+    lambda_ = np.array(
+        [
+            [0.0, 0.0],
+            [0.0, 0.0],
+        ]
+    )
+    marked = np.zeros(4, bool)
+    _random_inheritance_Hamilton_Kerr(genotypes, parent, tau, lambda_, marked, 1)
+    assert np.all(genotypes[1, 0:2] >= 0)
+    assert np.all(genotypes[1, 2:] == -2)  # correctly padded
+
+
 def test_random_inheritance_Hamilton_Kerr__raise_on_ploidy():
     seed_numba(0)
     genotypes = np.array(
@@ -186,6 +279,38 @@ def test_random_inheritance_Hamilton_Kerr__raise_on_non_zero_lambda():
         _random_inheritance_Hamilton_Kerr(genotypes, parent, tau, lambda_, marked, 3)
 
 
+def test_random_inheritance_Hamilton_Kerr__on_raise_half_founder():
+    seed_numba(0)
+    genotypes = np.array(
+        [
+            [0, 1, 2, 3],
+            [4, 5, 6, 7],
+        ]
+    )
+    parent = np.array(
+        [
+            [-1, -1],
+            [-1, 0],  # half-founder
+        ]
+    )
+    tau = np.array(
+        [
+            [2, 2],
+            [2, 2],
+        ],
+        dtype=np.uint64,
+    )
+    lambda_ = np.array(
+        [
+            [0.0, 0.0],
+            [0.0, 0.0],
+        ]
+    )
+    marked = np.zeros(4, bool)
+    with pytest.raises(ValueError, match="Pedigree contains half-founders."):
+        _random_inheritance_Hamilton_Kerr(genotypes, parent, tau, lambda_, marked, 1)
+
+
 @pytest.mark.parametrize("permute", [False, True])
 @pytest.mark.parametrize("chunks", [None, 5000])
 def test_simulate_genedrop__diploid_kinship(permute, chunks):