Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 24 additions & 2 deletions tests/test_pedigree.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ def simulate_pedigree(
num_generations=3,
sequence_length=1,
random_seed=42,
sample_gen=None,
) -> tskit.TableCollection:
"""
Simulates pedigree.
Expand All @@ -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:
Expand All @@ -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)

Expand Down Expand Up @@ -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):
Expand Down
Loading