Compute stationary distribution using power iteration to avoid numerical issue on Windows

Author: benjeffery

.github/workflows/tests.yml

@@ -95,4 +95,4 @@ jobs:
         run: |
           source ~/.profile
           conda activate anaconda-client-env
-          pytest -xvs -n0 tests/test_mutations.py::TestMatrixMutationModel::test_TPM[0.5-0.25-50-500]
+          pytest -xvs -n0

msprime/mutations.py

@@ -362,6 +362,7 @@ def __init__(
         hi=None,
         root_distribution=None,
     ):
+        print(np.show_config())
         s = 0.0 if s is None else s
         u = 0.5 if u is None else u
         v = 0.0 if v is None else v
@@ -388,13 +389,23 @@ def __init__(
             raise ValueError("u must be between 0 and 1")
         alleles = [str(int(x)) for x in range(lo, hi + 1)]
         transition_matrix = general_microsat_rate_matrix(s, u, v, p, m, lo, hi)
+
         if root_distribution is None:
-            # solve for stationary distribution
-            S, U = np.linalg.eig(transition_matrix.T)
-            U = np.real_if_close(U, tol=1)
-            stationary = np.array(U[:, np.where(np.abs(S - 1.0) < 1e-8)[0][0]])
-            stationary = stationary / np.sum(stationary)
-            root_distribution = stationary
+            # Compute the stationary distribution using power iteration
+            n = transition_matrix.shape[0]
+            pi = np.ones(n) / n
+            max_iter = 1000
+            tol = 1e-10
+            for _ in range(max_iter):
+                pi_next = pi @ transition_matrix
+                if np.max(np.abs(pi_next - pi)) < tol:
+                    root_distribution = pi_next
+                    break
+                pi = pi_next
+            else:
+                # If we haven't converged use the last approximation
+                root_distribution = pi
+
         super().__init__(alleles, root_distribution, transition_matrix)
 
 

requirements/CI-complete/requirements.txt

@@ -7,7 +7,7 @@ pytest==8.3.5
 pytest-cov==6.0.0
 pytest-xdist==3.6.1
 python_jsonschema_objects==0.5.7
-scipy==1.13.1
+scipy==1.14.0
 stdpopsim==0.1.2 #Pinned for OOA model
 tskit==0.6.0
 kastore==0.3.3