make testing more clear

Author: petrelharp

python/tests/test_ld_matrix.py

@@ -2063,38 +2063,36 @@ def compute_branch_stat(
         ts
         for ts in get_example_tree_sequences()
         if ts.id
-        not in {
-            "no_samples",
-            "empty_ts",
-            # We must skip these cases so that tests run in a reasonable
-            # amount of time. To get more complete testing, these filters
-            # can be commented out. (runtime ~1hr)
-            "gap_0",
-            "gap_0.1",
-            "gap_0.5",
-            "gap_0.75",
-            "n=2_m=32_rho=0",
-            "n=10_m=1_rho=0",
-            "n=10_m=1_rho=0.1",
-            "n=10_m=2_rho=0",
-            "n=10_m=2_rho=0.1",
-            "n=10_m=32_rho=0",
-            "n=10_m=32_rho=0.1",
-            "n=10_m=32_rho=0.5",
+        in {
+            # We run only these cases so that tests run in a reasonable
+            # amount of time. All examples takes ~1hr.
+            "decapitate_recomb",
+            "gap_at_end",
+            "all_nodes_samples",
+            "internal_nodes_samples",
+            "mixed_internal_leaf_samples",
+            "bottleneck_n=3_mutated",
+            "bottleneck_n=10_mutated",
+            "rev_node_order",
+            "empty_tree",
+            "n=3_m=2_rho=0.5",
+            "n=3_m=32_rho=0",
+            "n=3_m=32_rho=0.1",
+            "n=2_m=1_rho=0",
+            "n=2_m=1_rho=0.1",
+            "n=2_m=1_rho=0.5",
+            "n=2_m=2_rho=0",
+            "n=2_m=2_rho=0.1",
+            "n=2_m=2_rho=0.5",
+            "n=2_m=32_rho=0.1",
+            "n=2_m=32_rho=0.5",
+            "n=3_m=1_rho=0",
+            "n=3_m=1_rho=0.5",
+            "n=3_m=2_rho=0",
+            "n=10_m=1_rho=0.5",
+            "n=10_m=2_rho=0.5",
             # we keep one n=100 case to ensure bit arrays are working
-            "n=100_m=1_rho=0.1",
-            "n=100_m=1_rho=0.5",
-            "n=100_m=2_rho=0",
-            "n=100_m=2_rho=0.1",
-            "n=100_m=2_rho=0.5",
-            "n=100_m=32_rho=0",
-            "n=100_m=32_rho=0.1",
-            "n=100_m=32_rho=0.5",
-            "all_fields",
-            "back_mutations",
-            "multichar",
-            "multichar_no_metadata",
-            "bottleneck_n=100_mutated",
+            "n=100_m=1_rho=0",
         }
     ],
 )
@@ -2548,9 +2546,13 @@ def D2_ij_unbiased(X, n):
     "ts,stat",
     [
         (
-            ts := [
-                p for p in get_example_tree_sequences() if p.id == "n=100_m=32_rho=0.5"
-            ][0].values[0],
+            ts := tsutil.get_sim_example(
+                sample_size=100,
+                sequence_length=32,
+                recombination_rate=0.5,
+                mutation_rate=0.1,
+                seed=123,
+            ),
             "D",
         ),
         (ts, "D2"),
@@ -2575,9 +2577,13 @@ def test_general_two_locus_site_stat(ts, stat):
     "ts,stat",
     [
         (
-            ts := [
-                p for p in get_example_tree_sequences() if p.id == "n=100_m=32_rho=0.5"
-            ][0].values[0],
+            ts := tsutil.get_sim_example(
+                sample_size=100,
+                sequence_length=32,
+                recombination_rate=0.5,
+                mutation_rate=0.1,
+                seed=123,
+            ),
             "r2_ij",
         ),
         (ts, "D2_ij"),
@@ -2609,7 +2615,7 @@ def test_general_two_locus_two_way_site_stat(ts, stat):
     ],
 )
 def test_general_one_way_two_locus_stat_multiallelic(stat):
-    (ts,) = {t.id: t for t in get_example_tree_sequences()}["all_fields"].values
+    ts = tsutil.all_fields_ts()
     func = getattr(GeneralStatFuncs, stat)
     if stat == "r2":
         result = ts.two_locus_count_stat(
@@ -2632,7 +2638,7 @@ def test_general_one_way_two_locus_stat_multiallelic(stat):
     ],
 )
 def test_general_two_way_two_locus_stat_multiallelic(stat):
-    (ts,) = {t.id: t for t in get_example_tree_sequences()}["all_fields"].values
+    ts = tsutil.all_fields_ts()
     func = getattr(GeneralStatFuncs, stat)
     if stat == "r2_ij":
 

python/tests/tsutil.py

@@ -2463,6 +2463,28 @@ def get_back_mutation_examples():
         yield insert_branch_mutations(ts)
 
 
+@functools.lru_cache
+def get_sim_example(
+    sample_size, sequence_length, recombination_rate, mutation_rate, seed
+):
+    recomb_map = msprime.RecombinationMap.uniform_map(
+        sequence_length, recombination_rate
+    )
+    ts = msprime.simulate(
+        recombination_map=recomb_map,
+        mutation_rate=mutation_rate,
+        random_seed=seed,
+        population_configurations=[
+            msprime.PopulationConfiguration(sample_size),
+            msprime.PopulationConfiguration(0),
+        ],
+        migration_matrix=[[0, 1], [1, 0]],
+    )
+    ts = insert_random_ploidy_individuals(ts, 4, seed=seed)
+    ts = add_random_metadata(ts, seed=seed)
+    return ts
+
+
 def make_example_tree_sequences(custom_max=None):
     yield from get_decapitated_examples(custom_max=custom_max)
     yield from get_gap_examples(custom_max=custom_max)
@@ -2475,22 +2497,14 @@ def make_example_tree_sequences(custom_max=None):
     for n in n_list:
         for m in [1, 2, 32]:
             for rho in [0, 0.1, 0.5]:
-                recomb_map = msprime.RecombinationMap.uniform_map(m, rho, num_loci=m)
-                ts = msprime.simulate(
-                    recombination_map=recomb_map,
+                ts = get_sim_example(
+                    sample_size=n,
+                    sequence_length=m,
+                    recombination_rate=rho,
                     mutation_rate=0.1,
-                    random_seed=seed,
-                    population_configurations=[
-                        msprime.PopulationConfiguration(n),
-                        msprime.PopulationConfiguration(0),
-                    ],
-                    migration_matrix=[[0, 1], [1, 0]],
-                )
-                ts = insert_random_ploidy_individuals(ts, 4, seed=seed)
-                yield (
-                    f"n={n}_m={m}_rho={rho}",
-                    add_random_metadata(ts, seed=seed),
+                    seed=seed,
                 )
+                yield (f"n={n}_m={m}_rho={rho}", ts)
                 seed += 1
     for name, ts in get_bottleneck_examples(custom_max=custom_max):
         yield (