From 188d1ec9fafbafe7dc75d1e4e591a7ec441ba6c1 Mon Sep 17 00:00:00 2001 From: abureau Date: Tue, 22 Oct 2024 14:57:00 -0400 Subject: [PATCH 1/2] Initial test with internal samples Corrections to internal sample tests --- tests/test_pedigree.py | 22 ++++++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/tests/test_pedigree.py b/tests/test_pedigree.py index a9c4efa25..51bf5a91d 100644 --- a/tests/test_pedigree.py +++ b/tests/test_pedigree.py @@ -48,6 +48,7 @@ def simulate_pedigree( num_generations=3, sequence_length=1, random_seed=42, + internal_sample_gen=(None,None,None), ) -> tskit.TableCollection: """ Simulates pedigree. @@ -62,11 +63,16 @@ def simulate_pedigree( sequence_length: The sequence_length of the output tables. random_seed: Random seed. """ + # Fill-in internal_sample_gen with None if shorter than number of generations + if len(internal_sample_gen) < num_generations: + tmp = internal_sample_gen + internal_sample_gen = np.repeat(None,num_generations) + internal_sample_gen[0:len(tmp)] = tmp 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)] + curr_gen = [builder.add_individual(time=time,is_sample=internal_sample_gen[0]) 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 +86,7 @@ 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=internal_sample_gen[generation]) curr_gen.append(ind_id) return builder.finalise(sequence_length) @@ -544,6 +550,18 @@ def test_shallow(self, num_founders, recombination_rate): sequence_length=100, ) self.verify(tables, recombination_rate) + + @pytest.mark.parametrize("num_founders", [2, 3, 5, 100]) + @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, + internal_sample_gen=[True, False], + ) + self.verify(tables, recombination_rate) @pytest.mark.parametrize("num_founders", [2, 3, 10, 20]) @pytest.mark.parametrize("recombination_rate", [0, 0.01]) From 78255b6297dad71a1e602cc4996c121af1b46cdd Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Thu, 24 Jul 2025 16:15:48 +0100 Subject: [PATCH 2/2] Update the tests and API --- tests/test_pedigree.py | 28 ++++++++++++++++------------ 1 file changed, 16 insertions(+), 12 deletions(-) diff --git a/tests/test_pedigree.py b/tests/test_pedigree.py index 51bf5a91d..71a65ade2 100644 --- a/tests/test_pedigree.py +++ b/tests/test_pedigree.py @@ -48,7 +48,7 @@ def simulate_pedigree( num_generations=3, sequence_length=1, random_seed=42, - internal_sample_gen=(None,None,None), + sample_gen=None, ) -> tskit.TableCollection: """ Simulates pedigree. @@ -62,17 +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). """ - # Fill-in internal_sample_gen with None if shorter than number of generations - if len(internal_sample_gen) < num_generations: - tmp = internal_sample_gen - internal_sample_gen = np.repeat(None,num_generations) - internal_sample_gen[0:len(tmp)] = tmp rng = np.random.RandomState(random_seed) builder = msprime.PedigreeBuilder() time = num_generations - 1 - curr_gen = [builder.add_individual(time=time,is_sample=internal_sample_gen[0]) 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: @@ -86,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, is_sample=internal_sample_gen[generation]) + 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) @@ -550,16 +554,16 @@ def test_shallow(self, num_founders, recombination_rate): sequence_length=100, ) self.verify(tables, recombination_rate) - - @pytest.mark.parametrize("num_founders", [2, 3, 5, 100]) + + @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, - internal_sample_gen=[True, False], + sequence_length=100, + sample_gen=[0, 1], ) self.verify(tables, recombination_rate)