From c1203ab6414fb218cf49cbc9e5e3994162d226fc Mon Sep 17 00:00:00 2001
From: szhan <shing.zhan@gmail.com>
Date: Mon, 17 Jun 2024 12:42:31 +0100
Subject: [PATCH] Refactor getting recombination rates for test input

---
 tests/lsbase.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/tests/lsbase.py b/tests/lsbase.py
index a62ec28..4d16d7d 100644
--- a/tests/lsbase.py
+++ b/tests/lsbase.py
@@ -136,11 +136,14 @@ def get_examples_pars(
             H, ts.num_samples
         ), "Reference haplotypes have unexpected number of copiable entries."
 
-        rs = [
+        r_s = [
             np.zeros(m) + 0.01,  # Equal recombination and mutation
             np.random.rand(m),  # Random
             1e-5 * (np.random.rand(m) + 0.5) / 2,
         ]
+        for i in range(len(r_s)):
+            r_s[i][0] = 0
+
         mus = [
             np.zeros(m) + 0.01,  # Equal recombination and mutation
             np.random.rand(m) * 0.2,  # Random
@@ -148,13 +151,12 @@ def get_examples_pars(
         ]
 
         if include_extreme_rates:
-            rs.append(np.zeros(m) + 0.2)
-            rs.append(np.zeros(m) + 1e-6)
+            r_s.append(np.zeros(m) + 0.2)
+            r_s.append(np.zeros(m) + 1e-6)
             mus.append(np.zeros(m) + 0.2)
             mus.append(np.zeros(m) + 1e-6)
 
-        for query, r, mu in itertools.product(queries, rs, mus):
-            r[0] = 0
+        for query, r, mu in itertools.product(queries, r_s, mus):
             # Must be calculated from the genotype matrix,
             # because we can now get back mutations that
             # result in the number of alleles being higher