Merge pull request #2326 from abureau/internal-samples

Test_pedigree.py with internal samples

Author: jeromekelleher

tests/test_pedigree.py

@@ -48,6 +48,7 @@ def simulate_pedigree(
     num_generations=3,
     sequence_length=1,
     random_seed=42,
+    sample_gen=None,
 ) -> tskit.TableCollection:
     """
     Simulates pedigree.
@@ -61,12 +62,19 @@ def simulate_pedigree(
     num_generations: Number of generations to attempt to simulate
     sequence_length: The sequence_length of the output tables.
     random_seed: Random seed.
+    sample_gen: Generations at which all individuals are samples. Defaults
+        to the first generation (backwards in time).
     """
     rng = np.random.RandomState(random_seed)
     builder = msprime.PedigreeBuilder()
 
     time = num_generations - 1
-    curr_gen = [builder.add_individual(time=time) for _ in range(num_founders)]
+    if sample_gen is None:
+        sample_gen = [0]
+    curr_gen = [
+        builder.add_individual(time=time, is_sample=time in sample_gen)
+        for _ in range(num_founders)
+    ]
     for generation in range(1, num_generations):
         num_pairs = len(curr_gen) // 2
         if num_pairs == 0 and num_children_prob[0] != 1:
@@ -80,7 +88,9 @@ def simulate_pedigree(
             num_children = rng.choice(len(num_children_prob), p=num_children_prob)
             for _ in range(num_children):
                 parents = np.sort(parents).astype(np.int32)
-                ind_id = builder.add_individual(time=time, parents=parents)
+                ind_id = builder.add_individual(
+                    time=time, parents=parents, is_sample=time in sample_gen
+                )
                 curr_gen.append(ind_id)
     return builder.finalise(sequence_length)
 
@@ -545,6 +555,18 @@ def test_shallow(self, num_founders, recombination_rate):
         )
         self.verify(tables, recombination_rate)
 
+    @pytest.mark.parametrize("num_founders", [2, 3, 5])
+    @pytest.mark.parametrize("recombination_rate", [0, 0.01])
+    def test_shallow_internal(self, num_founders, recombination_rate):
+        tables = simulate_pedigree(
+            num_founders=num_founders,
+            num_children_prob=[0, 0, 1],
+            num_generations=2,
+            sequence_length=100,
+            sample_gen=[0, 1],
+        )
+        self.verify(tables, recombination_rate)
+
     @pytest.mark.parametrize("num_founders", [2, 3, 10, 20])
     @pytest.mark.parametrize("recombination_rate", [0, 0.01])
     def test_deep(self, num_founders, recombination_rate):