avoid negative probs, closes #2237

tskit-dev · Dec 12, 2023 · cb3a469 · cb3a469
1 parent 2d00f51
commit cb3a469
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 7 deletions.
diff --git a/msprime/mutations.py b/msprime/mutations.py
@@ -296,10 +296,11 @@ def gamma_microsat(m, i, j, lo, hi):
                 rate_matrix[ii, jj] = 0
     # scale by max row sum; then normalize to 1
     row_sums = rate_matrix.sum(axis=1, dtype="float64")
-    alpha = max(row_sums)
+    alpha = np.max(row_sums)
     rate_matrix /= alpha
     row_sums = rate_matrix.sum(axis=1, dtype="float64")
-    np.fill_diagonal(rate_matrix, 1.0 - row_sums)
+    # the max(0, *) is to avoid floating point error
+    np.fill_diagonal(rate_matrix, np.fmax(0.0, 1.0 - row_sums))
     return rate_matrix
 
 

diff --git a/tests/test_mutations.py b/tests/test_mutations.py
@@ -358,11 +358,14 @@ def test_SMM(self):
         self.validate_model(model)
         self.validate_stationary(model)
 
-    def test_TPM(self):
-        lo = 2
-        hi = 50
-        m = 0.9
-        p = 0.9
+    @pytest.mark.parametrize(
+        "p, m, lo, hi",
+        [
+            (0.9, 0.9, 2, 50),
+            (0.5, 0.25, 50, 500),
+        ],
+    )
+    def test_TPM(self, p, m, lo, hi):
         model = msprime.TPM(p=p, m=m, lo=lo, hi=hi)
         self.validate_model(model)
         self.validate_stationary(model)