From fa8b7e21886524cc1758a9c2d0c2224b6794ec52 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Tue, 29 Jul 2025 09:59:06 +0100 Subject: [PATCH 1/4] simulations above root implementation --- CHANGELOG.md | 2 + algorithms.py | 95 +++++++++++---- docs/ancestry.md | 42 +++++++ lib/msprime.c | 81 +++++++++++-- lib/msprime.h | 3 + lib/tests/test_ancestry.c | 190 +++++++++++++++++++++++++++++ msprime/__init__.py | 1 + msprime/_msprimemodule.c | 13 +- msprime/ancestry.py | 28 +++++ tests/test_algorithms.py | 24 ++++ tests/test_ancestry.py | 243 ++++++++++++++++++++++++++++++++++++++ 11 files changed, 689 insertions(+), 33 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index adf47bb10..2174db8f0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,8 @@ **Maintenance** - Add support and wheels for Python3.13 +- Support for simulations after local mrca, + ({pr}`2396`, {user}`jeromekelleher`, {user}`hossam26644`). - Drop Python 3.9 support, require Python >= 3.10 ({pr}`2418`, {user}`benjeffery`) diff --git a/algorithms.py b/algorithms.py index e17a8ed19..8521793de 100644 --- a/algorithms.py +++ b/algorithms.py @@ -19,6 +19,7 @@ logger = daiquiri.getLogger() +INFINITY = sys.float_info.max class FenwickTree: @@ -247,7 +248,7 @@ def _get_common_ancestor_waiting_time(self, np, t): Returns the random waiting time until a common ancestor event occurs within this population. """ - ret = sys.float_info.max + ret = INFINITY u = random.expovariate(2 * np) if self.growth_rate == 0: ret = self.start_size * u @@ -267,7 +268,7 @@ def _get_common_ancestor_waiting_time(self, np, t): def get_common_ancestor_waiting_time_hudson(self): def _get_common_ancestor_waiting_time_hudson(t): k = self.get_num_ancestors() - ret = sys.float_info.max + ret = INFINITY if k > 1: np = k * (k - 1) / 2 ret = self._get_common_ancestor_waiting_time(np, t) @@ -278,7 +279,7 @@ def _get_common_ancestor_waiting_time_hudson(t): def get_common_ancestor_waiting_time_smc_k(self): def _get_common_ancestor_waiting_time_smc_k(t): np = self.get_num_pairs() - ret = sys.float_info.max + ret = INFINITY if np > 0: ret = self._get_common_ancestor_waiting_time(np, t) return ret @@ -885,6 +886,7 @@ def __init__( gene_conversion_length=1, discrete_genome=True, hull_offset=None, + stop_at_local_mrca=True, ): # Must be a square matrix. N = len(migration_matrix) @@ -906,6 +908,7 @@ def __init__( self.migration_matrix = migration_matrix self.num_labels = num_labels self.num_populations = N + self.stop_at_local_mrca = stop_at_local_mrca self.max_segments = max_segments self.coalescing_segments_only = coalescing_segments_only self.additional_nodes = msprime.NodeType(additional_nodes) @@ -964,7 +967,7 @@ def __init__( self.sweep_trajectory = sweep_trajectory self.time_slice = time_slice - self.modifier_events = [(sys.float_info.max, None, None)] + self.modifier_events = [(INFINITY, None, None)] for time, pop_id, new_size in population_size_changes: self.modifier_events.append( (time, self.change_population_size, (int(pop_id), new_size)) @@ -1253,7 +1256,7 @@ def simulate(self, end_time): if self.model == "hudson": self.hudson_simulate(end_time) elif self.model == "dtwf": - self.dtwf_simulate() + self.dtwf_simulate(end_time) elif self.model == "fixed_pedigree": self.pedigree_simulate() elif self.model == "single_sweep": @@ -1315,7 +1318,6 @@ def hudson_simulate(self, end_time): """ Simulates the algorithm until all loci have coalesced. """ - infinity = sys.float_info.max non_empty_pops = {pop.id for pop in self.P if pop.get_num_ancestors() > 0} potential_destinations = self.get_potential_destinations() @@ -1326,23 +1328,23 @@ def hudson_simulate(self, end_time): break # self.print_state() re_rate = self.get_total_recombination_rate(label=0) - t_re = infinity + t_re = INFINITY if re_rate > 0: t_re = random.expovariate(re_rate) # Gene conversion can occur within segments .. gc_rate = self.get_total_gc_rate(label=0) - t_gcin = infinity + t_gcin = INFINITY if gc_rate > 0: t_gcin = random.expovariate(gc_rate) # ... or to the left of the first segment. gc_left_rate = self.get_total_gc_left_rate(label=0) - t_gc_left = infinity + t_gc_left = INFINITY if gc_left_rate > 0: t_gc_left = random.expovariate(gc_left_rate) # Common ancestor events occur within demes. - t_ca = infinity + t_ca = INFINITY for index in non_empty_pops: pop = self.P[index] assert pop.get_num_ancestors() > 0 @@ -1350,7 +1352,7 @@ def hudson_simulate(self, end_time): if t < t_ca: t_ca = t ca_population = index - t_mig = infinity + t_mig = INFINITY # Migration events happen at the rates in the matrix. for j in non_empty_pops: source_size = self.P[j].get_num_ancestors() @@ -1365,7 +1367,9 @@ def hudson_simulate(self, end_time): mig_source = j mig_dest = k min_time = min(t_re, t_ca, t_gcin, t_gc_left, t_mig) - assert min_time != infinity + assert (min_time != INFINITY) or ( + (min_time == INFINITY) and not self.stop_at_local_mrca + ) if self.t + min_time > self.modifier_events[0][0]: t, func, args = self.modifier_events.pop(0) self.t = t @@ -1379,7 +1383,14 @@ def hudson_simulate(self, end_time): event = "MOD" else: self.t += min_time - if min_time == t_re: + if min_time >= INFINITY: + assert not self.stop_at_local_mrca + event = "END" + if end_time >= INFINITY: + end_time = self.t + else: + self.t = end_time + elif min_time == t_re: event = "RE" self.hudson_recombination_event(0) elif min_time == t_gcin: @@ -1400,6 +1411,7 @@ def hudson_simulate(self, end_time): non_empty_pops.remove(mig_source) assert self.P[mig_dest].get_num_ancestors() > 0 non_empty_pops.add(mig_dest) + logger.info( "%s time=%f n=%d", event, @@ -1524,20 +1536,39 @@ def pedigree_simulate(self): self.pedigree = Pedigree(self.tables) self.dtwf_climb_pedigree() - def dtwf_simulate(self): + def dtwf_above_root_simulations_stop(self, events_happened): + return ( + not events_happened + and not self.stop_at_local_mrca + and sum(pop.get_num_ancestors() for pop in self.P) == 1 + ) + + def dtwf_simulate(self, end_time): """ Simulates the algorithm until all loci have coalesced. """ - while self.ancestors_remain(): + while self.ancestors_remain() and self.t < end_time: self.t += 1 self.verify() - self.dtwf_generation() + events_happened = self.dtwf_generation() + + if self.dtwf_above_root_simulations_stop(events_happened): + if end_time >= INFINITY: + end_time = self.t + else: + self.t = end_time + break def dtwf_generation(self): """ Evolves one generation of a Wright Fisher population + Returns True if any events occurred, False otherwise. """ # Migration events happen at the rates in the matrix. + nodes = self.tables.nodes.num_rows + edges = self.tables.edges.num_rows + no_migration = True + for j in range(len(self.P)): source_size = self.P[j].get_num_ancestors() for k in range(len(self.P)): @@ -1549,6 +1580,7 @@ def dtwf_generation(self): mig_source = j mig_dest = k self.migration_event(mig_source, mig_dest) + no_migration = False for pop_idx, pop in enumerate(self.P): # Cluster haploid inds by parent @@ -1605,6 +1637,14 @@ def dtwf_generation(self): h, pop_idx, 0, parent_nodes[ploid] ) # label 0 only self.verify() + if ( + self.tables.nodes.num_rows == nodes + and self.tables.edges.num_rows == edges + and no_migration + ): + # dtwf generation had no events + return False + return True def process_pedigree_common_ancestors(self, ind, ploid): """ @@ -2256,14 +2296,16 @@ def merge_ancestors(self, H, pop_id, label, new_node_id=-1): if r_max not in self.S: j = self.S.floor_key(r_max) self.S[r_max] = self.S[j] + + min_overlap = len(X) if self.stop_at_local_mrca else 0 # Update the number of extant segments. - if self.S[left] == len(X): + if self.S[left] == min_overlap: self.S[left] = 0 right = self.S.succ_key(left) else: right = left - while right < r_max and self.S[right] != len(X): - self.S[right] -= len(X) - 1 + while right < r_max and self.S[right] != min_overlap: + self.S[right] -= min_overlap - 1 right = self.S.succ_key(right) alpha = self.alloc_segment(left, right, new_node_id, pop_id) # Update the heaps and make the record. @@ -2394,6 +2436,7 @@ def merge_two_ancestors(self, population_index, label, x, y, u=-1): new_lineage = self.alloc_lineage(None, population_index, label=label) coalescence = False defrag_required = False + min_overlap = 2 if self.stop_at_local_mrca else 0 while x is not None or y is not None: alpha = None @@ -2435,12 +2478,12 @@ def merge_two_ancestors(self, population_index, label, x, y, u=-1): j = self.S.floor_key(r_max) self.S[r_max] = self.S[j] # Update the number of extant segments. - if self.S[left] == 2: + if self.S[left] == min_overlap: self.S[left] = 0 right = self.S.succ_key(left) else: right = left - while right < r_max and self.S[right] != 2: + while right < r_max and self.S[right] != min_overlap: self.S[right] -= 1 right = self.S.succ_key(right) alpha = self.alloc_segment( @@ -2860,6 +2903,7 @@ def run_simulate(args): gene_conversion_length=mean_tract_length, discrete_genome=args.discrete, hull_offset=args.offset, + stop_at_local_mrca=not args.continue_after_local_mrca, ) ts = s.simulate(args.end_time) ts.dump(args.output_file) @@ -2946,6 +2990,15 @@ def add_simulator_arguments(parser): help="The delta_t value for selective sweeps", ) parser.add_argument("--model", default="hudson") + parser.add_argument( + "--continue-after-local-mrca", + action="store_true", + default=False, + help=( + "If set, continue after local MRCA (i.e., do not stop). " + "Default: False (stop at local MRCA)." + ), + ) parser.add_argument("--offset", type=float, default=0.0) parser.add_argument( "--from-ts", diff --git a/docs/ancestry.md b/docs/ancestry.md index d41419d0b..12f4ac4dc 100644 --- a/docs/ancestry.md +++ b/docs/ancestry.md @@ -1129,6 +1129,48 @@ However, if `coalescing_segments_only=False`, then the edges recorded would be from `m` to `p` over the entire segment `[a, b)`, and from `n` to `p` over the entire segment `[c,d)`. The nodes `m` and `n` coalesce (in `p`) on only the overlapping segment `[c,b)`, and so the node `p` will be a unary node in the flanking regions: above `m` on the segment `[a,c)` and above `n` on the segment `[b,d)`. +(sec_simulation_after_mrca)= + +### Simulations after local MRCA + +By default, msprime stops simulating local trees when a local most recent common ancestor (MRCA) is found. This is because events that occur + above the common ancestor are shared across all samples, and we are usually interested in differences between samples. +```{code-cell} +import msprime + +ts = msprime.sim_ancestry(2, recombination_rate=0.1, sequence_length=10, random_seed=21) +ts.draw_svg(time_scale='rank') +``` +Here, for example, we see that simulations in the middle tree stopped at node 6, which is much younger than nodes 7 and 8. +And if we look closely, we can also see that simulations stopped at node 7 in the rightmost tree at a younger time than +node 8 in the leftmost tree. + + +However, for some specialised applications, simulations after the local MRCA might be needed. In this case, we set the +parameter `stop_at_local_mrca` of {func}`.sim_ancestry` to `False`, and we set the `end_time` parameter to stop simulations +when they reach this end time. + + +```{code-cell} +import msprime + +ts = msprime.sim_ancestry( + 2, recombination_rate=0.1, + sequence_length=10, + random_seed=21, + stop_at_local_mrca=False, + end_time=10) + +ts.draw_svg(time_scale='rank') +``` +It is possible to run the simulations without setting an `end_time`. In that case, simulations will +stop when no more events (recombination, coalescence, migration, etc.) could happen. However, since in many cases this state is impossible to reach, we recommend setting an `end_time`. + +Note that you might want to set an `end_time` that is larger than possible end times for your simulations to avoid stopping +your simulations before full coalescence. However, a larger `end_time` means that simulations will take longer and consume +more memory. + + (sec_additional_nodes_ca)= ### Common ancestor events diff --git a/lib/msprime.c b/lib/msprime.c index fc84f3cc3..bc23cb306 100644 --- a/lib/msprime.c +++ b/lib/msprime.c @@ -328,6 +328,13 @@ msp_set_ploidy(msp_t *self, int ploidy) return ret; } +int +msp_set_stop_at_local_mrca(msp_t *self, bool stop_at_local_mrca) +{ + self->stop_at_local_mrca = stop_at_local_mrca; + return 0; +} + int msp_set_discrete_genome(msp_t *self, bool is_discrete) { @@ -1030,6 +1037,8 @@ msp_alloc(msp_t *self, tsk_table_collection_t *tables, gsl_rng *rng) self->store_migrations = false; self->additional_nodes = 0; self->coalescing_segments_only = true; + self->stop_at_local_mrca = true; + self->dtwf_no_events_generation = false; self->avl_node_block_size = 1024; self->node_mapping_block_size = 1024; self->segment_block_size = 1024; @@ -3777,6 +3786,7 @@ msp_merge_two_ancestors(msp_t *self, population_id_t population_id, label_id_t l bool defrag_required = false; tsk_id_t v; double l, r, l_min, r_max; + uint32_t min_overlap; avl_node_t *node; node_mapping_t *nm, search; segment_t *x, *y, *z, *alpha, *beta; @@ -3787,6 +3797,12 @@ msp_merge_two_ancestors(msp_t *self, population_id_t population_id, label_id_t l goto out; } + if (self->stop_at_local_mrca) { + min_overlap = 2; + } else { + min_overlap = 0; + } + x = a; y = b; /* Keep GCC happy */ @@ -3862,7 +3878,7 @@ msp_merge_two_ancestors(msp_t *self, population_id_t population_id, label_id_t l node = avl_search(&self->overlap_counts, &search); tsk_bug_assert(node != NULL); nm = (node_mapping_t *) node->item; - if (nm->value == 2) { + if (nm->value == min_overlap) { nm->value = 0; node = node->next; tsk_bug_assert(node != NULL); @@ -3870,7 +3886,7 @@ msp_merge_two_ancestors(msp_t *self, population_id_t population_id, label_id_t l r = nm->position; } else { r = l; - while (nm->value != 2 && r < r_max) { + while (nm->value != min_overlap && r < r_max) { nm->value--; node = node->next; tsk_bug_assert(node != NULL); @@ -3983,6 +3999,8 @@ msp_merge_ancestors(msp_t *self, avl_tree_t *Q, population_id_t population_id, bool defrag_required = false; uint32_t j, h; double l, r, r_max, next_l, l_min; + uint32_t min_overlap; + avl_node_t *node; node_mapping_t *nm, search; segment_t *x, *z, *alpha; @@ -4072,7 +4090,12 @@ msp_merge_ancestors(msp_t *self, avl_tree_t *Q, population_id_t population_id, node = avl_search(&self->overlap_counts, &search); tsk_bug_assert(node != NULL); nm = (node_mapping_t *) node->item; - if (nm->value == h) { + if (self->stop_at_local_mrca) { + min_overlap = h; + } else { + min_overlap = 0; + } + if (nm->value == min_overlap) { nm->value = 0; node = node->next; tsk_bug_assert(node != NULL); @@ -4080,8 +4103,8 @@ msp_merge_ancestors(msp_t *self, avl_tree_t *Q, population_id_t population_id, r = nm->position; } else { r = l; - while (nm->value != h && r < r_max) { - nm->value -= h - 1; + while (nm->value != min_overlap && r < r_max) { + nm->value -= min_overlap - 1; node = node->next; tsk_bug_assert(node != NULL); nm = (node_mapping_t *) node->item; @@ -5189,9 +5212,20 @@ msp_run_coalescent(msp_t *self, double max_time, unsigned long max_events) t_wait = GSL_MIN(mig_t_wait, GSL_MIN(gc_t_wait, GSL_MIN(gc_left_t_wait, GSL_MIN(re_t_wait, ca_t_wait)))); - if (fixed_event_time == DBL_MAX && t_wait == DBL_MAX) { + if (fixed_event_time == DBL_MAX && t_wait == DBL_MAX + && self->stop_at_local_mrca) { // TODO add condition ret = MSP_ERR_INFINITE_WAITING_TIME; goto out; + } else if ((fixed_event_time == DBL_MAX && t_wait == DBL_MAX) + && !self->stop_at_local_mrca) { + + if (max_time >= DBL_MAX) { + max_time = self->time; + } else { + self->time = max_time; + } + ret = MSP_EXIT_COALESCENCE; + break; } random_event_time = self->time + t_wait; @@ -5442,6 +5476,10 @@ msp_dtwf_generation(msp_t *self) label_id_t label = 0; tsk_id_t parent_nodes[MSP_MAX_PED_PLOIDY]; + unsigned long int edgs = self->tables->edges.num_rows; + unsigned long int nodes = self->tables->nodes.num_rows; + + self->dtwf_no_events_generation = false; for (i = 0; i < 2; i++) { avl_init_tree(&Q[i], cmp_segment_queue, NULL); } @@ -5534,6 +5572,16 @@ msp_dtwf_generation(msp_t *self) } } } + + if (self->tables->edges.num_rows == edgs + && self->tables->nodes.num_rows == nodes) { + // No edges or nodes added this generation + // i.e., no coalescence or recombination + self->dtwf_no_events_generation = true; + } else { + edgs = self->tables->edges.num_rows; + nodes = self->tables->nodes.num_rows; + } free(parents); free(segment_mem); segment_mem = NULL; @@ -5618,6 +5666,7 @@ msp_run_dtwf(msp_t *self, double max_time, unsigned long max_events) avl_tree_t *nodes; /* Only support a single structured coalescent label at the moment */ label_id_t label = 0; + bool migration_occurred; tsk_bug_assert(self->recomb_mass_index == NULL); tsk_bug_assert(self->gc_mass_index == NULL); @@ -5650,6 +5699,8 @@ msp_run_dtwf(msp_t *self, double max_time, unsigned long max_events) } self->time++; + migration_occurred = false; + /* Following SLiM, we perform migrations prior to selecting * parents for the current generation */ node_trees @@ -5699,6 +5750,7 @@ msp_run_dtwf(msp_t *self, double max_time, unsigned long max_events) mig_dest_pop = (population_id_t) k; for (i = 0; i < n[k]; i++) { + migration_occurred = true; ret = msp_store_simultaneous_migration_events( self, nodes, mig_source_pop, label); if (ret != 0) { @@ -5715,6 +5767,7 @@ msp_run_dtwf(msp_t *self, double max_time, unsigned long max_events) nodes = &node_trees[j * self->num_populations + k]; mig_source_pop = (population_id_t) j; mig_dest_pop = (population_id_t) k; + migration_occurred = true; ret = msp_simultaneous_migration_event( self, nodes, mig_source_pop, mig_dest_pop); if (ret != 0) { @@ -5737,13 +5790,15 @@ msp_run_dtwf(msp_t *self, double max_time, unsigned long max_events) goto out; } ret = msp_apply_demographic_events(self, self->next_demographic_event->time); + // Demographic events can change population structure + migration_occurred = true; if (ret != 0) { goto out; } } self->time = cur_time; ret = msp_dtwf_generation(self); - if (ret != 0) { + if (ret < 0) { goto out; } @@ -5758,6 +5813,18 @@ msp_run_dtwf(msp_t *self, double max_time, unsigned long max_events) } self->next_sampling_event++; } + + if (self->dtwf_no_events_generation && (!migration_occurred) + && (!self->stop_at_local_mrca) && msp_get_num_ancestors(self) == 1) { + /* No events occurred: we are done */ + if (max_time >= DBL_MAX) { + max_time = self->time; + } else { + self->time = max_time; + } + ret = MSP_EXIT_COALESCENCE; + break; + } } out: msp_safe_free(node_trees); diff --git a/lib/msprime.h b/lib/msprime.h index c352ef539..30ec2d4d4 100644 --- a/lib/msprime.h +++ b/lib/msprime.h @@ -256,6 +256,8 @@ typedef struct _msp_t { bool store_full_arg; uint32_t additional_nodes; bool coalescing_segments_only; + bool stop_at_local_mrca; + bool dtwf_no_events_generation; double sequence_length; bool discrete_genome; rate_map_t recomb_map; @@ -497,6 +499,7 @@ int msp_set_store_migrations(msp_t *self, bool store_migrations); int msp_set_store_full_arg(msp_t *self, bool store_full_arg); int msp_set_additional_nodes(msp_t *self, uint32_t additional_nodes); int msp_set_coalescing_segments_only(msp_t *self, bool coalescing_segments_only); +int msp_set_stop_at_local_mrca(msp_t *self, bool stop_at_local_mrca); int msp_set_ploidy(msp_t *self, int ploidy); int msp_set_recombination_map(msp_t *self, size_t size, double *position, double *rate); int msp_set_recombination_rate(msp_t *self, double rate); diff --git a/lib/tests/test_ancestry.c b/lib/tests/test_ancestry.c index 19ef82dea..751ec0ce0 100644 --- a/lib/tests/test_ancestry.c +++ b/lib/tests/test_ancestry.c @@ -666,6 +666,192 @@ test_multi_locus_simulation(void) gsl_rng_free(rng); } +static void +test_single_locus_continue_after_local_mrca_with_end_time(void) +{ + int ret; + uint32_t n = 10; + uint32_t m = 10; + long t = 100000.0; // end time + long seed = 10; + int model; + int models[] = { MSP_MODEL_HUDSON, MSP_MODEL_SMC, MSP_MODEL_SMC_PRIME, + MSP_MODEL_SMC_K, MSP_MODEL_DTWF }; + const char *model_names[] = { "hudson", "smc", "smc_prime", "smc_k", "dtwf" }; + const char *model_name; + size_t j; + tsk_table_collection_t tables; + msp_t msp; + tsk_treeseq_t ts; + tsk_tree_t tree; + tsk_id_t root; + + gsl_rng *rng = safe_rng_alloc(); + + gsl_rng_set(rng, seed); + + for (j = 0; j < sizeof(models) / sizeof(int); j++) { + + ret = build_sim(&msp, &tables, rng, m, 1, NULL, n); + ret = msp_set_recombination_rate(&msp, 0.01); + CU_ASSERT_EQUAL_FATAL(ret, 0); + + ret = msp_set_stop_at_local_mrca(&msp, false); + CU_ASSERT_EQUAL_FATAL(ret, 0); + + switch (j) { + case 0: + ret = msp_set_simulation_model_hudson(&msp); + break; + case 1: + ret = msp_set_simulation_model_smc(&msp); + break; + case 2: + ret = msp_set_simulation_model_smc_prime(&msp); + break; + case 3: + ret = msp_set_simulation_model_smc_k(&msp, 0); + break; + case 4: + ret = msp_set_simulation_model_dtwf(&msp); + break; + } + ret = msp_initialise(&msp); + CU_ASSERT_EQUAL(ret, 0); + + ret = msp_run(&msp, t, ULONG_MAX); + CU_ASSERT(ret == 0 || ret == 2); // exit on max time or simulations done + + ret = msp_finalise_tables(&msp); + CU_ASSERT_EQUAL(ret, 0); + + ret = msp_get_time(&msp); + CU_ASSERT_EQUAL(ret, t); + + ret = tsk_treeseq_init(&ts, &tables, TSK_TS_INIT_BUILD_INDEXES); + CU_ASSERT_EQUAL_FATAL(ret, 0); + + ret = tsk_tree_init(&tree, &ts, 0); + CU_ASSERT_EQUAL_FATAL(ret, 0); + + for (ret = tsk_tree_first(&tree); ret == 1; ret = tsk_tree_next(&tree)) { + + // verrify each tree has a single root + CU_ASSERT_EQUAL_FATAL(tsk_tree_get_num_roots(&tree), 1); + + // verify that all roots end at end_time + root = tsk_tree_get_left_root(&tree); + CU_ASSERT_EQUAL(msp.tables->nodes.time[root], t); + } + + model = msp_get_model(&msp)->type; + CU_ASSERT_EQUAL(model, models[j]); + model_name = msp_get_model_name(&msp); + CU_ASSERT_STRING_EQUAL(model_name, model_names[j]); + + ret = msp_free(&msp); + CU_ASSERT_EQUAL(ret, 0); + tsk_table_collection_free(&tables); + tsk_treeseq_free(&ts); + tsk_tree_free(&tree); + } + gsl_rng_free(rng); +} + +static void +test_single_locus_continue_after_local_mrca(void) +{ + int ret; + uint32_t n = 10; + uint32_t m = 10; + long check_time; + long seed = 10; + int model; + int models[] = { MSP_MODEL_HUDSON, MSP_MODEL_SMC, MSP_MODEL_SMC_PRIME, + MSP_MODEL_SMC_K, MSP_MODEL_DTWF }; + const char *model_names[] = { "hudson", "smc", "smc_prime", "smc_k", "dtwf" }; + const char *model_name; + size_t j; + int k; + tsk_table_collection_t tables; + msp_t msp; + tsk_treeseq_t ts; + tsk_tree_t tree; + tsk_id_t root; + + gsl_rng *rng = safe_rng_alloc(); + + gsl_rng_set(rng, seed); + + for (j = 0; j < sizeof(models) / sizeof(int); j++) { + + ret = build_sim(&msp, &tables, rng, m, 1, NULL, n); + ret = msp_set_recombination_rate(&msp, 0); + CU_ASSERT_EQUAL_FATAL(ret, 0); + + ret = msp_set_stop_at_local_mrca(&msp, false); + CU_ASSERT_EQUAL_FATAL(ret, 0); + + switch (j) { + case 0: + ret = msp_set_simulation_model_hudson(&msp); + break; + case 1: + ret = msp_set_simulation_model_smc(&msp); + break; + case 2: + ret = msp_set_simulation_model_smc_prime(&msp); + break; + case 3: + ret = msp_set_simulation_model_smc_k(&msp, 0); + break; + case 4: + ret = msp_set_simulation_model_dtwf(&msp); + break; + } + ret = msp_initialise(&msp); + CU_ASSERT_EQUAL(ret, 0); + + ret = msp_run(&msp, DBL_MAX, ULONG_MAX); + CU_ASSERT_EQUAL(ret, 0); + + ret = msp_finalise_tables(&msp); + CU_ASSERT_EQUAL(ret, 0); + + ret = tsk_treeseq_init(&ts, &tables, TSK_TS_INIT_BUILD_INDEXES); + CU_ASSERT_EQUAL_FATAL(ret, 0); + + ret = tsk_tree_init(&tree, &ts, 0); + CU_ASSERT_EQUAL_FATAL(ret, 0); + + k = 0; + for (ret = tsk_tree_first(&tree); ret == 1; ret = tsk_tree_next(&tree)) { + k++; + // verrify each tree has a single root + CU_ASSERT_EQUAL_FATAL(tsk_tree_get_num_roots(&tree), 1); + + root = tsk_tree_get_left_root(&tree); + if (k == 1) { + check_time = msp.tables->nodes.time[root]; + } else { + CU_ASSERT_EQUAL(msp.tables->nodes.time[root], check_time); + } + } + + model = msp_get_model(&msp)->type; + CU_ASSERT_EQUAL(model, models[j]); + model_name = msp_get_model_name(&msp); + CU_ASSERT_STRING_EQUAL(model_name, model_names[j]); + + ret = msp_free(&msp); + CU_ASSERT_EQUAL(ret, 0); + tsk_table_collection_free(&tables); + tsk_treeseq_free(&ts); + tsk_tree_free(&tree); + } + gsl_rng_free(rng); +} + static void test_multi_locus_bottleneck_arg(void) { @@ -4484,6 +4670,10 @@ main(int argc, char **argv) test_single_locus_historical_sample_end_time }, { "test_multi_locus_simulation", test_multi_locus_simulation }, + { "test_single_locus_continue_after_local_mrca", + test_single_locus_continue_after_local_mrca }, + { "test_single_locus_continue_after_local_mrca_with_end_time", + test_single_locus_continue_after_local_mrca_with_end_time }, { "test_multi_locus_bottleneck_arg", test_multi_locus_bottleneck_arg }, { "test_multi_locus_store_unary_simple", test_multi_locus_store_unary_simple }, diff --git a/msprime/__init__.py b/msprime/__init__.py index 63b18fe4b..188e06a65 100644 --- a/msprime/__init__.py +++ b/msprime/__init__.py @@ -48,6 +48,7 @@ SweepGenicSelection, FixedPedigree, TimeUnitsMismatchWarning, + PotentialInfiniteSimulationWarning, NodeType, ) diff --git a/msprime/_msprimemodule.c b/msprime/_msprimemodule.c index b3de9398e..952f6ef27 100644 --- a/msprime/_msprimemodule.c +++ b/msprime/_msprimemodule.c @@ -1790,8 +1790,8 @@ Simulator_init(Simulator *self, PyObject *args, PyObject *kwds) "demographic_events", "model", "avl_node_block_size", "segment_block_size", "node_mapping_block_size", "store_migrations", "start_time", "additional_nodes", "coalescing_segments_only", - "num_labels", "gene_conversion_rate", "gene_conversion_tract_length", - "discrete_genome", "ploidy", NULL}; + "num_labels", "gene_conversion_rate", "gene_conversion_tract_length", + "discrete_genome", "ploidy", "stop_at_local_mrca", NULL}; PyObject *migration_matrix = NULL; PyObject *population_configuration = NULL; PyObject *demographic_events = NULL; @@ -1813,11 +1813,13 @@ Simulator_init(Simulator *self, PyObject *args, PyObject *kwds) double gene_conversion_rate = 0; double gene_conversion_tract_length = 1.0; int ploidy = 2; + int stop_at_local_mrca = true; + self->sim = NULL; self->random_generator = NULL; if (!PyArg_ParseTupleAndKeywords(args, kwds, - "O!O!|O!O!OO!O!nnnidkinddii", kwlist, + "O!O!|O!O!OO!O!nnnidkinddiii", kwlist, &LightweightTableCollectionType, &tables, &RandomGeneratorType, &random_generator, /* optional */ @@ -1830,7 +1832,7 @@ Simulator_init(Simulator *self, PyObject *args, PyObject *kwds) &node_mapping_block_size, &store_migrations, &start_time, &additional_nodes, &coalescing_segments_only, &num_labels, &gene_conversion_rate, &gene_conversion_tract_length, - &discrete_genome, &ploidy)) { + &discrete_genome, &ploidy, &stop_at_local_mrca)) { goto out; } self->random_generator = random_generator; @@ -1951,7 +1953,8 @@ Simulator_init(Simulator *self, PyObject *args, PyObject *kwds) } msp_set_additional_nodes(self->sim, (uint32_t) additional_nodes); msp_set_coalescing_segments_only(self->sim, coalescing_segments_only); - + msp_set_stop_at_local_mrca(self->sim, stop_at_local_mrca); + sim_ret = msp_initialise(self->sim); if (sim_ret != 0) { handle_input_error("initialise", sim_ret); diff --git a/msprime/ancestry.py b/msprime/ancestry.py index 8974f8fd5..667c783bf 100644 --- a/msprime/ancestry.py +++ b/msprime/ancestry.py @@ -56,6 +56,13 @@ class TimeUnitsMismatchWarning(UserWarning): """ +class PotentialInfiniteSimulationWarning(UserWarning): + """ + Warning raised when a simulation is likely to run indefinitely: + stop_at_local_mrca false and no end_time set. + """ + + def _model_factory(model: None | str | AncestryModel) -> AncestryModel: """ Returns an AncestryModel corresponding to the specified model @@ -815,6 +822,7 @@ def _parse_sim_ancestry( coalescing_segments_only=None, num_labels=None, random_seed=None, + stop_at_local_mrca=None, init_for_debugger=False, ): """ @@ -1087,6 +1095,17 @@ def _parse_sim_ancestry( random_seed = _parse_random_seed(random_seed) random_generator = _msprime.RandomGenerator(random_seed) + if stop_at_local_mrca is None: + stop_at_local_mrca = True + if not stop_at_local_mrca: + if end_time == math.inf or end_time is None: + message = ( + "Warning: stop_at_local_mrca=False with no set end_time may result " + "in long (or infinite) run times, especially with a large recombination " + "rate. You can set an end_time to avoid this warning." + ) + warnings.warn(message, PotentialInfiniteSimulationWarning, stacklevel=2) + return Simulator( tables=initial_state, recombination_map=recombination_map, @@ -1103,6 +1122,7 @@ def _parse_sim_ancestry( end_time=end_time, num_labels=num_labels, random_generator=random_generator, + stop_at_local_mrca=stop_at_local_mrca, ) @@ -1130,6 +1150,7 @@ def sim_ancestry( num_replicates=None, replicate_index=None, record_provenance=None, + stop_at_local_mrca=None, ): """ Simulates an ancestral process described by the specified model, demography and @@ -1251,6 +1272,9 @@ def sim_ancestry( the :ref:`sec_ancestry_models_specifying` section for more details, and the :ref:`sec_ancestry_models` section for the available models and examples. + :param stop_at_local_mrca: If True (the default), the simulation will stop for a + tree when local mrca is reached. If False, simulations will continue on the path + to the grand mrca. :type model: str or msprime.AncestryModel or list :return: The :class:`tskit.TreeSequence` object representing the results of the simulation if no replication is performed, or an @@ -1289,6 +1313,7 @@ def sim_ancestry( coalescing_segments_only=coalescing_segments_only, num_labels=num_labels, random_seed=random_seed, + stop_at_local_mrca=stop_at_local_mrca, # num_replicates is excluded as provenance is per replicate # replicate index is excluded as it is inserted for each replicate ) @@ -1313,6 +1338,7 @@ def sim_ancestry( coalescing_segments_only=coalescing_segments_only, num_labels=num_labels, random_seed=random_seed, + stop_at_local_mrca=stop_at_local_mrca, ) return _wrap_replicates( sim, @@ -1398,6 +1424,7 @@ def __init__( start_time=None, end_time=None, num_labels=None, + stop_at_local_mrca=True, ): # We always need at least n segments, so no point in making # allocation any smaller than this. @@ -1448,6 +1475,7 @@ def __init__( gene_conversion_tract_length=gene_conversion_tract_length, discrete_genome=discrete_genome, ploidy=ploidy, + stop_at_local_mrca=stop_at_local_mrca, ) # Highlevel attributes used externally that have no lowlevel equivalent self.end_time = np.inf if end_time is None else end_time diff --git a/tests/test_algorithms.py b/tests/test_algorithms.py index f547a1996..9e5143c5e 100644 --- a/tests/test_algorithms.py +++ b/tests/test_algorithms.py @@ -267,6 +267,30 @@ def test_store_ca_re(self): assert np.sum(node_flags == msprime.NODE_IS_CA_EVENT) > 0 assert np.sum(node_flags == msprime.NODE_IS_RE_EVENT) > 0 + def test_stop_at_local_mrca_end_time(self): + end_time = 100 + r = 0.1 + ts = self.run_script( + f"10 --continue-after-local-mrca --end-time={end_time} -r {r}" + ) + + # test that simulations continue for all trees to the end_time + old_time = None + for tree in ts.trees(): + assert len(tree.roots) == 1 # otherwise the test is not valid + u = tree.roots[0] + assert tree.time(u) >= end_time + if old_time is None: + old_time = tree.time(u) + # check that end time is the same for all roots + assert tree.time(u) == old_time + + def test_stop_at_local_mrca(self): + ts = self.run_script("10 --continue-after-local-mrca -r 0") + + root_times = [tree.time(tree.root) for tree in ts.trees()] + assert len(set(root_times)) == 1 + def test_gc(self): ts = self.run_script("10 -c 0.4 2 -d") assert ts.num_trees > 1 diff --git a/tests/test_ancestry.py b/tests/test_ancestry.py index 7b0510969..b7cd2184b 100644 --- a/tests/test_ancestry.py +++ b/tests/test_ancestry.py @@ -891,6 +891,249 @@ def test_str(self): assert len(s) > 0 +class TestSimulateAfterLocalMRCA: + + def check_roots(self, ts, end_time=None, allow_multiple_roots=False): + if allow_multiple_roots: + root_times = [tree.time(tree.roots[0]) for tree in ts.trees()] + else: + root_times = [tree.time(tree.root) for tree in ts.trees()] + assert len(set(root_times)) == 1 + if end_time is not None: + assert root_times[0] >= end_time + + def verify_simulation(self, n, m, r): + """ + Verifies a simulation for the specified parameters. + """ + end_time = 1000 + sim = ancestry._parse_sim_ancestry( + n, + sequence_length=m, + recombination_rate=r, + stop_at_local_mrca=False, + end_time=end_time, + ) + sim.run() + assert sim.num_breakpoints == len(sim.breakpoints) + assert sim.time > 0 + assert sim.num_avl_node_blocks > 0 + assert sim.num_segment_blocks > 0 + assert sim.num_node_mapping_blocks > 0 + tree_sequence = tskit.TableCollection.fromdict( + sim.tables.asdict() + ).tree_sequence() + t = 0.0 + for record in tree_sequence.nodes(): + if record.time > t: + t = record.time + assert t >= end_time + assert sim.time == t + assert sim.num_common_ancestor_events > 0 + assert sim.num_recombination_events >= 0 + assert np.sum(sim.num_migration_events) >= 0 + assert sim.num_multiple_recombination_events >= 0 + self.check_roots(tree_sequence, end_time) + + def test_random_parameters(self): + num_random_sims = 10 + for _ in range(num_random_sims): + n = random.randint(2, 100) + m = random.randint(10, 100) + r = random.random() + self.verify_simulation(n, m, r) + + @pytest.mark.parametrize( + "model", + [ + "dtwf", + "hudson", + msprime.BetaCoalescent(alpha=1.5), + msprime.DiracCoalescent(psi=0.1, c=2), + ], + ) + def test_single_models(self, model): + """ + Tests that simulations run after the local MRCA when the flag is set + """ + + end_time = 1000 + tss = msprime.sim_ancestry( + 10, + population_size=10, + stop_at_local_mrca=False, + end_time=end_time, + random_seed=1, + sequence_length=10, + recombination_rate=0.1, + model=model, + num_replicates=10, + ) + for ts in tss: + self.check_roots(ts, end_time) + + def test_fixed_pedigree(self): + """ + To test pedigree model, we make diamond shaped pedigree and make + sure that simulations run until the eldest individual. + There is a chance that no coalescence happens until the eldest individual, + so, we use longer sequence with a high recombination rate and confirm that + all trees does not end before the eldest individual. + + The pedigree used has this shape: + O + | + O + / \ + O O + \\ // + O + / \ + O O + \\ // + O + """ + eldest_time = 10 + pb = msprime.PedigreeBuilder() + + gggpa = pb.add_individual(time=eldest_time) + ggpa = pb.add_individual(time=4, parents=[gggpa, gggpa]) + + # first diamound + gmom = pb.add_individual(time=3, parents=[ggpa, ggpa]) + gdad = pb.add_individual(time=3, parents=[ggpa, ggpa]) + grandpa_id = pb.add_individual(time=2, parents=[gmom, gdad]) + + # second diamond + mom_id = pb.add_individual(time=1, parents=[grandpa_id, grandpa_id]) + dad_id = pb.add_individual(time=1, parents=[grandpa_id, grandpa_id]) + pb.add_individual(time=0, parents=[mom_id, dad_id], is_sample=True) + + pedigree = pb.finalise(100) + + ts = msprime.sim_ancestry( + initial_state=pedigree, + stop_at_local_mrca=False, + recombination_rate=0.1, + end_time=eldest_time, + model="fixed_pedigree", + ) + + self.check_roots(ts, eldest_time, allow_multiple_roots=True) + + def test_multiple_models(self): + """ + To test multiple models, we run simulations with end time < first models + duration and test (that would be using the first model). Then we run again + until end_time too large to cover the second model. + """ + end_time = 499 + ts = msprime.sim_ancestry( + 2, + population_size=10, + model=[ + msprime.DiscreteTimeWrightFisher(duration=500), + msprime.StandardCoalescent(), + ], + random_seed=2, + recombination_rate=0.1, + sequence_length=10, + stop_at_local_mrca=False, + end_time=end_time, + ) + + self.check_roots(ts, end_time, allow_multiple_roots=True) + + end_time = 100000 + ts = msprime.sim_ancestry( + 2, + population_size=10, + model=[ + msprime.DiscreteTimeWrightFisher(duration=500), + msprime.StandardCoalescent(), + ], + random_seed=2, + recombination_rate=0.1, + sequence_length=10, + stop_at_local_mrca=False, + end_time=end_time, + ) + self.check_roots(ts, end_time, allow_multiple_roots=True) + + @pytest.mark.parametrize( + "model", + [ + "dtwf", + "hudson", + msprime.BetaCoalescent(alpha=1.5), + msprime.DiracCoalescent(psi=0.1, c=2), + ], + ) + def test_no_end_time(self, model): + with pytest.warns(msprime.PotentialInfiniteSimulationWarning): + ts = msprime.sim_ancestry( + 10, + stop_at_local_mrca=False, + population_size=10, + random_seed=1, + sequence_length=2, + recombination_rate=0, + model=model, + ) + self.check_roots(ts) + + @pytest.mark.parametrize( + "model", + [ + "dtwf", + "hudson", + ], + ) + def test_migration(self, model): + end_time = 10000 + demography = msprime.Demography.stepping_stone_model( + [100, 100], migration_rate=0.01 + ) + ts = msprime.sim_ancestry( + {0: 2, 1: 1}, + demography=demography, + stop_at_local_mrca=False, + end_time=end_time, + recombination_rate=0.01, + sequence_length=10, + model=model, + ) + self.check_roots(ts, end_time, allow_multiple_roots=True) + + @pytest.mark.parametrize( + "model", + [ + "dtwf", + "hudson", + ], + ) + def test_additional_nodes(self, model): + + end_time = 1000 + ts = msprime.sim_ancestry( + 2, + population_size=10, + random_seed=2, + additional_nodes=( + msprime.NodeType.RECOMBINANT + | msprime.NodeType.PASS_THROUGH + | msprime.NodeType.COMMON_ANCESTOR + ), + coalescing_segments_only=False, + stop_at_local_mrca=False, + sequence_length=5, + end_time=end_time, + model=model, + ) + + self.check_roots(ts, end_time) + + class TestParseRandomSeed: """ Tests for parsing the random seed values. From 96df64f72fec143e5da000ba63733b6320ed9191 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Tue, 2 Dec 2025 11:44:52 +0000 Subject: [PATCH 2/4] Change stopping condition for stop_at_local_mrca=False --- algorithms.py | 71 ++++---------- docs/ancestry.md | 79 ++++++++-------- lib/msprime.c | 132 +++++++++++--------------- lib/msprime.h | 1 - lib/tests/test_ancestry.c | 192 ++++++++++++++++++++------------------ msprime/__init__.py | 1 - msprime/ancestry.py | 15 --- tests/test_algorithms.py | 62 ++++++------ tests/test_ancestry.py | 90 ++---------------- 9 files changed, 253 insertions(+), 390 deletions(-) diff --git a/algorithms.py b/algorithms.py index 8521793de..26369c309 100644 --- a/algorithms.py +++ b/algorithms.py @@ -1049,13 +1049,6 @@ def initialise(self, ts): seg = seg.next self.add_lineage(lineage) - def ancestors_remain(self): - """ - Returns True if the simulation is not finished, i.e., there is some ancestral - material that has not fully coalesced. - """ - return sum(pop.get_num_ancestors() for pop in self.P) != 0 - def change_population_size(self, pop_id, size): self.P[pop_id].set_start_size(size) @@ -1314,6 +1307,12 @@ def find_cleft_individual(self, label, cleft_value): individual_index -= num_ancestors raise AssertionError() + def is_completed(self): + for x in self.S.values(): + if x > 1: + return False + return True + def hudson_simulate(self, end_time): """ Simulates the algorithm until all loci have coalesced. @@ -1322,7 +1321,7 @@ def hudson_simulate(self, end_time): potential_destinations = self.get_potential_destinations() # only worried about label 0 below - while len(non_empty_pops) > 0: + while not self.is_completed(): self.verify() if self.t >= end_time: break @@ -1367,9 +1366,7 @@ def hudson_simulate(self, end_time): mig_source = j mig_dest = k min_time = min(t_re, t_ca, t_gcin, t_gc_left, t_mig) - assert (min_time != INFINITY) or ( - (min_time == INFINITY) and not self.stop_at_local_mrca - ) + assert min_time != INFINITY if self.t + min_time > self.modifier_events[0][0]: t, func, args = self.modifier_events.pop(0) self.t = t @@ -1383,14 +1380,7 @@ def hudson_simulate(self, end_time): event = "MOD" else: self.t += min_time - if min_time >= INFINITY: - assert not self.stop_at_local_mrca - event = "END" - if end_time >= INFINITY: - end_time = self.t - else: - self.t = end_time - elif min_time == t_re: + if min_time == t_re: event = "RE" self.hudson_recombination_event(0) elif min_time == t_gcin: @@ -1453,7 +1443,7 @@ def single_sweep_simulate(self): # main loop time t_inc_orig = self.time_slice e_time = 0.0 - while self.ancestors_remain() and sweep_traj_step < len(times) - 1: + while sweep_traj_step < len(times) - 1 and not self.is_completed(): self.verify() event_prob = 1.0 while event_prob > random.random() and sweep_traj_step < len(times) - 1: @@ -1536,28 +1526,17 @@ def pedigree_simulate(self): self.pedigree = Pedigree(self.tables) self.dtwf_climb_pedigree() - def dtwf_above_root_simulations_stop(self, events_happened): - return ( - not events_happened - and not self.stop_at_local_mrca - and sum(pop.get_num_ancestors() for pop in self.P) == 1 - ) - def dtwf_simulate(self, end_time): """ Simulates the algorithm until all loci have coalesced. """ - while self.ancestors_remain() and self.t < end_time: - self.t += 1 + while not self.is_completed(): self.verify() - events_happened = self.dtwf_generation() - - if self.dtwf_above_root_simulations_stop(events_happened): - if end_time >= INFINITY: - end_time = self.t - else: - self.t = end_time + if self.t >= end_time: break + self.t += 1 + # print("DTWF", self.t) + self.dtwf_generation() def dtwf_generation(self): """ @@ -1565,9 +1544,6 @@ def dtwf_generation(self): Returns True if any events occurred, False otherwise. """ # Migration events happen at the rates in the matrix. - nodes = self.tables.nodes.num_rows - edges = self.tables.edges.num_rows - no_migration = True for j in range(len(self.P)): source_size = self.P[j].get_num_ancestors() @@ -1580,7 +1556,6 @@ def dtwf_generation(self): mig_source = j mig_dest = k self.migration_event(mig_source, mig_dest) - no_migration = False for pop_idx, pop in enumerate(self.P): # Cluster haploid inds by parent @@ -1637,14 +1612,6 @@ def dtwf_generation(self): h, pop_idx, 0, parent_nodes[ploid] ) # label 0 only self.verify() - if ( - self.tables.nodes.num_rows == nodes - and self.tables.edges.num_rows == edges - and no_migration - ): - # dtwf generation had no events - return False - return True def process_pedigree_common_ancestors(self, ind, ploid): """ @@ -2305,7 +2272,7 @@ def merge_ancestors(self, H, pop_id, label, new_node_id=-1): else: right = left while right < r_max and self.S[right] != min_overlap: - self.S[right] -= min_overlap - 1 + self.S[right] -= len(X) - 1 right = self.S.succ_key(right) alpha = self.alloc_segment(left, right, new_node_id, pop_id) # Update the heaps and make the record. @@ -2775,10 +2742,10 @@ def verify(self): Checks that the state of the simulator is consistent. """ self.verify_segments() - if self.model not in ["fixed_pedigree", "dtwf"]: - # The fixed_pedigree model doesn't maintain a bunch of stuff. - # It would probably be simpler if it did. + # The fixed_pedigree model doesn't maintain a bunch of stuff. + if self.model != "fixed_pedigree": self.verify_overlaps() + if self.model not in ["fixed_pedigree", "dtwf"]: for label in range(self.num_labels): if self.recomb_mass_index is None: assert self.recomb_map.total_mass == 0 diff --git a/docs/ancestry.md b/docs/ancestry.md index 12f4ac4dc..07df54724 100644 --- a/docs/ancestry.md +++ b/docs/ancestry.md @@ -1129,47 +1129,6 @@ However, if `coalescing_segments_only=False`, then the edges recorded would be from `m` to `p` over the entire segment `[a, b)`, and from `n` to `p` over the entire segment `[c,d)`. The nodes `m` and `n` coalesce (in `p`) on only the overlapping segment `[c,b)`, and so the node `p` will be a unary node in the flanking regions: above `m` on the segment `[a,c)` and above `n` on the segment `[b,d)`. -(sec_simulation_after_mrca)= - -### Simulations after local MRCA - -By default, msprime stops simulating local trees when a local most recent common ancestor (MRCA) is found. This is because events that occur - above the common ancestor are shared across all samples, and we are usually interested in differences between samples. -```{code-cell} -import msprime - -ts = msprime.sim_ancestry(2, recombination_rate=0.1, sequence_length=10, random_seed=21) -ts.draw_svg(time_scale='rank') -``` -Here, for example, we see that simulations in the middle tree stopped at node 6, which is much younger than nodes 7 and 8. -And if we look closely, we can also see that simulations stopped at node 7 in the rightmost tree at a younger time than -node 8 in the leftmost tree. - - -However, for some specialised applications, simulations after the local MRCA might be needed. In this case, we set the -parameter `stop_at_local_mrca` of {func}`.sim_ancestry` to `False`, and we set the `end_time` parameter to stop simulations -when they reach this end time. - - -```{code-cell} -import msprime - -ts = msprime.sim_ancestry( - 2, recombination_rate=0.1, - sequence_length=10, - random_seed=21, - stop_at_local_mrca=False, - end_time=10) - -ts.draw_svg(time_scale='rank') -``` -It is possible to run the simulations without setting an `end_time`. In that case, simulations will -stop when no more events (recombination, coalescence, migration, etc.) could happen. However, since in many cases this state is impossible to reach, we recommend setting an `end_time`. - -Note that you might want to set an `end_time` that is larger than possible end times for your simulations to avoid stopping -your simulations before full coalescence. However, a larger `end_time` means that simulations will take longer and consume -more memory. - (sec_additional_nodes_ca)= @@ -1522,6 +1481,44 @@ Migrations nodes are also recorded in the ARG using the `NodeType.MIGRANT` flag. See the {ref}`sec_api_node_flags` section for more details. +(sec_simulation_after_mrca)= + +### Simulations after local MRCA + +By default, msprime stops simulating local trees when a local most recent +common ancestor (MRCA) is found. This is because events that occur above the +common ancestor are shared across all samples, and we are usually interested in +differences between samples. + +However, for some specialised applications, simulations after the local MRCA +might be needed. In this case, we set the parameter `stop_at_local_mrca` of +{func}`.sim_ancestry` to `False`. + +```{code-cell} +ts = msprime.sim_ancestry( + 2, recombination_rate=0.1, sequence_length=10, random_seed=75, + stop_at_local_mrca=False, record_full_arg=True) +ts.draw_svg(time_scale='rank') +``` + +In this example we've continued the simulation *after* we've found roots for the local +trees, and we can see, for example, that node 14 is still the ancestor of 11 +in the rightmost tree. This would be not recorded with the default ``stop_at_local_mrca=False`` +because 11 is the root of that tree. + +The stopping condition remains the same---we keep simulating until an MRCA has been +found at *all* local trees. The interpretation of the the final "root" nodes +is a little subtle + +:::{note} +When using the ``stop_at_local_mrca=False`` option you usually need to specify an +``end_time`` argument (see {ref}`sec_ancestry_end_time`) or the simulation will +run indefinitely! It is in general not possible to know what the "right" end +time for a given simulation is, and you may need to specify a "large" value to be +sure that all trees have coalesced locally. +::: + + ## Manipulating simulation time (sec_ancestry_end_time)= diff --git a/lib/msprime.c b/lib/msprime.c index bc23cb306..bec53771d 100644 --- a/lib/msprime.c +++ b/lib/msprime.c @@ -1038,7 +1038,6 @@ msp_alloc(msp_t *self, tsk_table_collection_t *tables, gsl_rng *rng) self->additional_nodes = 0; self->coalescing_segments_only = true; self->stop_at_local_mrca = true; - self->dtwf_no_events_generation = false; self->avl_node_block_size = 1024; self->node_mapping_block_size = 1024; self->segment_block_size = 1024; @@ -1107,17 +1106,6 @@ msp_alloc_memory_blocks(msp_t *self) return ret; } -/* - * Returns true if the simulation has completed. - */ -int -msp_is_completed(msp_t *self) -{ - size_t n = msp_get_num_ancestors(self); - - return self->state == MSP_STATE_SIMULATING && n == 0; -} - int msp_free(msp_t *self) { @@ -3786,7 +3774,7 @@ msp_merge_two_ancestors(msp_t *self, population_id_t population_id, label_id_t l bool defrag_required = false; tsk_id_t v; double l, r, l_min, r_max; - uint32_t min_overlap; + uint32_t min_overlap = self->stop_at_local_mrca ? 2 : 0; avl_node_t *node; node_mapping_t *nm, search; segment_t *x, *y, *z, *alpha, *beta; @@ -3797,12 +3785,6 @@ msp_merge_two_ancestors(msp_t *self, population_id_t population_id, label_id_t l goto out; } - if (self->stop_at_local_mrca) { - min_overlap = 2; - } else { - min_overlap = 0; - } - x = a; y = b; /* Keep GCC happy */ @@ -3997,10 +3979,8 @@ msp_merge_ancestors(msp_t *self, avl_tree_t *Q, population_id_t population_id, int ret = MSP_ERR_GENERIC; bool coalescence = false; bool defrag_required = false; - uint32_t j, h; + uint32_t j, h, min_overlap; double l, r, r_max, next_l, l_min; - uint32_t min_overlap; - avl_node_t *node; node_mapping_t *nm, search; segment_t *x, *z, *alpha; @@ -4090,11 +4070,7 @@ msp_merge_ancestors(msp_t *self, avl_tree_t *Q, population_id_t population_id, node = avl_search(&self->overlap_counts, &search); tsk_bug_assert(node != NULL); nm = (node_mapping_t *) node->item; - if (self->stop_at_local_mrca) { - min_overlap = h; - } else { - min_overlap = 0; - } + min_overlap = self->stop_at_local_mrca ? h : 0; if (nm->value == min_overlap) { nm->value = 0; node = node->next; @@ -4104,7 +4080,7 @@ msp_merge_ancestors(msp_t *self, avl_tree_t *Q, population_id_t population_id, } else { r = l; while (nm->value != min_overlap && r < r_max) { - nm->value -= min_overlap - 1; + nm->value -= h - 1; node = node->next; tsk_bug_assert(node != NULL); nm = (node_mapping_t *) node->item; @@ -5142,7 +5118,7 @@ msp_run_coalescent(msp_t *self, double max_time, unsigned long max_events) goto out; } - while (msp_get_num_ancestors(self) > 0) { + while (!msp_is_completed(self)) { if (events == max_events) { ret = MSP_EXIT_MAX_EVENTS; break; @@ -5212,20 +5188,10 @@ msp_run_coalescent(msp_t *self, double max_time, unsigned long max_events) t_wait = GSL_MIN(mig_t_wait, GSL_MIN(gc_t_wait, GSL_MIN(gc_left_t_wait, GSL_MIN(re_t_wait, ca_t_wait)))); - if (fixed_event_time == DBL_MAX && t_wait == DBL_MAX - && self->stop_at_local_mrca) { // TODO add condition + if (fixed_event_time == DBL_MAX && t_wait == DBL_MAX) { + msp_print_state(self, stdout); ret = MSP_ERR_INFINITE_WAITING_TIME; goto out; - } else if ((fixed_event_time == DBL_MAX && t_wait == DBL_MAX) - && !self->stop_at_local_mrca) { - - if (max_time >= DBL_MAX) { - max_time = self->time; - } else { - self->time = max_time; - } - ret = MSP_EXIT_COALESCENCE; - break; } random_event_time = self->time + t_wait; @@ -5476,10 +5442,6 @@ msp_dtwf_generation(msp_t *self) label_id_t label = 0; tsk_id_t parent_nodes[MSP_MAX_PED_PLOIDY]; - unsigned long int edgs = self->tables->edges.num_rows; - unsigned long int nodes = self->tables->nodes.num_rows; - - self->dtwf_no_events_generation = false; for (i = 0; i < 2; i++) { avl_init_tree(&Q[i], cmp_segment_queue, NULL); } @@ -5573,15 +5535,6 @@ msp_dtwf_generation(msp_t *self) } } - if (self->tables->edges.num_rows == edgs - && self->tables->nodes.num_rows == nodes) { - // No edges or nodes added this generation - // i.e., no coalescence or recombination - self->dtwf_no_events_generation = true; - } else { - edgs = self->tables->edges.num_rows; - nodes = self->tables->nodes.num_rows; - } free(parents); free(segment_mem); segment_mem = NULL; @@ -5666,7 +5619,6 @@ msp_run_dtwf(msp_t *self, double max_time, unsigned long max_events) avl_tree_t *nodes; /* Only support a single structured coalescent label at the moment */ label_id_t label = 0; - bool migration_occurred; tsk_bug_assert(self->recomb_mass_index == NULL); tsk_bug_assert(self->gc_mass_index == NULL); @@ -5687,7 +5639,7 @@ msp_run_dtwf(msp_t *self, double max_time, unsigned long max_events) goto out; } - while (msp_get_num_ancestors(self) > 0) { + while (!msp_is_completed(self)) { if (events == max_events) { ret = MSP_EXIT_MAX_EVENTS; break; @@ -5699,8 +5651,6 @@ msp_run_dtwf(msp_t *self, double max_time, unsigned long max_events) } self->time++; - migration_occurred = false; - /* Following SLiM, we perform migrations prior to selecting * parents for the current generation */ node_trees @@ -5750,7 +5700,6 @@ msp_run_dtwf(msp_t *self, double max_time, unsigned long max_events) mig_dest_pop = (population_id_t) k; for (i = 0; i < n[k]; i++) { - migration_occurred = true; ret = msp_store_simultaneous_migration_events( self, nodes, mig_source_pop, label); if (ret != 0) { @@ -5767,7 +5716,6 @@ msp_run_dtwf(msp_t *self, double max_time, unsigned long max_events) nodes = &node_trees[j * self->num_populations + k]; mig_source_pop = (population_id_t) j; mig_dest_pop = (population_id_t) k; - migration_occurred = true; ret = msp_simultaneous_migration_event( self, nodes, mig_source_pop, mig_dest_pop); if (ret != 0) { @@ -5791,7 +5739,6 @@ msp_run_dtwf(msp_t *self, double max_time, unsigned long max_events) } ret = msp_apply_demographic_events(self, self->next_demographic_event->time); // Demographic events can change population structure - migration_occurred = true; if (ret != 0) { goto out; } @@ -5813,18 +5760,6 @@ msp_run_dtwf(msp_t *self, double max_time, unsigned long max_events) } self->next_sampling_event++; } - - if (self->dtwf_no_events_generation && (!migration_occurred) - && (!self->stop_at_local_mrca) && msp_get_num_ancestors(self) == 1) { - /* No events occurred: we are done */ - if (max_time >= DBL_MAX) { - max_time = self->time; - } else { - self->time = max_time; - } - ret = MSP_EXIT_COALESCENCE; - break; - } } out: msp_safe_free(node_trees); @@ -6009,7 +5944,7 @@ msp_run_sweep(msp_t *self) } curr_step = 1; - while (msp_get_num_ancestors(self) > 0 && curr_step < num_steps) { + while (curr_step < num_steps && !msp_is_completed(self)) { events++; /* Set pop sizes & rec_rates */ for (j = 0; j < self->num_labels; j++) { @@ -6273,7 +6208,7 @@ msp_finalise_tables(msp_t *self) tsk_bookmark_t bookmark; /* We don't want to add unary edges for the pedigree simulation model */ - if (!msp_is_completed(self) && self->model.type != MSP_MODEL_WF_PED) { + if (self->model.type != MSP_MODEL_WF_PED) { ret = msp_insert_uncoalesced_edges(self); if (ret != 0) { goto out; @@ -6574,6 +6509,45 @@ msp_get_time(msp_t *self) return self->time; } +/* + * Returns true if the simulation has completed. + */ +int +msp_is_completed(msp_t *self) +{ + /* int model = self->model.type; */ + bool completed; + avl_node_t *node; + node_mapping_t *nm; + + /* NOTE! This is a quick first implementation. This could be a performance + * bottleneck and maybe we keep track of whether there's any non-1 values + * instead + */ + completed = true; + for (node = self->overlap_counts.head; node->next != NULL; node = node->next) { + nm = (node_mapping_t *) node->item; + if (nm->value > 1) { + completed = false; + break; + } + } + + /* } */ + + /* if (model == MSP_MODEL_HUDSON || model == MSP_MODEL_SMC_K || model == + * MSP_MODEL_BETA */ + /* || model == MSP_MODEL_DIRAC || model == MSP_MODEL_SMC */ + /* || model == MSP_MODEL_SMC_PRIME */ + /* || model == MSP_MODEL_DTWF) { */ + /* completed = msp_coalescent_is_complete(self); */ + /* } else { */ + /* completed = msp_get_num_ancestors(self) == 0; */ + /* } */ + + return self->state == MSP_STATE_SIMULATING && completed; +} + /* Demographic events. All times and input parameters are specified in units * of generations. When we store these values, we must rescale them into * model time, as appropriate. */ @@ -8109,9 +8083,11 @@ msp_beta_common_ancestor_event(msp_t *self, population_id_t pop_id, label_id_t l * merged. */ for (j = 0; j < num_parental_copies; j++) { - ret = msp_merge_ancestors(self, &Q[j], pop_id, label, TSK_NULL, NULL); - if (ret < 0) { - goto out; + if (avl_count(&Q[j]) > 0) { + ret = msp_merge_ancestors(self, &Q[j], pop_id, label, TSK_NULL, NULL); + if (ret < 0) { + goto out; + } } } } diff --git a/lib/msprime.h b/lib/msprime.h index 30ec2d4d4..25c5da47b 100644 --- a/lib/msprime.h +++ b/lib/msprime.h @@ -257,7 +257,6 @@ typedef struct _msp_t { uint32_t additional_nodes; bool coalescing_segments_only; bool stop_at_local_mrca; - bool dtwf_no_events_generation; double sequence_length; bool discrete_genome; rate_map_t recomb_map; diff --git a/lib/tests/test_ancestry.c b/lib/tests/test_ancestry.c index 751ec0ce0..cd51c54e8 100644 --- a/lib/tests/test_ancestry.c +++ b/lib/tests/test_ancestry.c @@ -96,6 +96,78 @@ test_single_locus_simulation(void) tsk_table_collection_free(&tables); } +static void +check_single_locus_event_by_event(int model, bool stop_at_local_mrca) +{ + int ret; + uint32_t n = 10; + uint32_t m = 100; + long seed = 103; + double t; + tsk_table_collection_t tables; + msp_t msp; + tsk_treeseq_t ts; + tsk_tree_t tree; + tsk_id_t root; + gsl_rng *rng = safe_rng_alloc(); + + gsl_rng_set(rng, seed); + ret = build_sim(&msp, &tables, rng, m, 1, NULL, n); + CU_ASSERT_EQUAL_FATAL(ret, 0); + set_simulation_model(&msp, model); + ret = msp_initialise(&msp); + CU_ASSERT_EQUAL_FATAL(ret, 0); + ret = msp_set_stop_at_local_mrca(&msp, stop_at_local_mrca); + CU_ASSERT_EQUAL_FATAL(ret, 0); + + while ((ret = msp_run(&msp, DBL_MAX, 1)) == MSP_EXIT_MAX_EVENTS) { + /* msp_print_state(&msp, stdout); */ + msp_verify(&msp, 0); + CU_ASSERT_FALSE(msp_is_completed(&msp)); + } + CU_ASSERT_EQUAL_FATAL(ret, 0); + CU_ASSERT_TRUE(msp_is_completed(&msp)); + + t = msp_get_time(&msp); + ret = msp_finalise_tables(&msp); + CU_ASSERT_EQUAL(ret, 0); + + ret = tsk_treeseq_init(&ts, &tables, TSK_TS_INIT_BUILD_INDEXES); + CU_ASSERT_EQUAL_FATAL(ret, 0); + CU_ASSERT_EQUAL_FATAL(tsk_treeseq_get_num_trees(&ts), 1); + ret = tsk_tree_init(&tree, &ts, 0); + CU_ASSERT_EQUAL_FATAL(ret, 0); + ret = tsk_tree_first(&tree); + CU_ASSERT_EQUAL_FATAL(ret, TSK_TREE_OK); + CU_ASSERT_EQUAL_FATAL(tsk_tree_get_num_roots(&tree), 1); + root = tsk_tree_get_left_root(&tree); + CU_ASSERT_EQUAL(tables.nodes.time[root], t); + + CU_ASSERT_EQUAL(model, msp_get_model(&msp)->type); + + ret = msp_free(&msp); + CU_ASSERT_EQUAL(ret, 0); + tsk_tree_free(&tree); + tsk_treeseq_free(&ts); + tsk_table_collection_free(&tables); + gsl_rng_free(rng); +} + +static void +test_single_locus_event_by_event(void) +{ + size_t j; + int models[] = { MSP_MODEL_HUDSON, MSP_MODEL_SMC, MSP_MODEL_SMC_K, MSP_MODEL_DTWF, + MSP_MODEL_DIRAC, MSP_MODEL_BETA }; + int stop_at_local_mrca; + + for (j = 0; j < sizeof(models) / sizeof(int); j++) { + for (stop_at_local_mrca = 0; stop_at_local_mrca < 2; stop_at_local_mrca++) { + check_single_locus_event_by_event(models[j], stop_at_local_mrca); + } + } +} + static void test_single_locus_two_populations(void) { @@ -667,87 +739,58 @@ test_multi_locus_simulation(void) } static void -test_single_locus_continue_after_local_mrca_with_end_time(void) +test_multi_locus_continue_after_local_mrca(void) { int ret; uint32_t n = 10; - uint32_t m = 10; - long t = 100000.0; // end time - long seed = 10; + uint32_t m = 100; + long seed = 11; + double t; int model; - int models[] = { MSP_MODEL_HUDSON, MSP_MODEL_SMC, MSP_MODEL_SMC_PRIME, - MSP_MODEL_SMC_K, MSP_MODEL_DTWF }; - const char *model_names[] = { "hudson", "smc", "smc_prime", "smc_k", "dtwf" }; - const char *model_name; + int models[] = { MSP_MODEL_HUDSON, MSP_MODEL_SMC, MSP_MODEL_SMC_K, MSP_MODEL_DTWF, + MSP_MODEL_BETA, MSP_MODEL_DIRAC }; size_t j; tsk_table_collection_t tables; msp_t msp; tsk_treeseq_t ts; tsk_tree_t tree; tsk_id_t root; - gsl_rng *rng = safe_rng_alloc(); - gsl_rng_set(rng, seed); - for (j = 0; j < sizeof(models) / sizeof(int); j++) { + gsl_rng_set(rng, seed); ret = build_sim(&msp, &tables, rng, m, 1, NULL, n); + CU_ASSERT_EQUAL_FATAL(ret, 0); ret = msp_set_recombination_rate(&msp, 0.01); CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = msp_set_stop_at_local_mrca(&msp, false); CU_ASSERT_EQUAL_FATAL(ret, 0); - switch (j) { - case 0: - ret = msp_set_simulation_model_hudson(&msp); - break; - case 1: - ret = msp_set_simulation_model_smc(&msp); - break; - case 2: - ret = msp_set_simulation_model_smc_prime(&msp); - break; - case 3: - ret = msp_set_simulation_model_smc_k(&msp, 0); - break; - case 4: - ret = msp_set_simulation_model_dtwf(&msp); - break; - } + set_simulation_model(&msp, models[j]); + ret = msp_initialise(&msp); CU_ASSERT_EQUAL(ret, 0); - - ret = msp_run(&msp, t, ULONG_MAX); - CU_ASSERT(ret == 0 || ret == 2); // exit on max time or simulations done - + ret = msp_run(&msp, DBL_MAX, ULONG_MAX); + CU_ASSERT(ret == MSP_EXIT_COALESCENCE); + t = msp_get_time(&msp); ret = msp_finalise_tables(&msp); CU_ASSERT_EQUAL(ret, 0); - ret = msp_get_time(&msp); - CU_ASSERT_EQUAL(ret, t); - + /* msp_print_state(&msp, stdout); */ ret = tsk_treeseq_init(&ts, &tables, TSK_TS_INIT_BUILD_INDEXES); CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = tsk_tree_init(&tree, &ts, 0); CU_ASSERT_EQUAL_FATAL(ret, 0); - + /* printf("%d %d\n", models[j], (int) tsk_treeseq_get_num_trees(&ts)); */ + CU_ASSERT_FATAL(tsk_treeseq_get_num_trees(&ts) > 1); for (ret = tsk_tree_first(&tree); ret == 1; ret = tsk_tree_next(&tree)) { - - // verrify each tree has a single root CU_ASSERT_EQUAL_FATAL(tsk_tree_get_num_roots(&tree), 1); - - // verify that all roots end at end_time root = tsk_tree_get_left_root(&tree); - CU_ASSERT_EQUAL(msp.tables->nodes.time[root], t); + CU_ASSERT_EQUAL(tables.nodes.time[root], t); } - model = msp_get_model(&msp)->type; CU_ASSERT_EQUAL(model, models[j]); - model_name = msp_get_model_name(&msp); - CU_ASSERT_STRING_EQUAL(model_name, model_names[j]); ret = msp_free(&msp); CU_ASSERT_EQUAL(ret, 0); @@ -764,15 +807,11 @@ test_single_locus_continue_after_local_mrca(void) int ret; uint32_t n = 10; uint32_t m = 10; - long check_time; long seed = 10; int model; - int models[] = { MSP_MODEL_HUDSON, MSP_MODEL_SMC, MSP_MODEL_SMC_PRIME, - MSP_MODEL_SMC_K, MSP_MODEL_DTWF }; - const char *model_names[] = { "hudson", "smc", "smc_prime", "smc_k", "dtwf" }; - const char *model_name; + int models[] = { MSP_MODEL_HUDSON, MSP_MODEL_SMC, MSP_MODEL_SMC_K, MSP_MODEL_DTWF, + MSP_MODEL_BETA }; size_t j; - int k; tsk_table_collection_t tables; msp_t msp; tsk_treeseq_t ts; @@ -792,56 +831,30 @@ test_single_locus_continue_after_local_mrca(void) ret = msp_set_stop_at_local_mrca(&msp, false); CU_ASSERT_EQUAL_FATAL(ret, 0); - switch (j) { - case 0: - ret = msp_set_simulation_model_hudson(&msp); - break; - case 1: - ret = msp_set_simulation_model_smc(&msp); - break; - case 2: - ret = msp_set_simulation_model_smc_prime(&msp); - break; - case 3: - ret = msp_set_simulation_model_smc_k(&msp, 0); - break; - case 4: - ret = msp_set_simulation_model_dtwf(&msp); - break; - } + set_simulation_model(&msp, models[j]); + ret = msp_initialise(&msp); CU_ASSERT_EQUAL(ret, 0); - ret = msp_run(&msp, DBL_MAX, ULONG_MAX); - CU_ASSERT_EQUAL(ret, 0); - + CU_ASSERT_EQUAL_FATAL(ret, 0); ret = msp_finalise_tables(&msp); CU_ASSERT_EQUAL(ret, 0); - ret = tsk_treeseq_init(&ts, &tables, TSK_TS_INIT_BUILD_INDEXES); CU_ASSERT_EQUAL_FATAL(ret, 0); ret = tsk_tree_init(&tree, &ts, 0); CU_ASSERT_EQUAL_FATAL(ret, 0); + CU_ASSERT_EQUAL_FATAL(tsk_treeseq_get_num_trees(&ts), 1); - k = 0; - for (ret = tsk_tree_first(&tree); ret == 1; ret = tsk_tree_next(&tree)) { - k++; - // verrify each tree has a single root - CU_ASSERT_EQUAL_FATAL(tsk_tree_get_num_roots(&tree), 1); + ret = tsk_tree_first(&tree); + CU_ASSERT_EQUAL_FATAL(ret, TSK_TREE_OK); + CU_ASSERT_EQUAL_FATAL(tsk_tree_get_num_roots(&tree), 1); - root = tsk_tree_get_left_root(&tree); - if (k == 1) { - check_time = msp.tables->nodes.time[root]; - } else { - CU_ASSERT_EQUAL(msp.tables->nodes.time[root], check_time); - } - } + root = tsk_tree_get_left_root(&tree); + CU_ASSERT_EQUAL(msp.tables->nodes.time[root], msp_get_time(&msp)); model = msp_get_model(&msp)->type; CU_ASSERT_EQUAL(model, models[j]); - model_name = msp_get_model_name(&msp); - CU_ASSERT_STRING_EQUAL(model_name, model_names[j]); ret = msp_free(&msp); CU_ASSERT_EQUAL(ret, 0); @@ -4657,6 +4670,7 @@ main(int argc, char **argv) { CU_TestInfo tests[] = { { "test_single_locus_simulation", test_single_locus_simulation }, + { "test_single_locus_event_by_event", test_single_locus_event_by_event }, { "test_single_locus_two_populations", test_single_locus_two_populations }, { "test_single_locus_many_populations", test_single_locus_many_populations }, { "test_single_locus_labels", test_single_locus_labels }, @@ -4672,8 +4686,8 @@ main(int argc, char **argv) { "test_multi_locus_simulation", test_multi_locus_simulation }, { "test_single_locus_continue_after_local_mrca", test_single_locus_continue_after_local_mrca }, - { "test_single_locus_continue_after_local_mrca_with_end_time", - test_single_locus_continue_after_local_mrca_with_end_time }, + { "test_multi_locus_continue_after_local_mrca", + test_multi_locus_continue_after_local_mrca }, { "test_multi_locus_bottleneck_arg", test_multi_locus_bottleneck_arg }, { "test_multi_locus_store_unary_simple", test_multi_locus_store_unary_simple }, diff --git a/msprime/__init__.py b/msprime/__init__.py index 188e06a65..63b18fe4b 100644 --- a/msprime/__init__.py +++ b/msprime/__init__.py @@ -48,7 +48,6 @@ SweepGenicSelection, FixedPedigree, TimeUnitsMismatchWarning, - PotentialInfiniteSimulationWarning, NodeType, ) diff --git a/msprime/ancestry.py b/msprime/ancestry.py index 667c783bf..3de63a3ea 100644 --- a/msprime/ancestry.py +++ b/msprime/ancestry.py @@ -56,13 +56,6 @@ class TimeUnitsMismatchWarning(UserWarning): """ -class PotentialInfiniteSimulationWarning(UserWarning): - """ - Warning raised when a simulation is likely to run indefinitely: - stop_at_local_mrca false and no end_time set. - """ - - def _model_factory(model: None | str | AncestryModel) -> AncestryModel: """ Returns an AncestryModel corresponding to the specified model @@ -1097,14 +1090,6 @@ def _parse_sim_ancestry( if stop_at_local_mrca is None: stop_at_local_mrca = True - if not stop_at_local_mrca: - if end_time == math.inf or end_time is None: - message = ( - "Warning: stop_at_local_mrca=False with no set end_time may result " - "in long (or infinite) run times, especially with a large recombination " - "rate. You can set an end_time to avoid this warning." - ) - warnings.warn(message, PotentialInfiniteSimulationWarning, stacklevel=2) return Simulator( tables=initial_state, diff --git a/tests/test_algorithms.py b/tests/test_algorithms.py index 9e5143c5e..f55da127c 100644 --- a/tests/test_algorithms.py +++ b/tests/test_algorithms.py @@ -112,14 +112,16 @@ def test_discrete(self): assert ts.num_trees > 1 assert has_discrete_genome(ts) - def test_dtwf(self): - ts = self.run_script("10 --model=dtwf") + @pytest.mark.parametrize("options", ["", "--continue-after-local-mrca"]) + def test_dtwf(self, options): + ts = self.run_script(f"10 --model=dtwf {options}") assert ts.num_trees > 1 assert not has_discrete_genome(ts) assert ts.sequence_length == 100 - def test_dtwf_migration(self): - ts = self.run_script("10 -r 0 --model=dtwf -p 2 -g 0.1") + @pytest.mark.parametrize("options", ["", "--continue-after-local-mrca"]) + def test_dtwf_migration(self, options): + ts = self.run_script(f"10 -r 0 --model=dtwf -p 2 -g 0.1 {options}") assert ts.num_trees == 1 assert ts.sequence_length == 100 assert ts.num_populations == 2 @@ -129,10 +131,11 @@ def test_dtwf_discrete(self): assert ts.num_trees > 1 assert has_discrete_genome(ts) - def test_full_arg(self): + @pytest.mark.parametrize("options", ["", "--continue-after-local-mrca"]) + def test_full_arg(self, options): node_value = sum(2**i for i in (17, 18, 19, 21)) ts = self.run_script( - f"30 -L 200 --additional-nodes {node_value} --all-segments" + f"30 -L 200 --additional-nodes {node_value} --all-segments {options}" ) assert ts.num_trees > 1 node_flags = ts.tables.nodes.flags @@ -156,21 +159,27 @@ def test_store_unary(self): assert ts.sequence_length == 100 verify_unary(ts) - def test_store_unary_dtwf(self): + @pytest.mark.parametrize("options", ["", "--continue-after-local-mrca"]) + def test_store_unary_dtwf(self, options): node_value = 1 << 18 | 1 << 22 - ts = self.run_script( - f"10 --all-segments -d --additional-nodes {node_value} --model=dtwf" + cmd = ( + f"10 --all-segments -d --additional-nodes {node_value} " + f"--model=dtwf {options}" ) + ts = self.run_script(cmd) assert ts.num_samples == 10 assert ts.num_trees > 1 assert ts.sequence_length == 100 verify_dtwf_unary(ts) - def test_store_unary_dtwf_re(self): + @pytest.mark.parametrize("options", ["", "--continue-after-local-mrca"]) + def test_store_unary_dtwf_re(self, options): node_value = 1 << 17 | 1 << 18 | 1 << 22 - ts = self.run_script( - f"10 --all-segments -d --additional-nodes {node_value} --model=dtwf" + cmd = ( + f"10 --all-segments -d --additional-nodes {node_value} " + f"--model=dtwf {options}" ) + ts = self.run_script(cmd) assert ts.num_samples == 10 assert ts.num_trees > 1 assert ts.sequence_length == 100 @@ -267,27 +276,20 @@ def test_store_ca_re(self): assert np.sum(node_flags == msprime.NODE_IS_CA_EVENT) > 0 assert np.sum(node_flags == msprime.NODE_IS_RE_EVENT) > 0 - def test_stop_at_local_mrca_end_time(self): - end_time = 100 + @pytest.mark.parametrize("model", ["hudson", "smc_k", "dtwf"]) + def test_continue_after_local_mrca_multi_tree(self, model): r = 0.1 - ts = self.run_script( - f"10 --continue-after-local-mrca --end-time={end_time} -r {r}" - ) - - # test that simulations continue for all trees to the end_time - old_time = None + ts = self.run_script(f"10 --continue-after-local-mrca -r {r} --model={model}") + assert ts.num_trees >= 2 + # All roots should have the same time + root_times = [] for tree in ts.trees(): - assert len(tree.roots) == 1 # otherwise the test is not valid - u = tree.roots[0] - assert tree.time(u) >= end_time - if old_time is None: - old_time = tree.time(u) - # check that end time is the same for all roots - assert tree.time(u) == old_time - - def test_stop_at_local_mrca(self): - ts = self.run_script("10 --continue-after-local-mrca -r 0") + root_times.append(tree.time(tree.root)) + assert len(set(root_times)) == 1 + @pytest.mark.parametrize("model", ["hudson", "smc_k", "dtwf"]) + def test_continue_after_local_mrca_single_tree(self, model): + ts = self.run_script(f"10 --continue-after-local-mrca -r 0 --model={model}") root_times = [tree.time(tree.root) for tree in ts.trees()] assert len(set(root_times)) == 1 diff --git a/tests/test_ancestry.py b/tests/test_ancestry.py index b7cd2184b..1f7f5fdca 100644 --- a/tests/test_ancestry.py +++ b/tests/test_ancestry.py @@ -900,48 +900,7 @@ def check_roots(self, ts, end_time=None, allow_multiple_roots=False): root_times = [tree.time(tree.root) for tree in ts.trees()] assert len(set(root_times)) == 1 if end_time is not None: - assert root_times[0] >= end_time - - def verify_simulation(self, n, m, r): - """ - Verifies a simulation for the specified parameters. - """ - end_time = 1000 - sim = ancestry._parse_sim_ancestry( - n, - sequence_length=m, - recombination_rate=r, - stop_at_local_mrca=False, - end_time=end_time, - ) - sim.run() - assert sim.num_breakpoints == len(sim.breakpoints) - assert sim.time > 0 - assert sim.num_avl_node_blocks > 0 - assert sim.num_segment_blocks > 0 - assert sim.num_node_mapping_blocks > 0 - tree_sequence = tskit.TableCollection.fromdict( - sim.tables.asdict() - ).tree_sequence() - t = 0.0 - for record in tree_sequence.nodes(): - if record.time > t: - t = record.time - assert t >= end_time - assert sim.time == t - assert sim.num_common_ancestor_events > 0 - assert sim.num_recombination_events >= 0 - assert np.sum(sim.num_migration_events) >= 0 - assert sim.num_multiple_recombination_events >= 0 - self.check_roots(tree_sequence, end_time) - - def test_random_parameters(self): - num_random_sims = 10 - for _ in range(num_random_sims): - n = random.randint(2, 100) - m = random.randint(10, 100) - r = random.random() - self.verify_simulation(n, m, r) + assert root_times[0] == end_time @pytest.mark.parametrize( "model", @@ -953,16 +912,10 @@ def test_random_parameters(self): ], ) def test_single_models(self, model): - """ - Tests that simulations run after the local MRCA when the flag is set - """ - - end_time = 1000 tss = msprime.sim_ancestry( 10, population_size=10, stop_at_local_mrca=False, - end_time=end_time, random_seed=1, sequence_length=10, recombination_rate=0.1, @@ -970,7 +923,7 @@ def test_single_models(self, model): num_replicates=10, ) for ts in tss: - self.check_roots(ts, end_time) + self.check_roots(ts) def test_fixed_pedigree(self): """ @@ -1015,7 +968,6 @@ def test_fixed_pedigree(self): initial_state=pedigree, stop_at_local_mrca=False, recombination_rate=0.1, - end_time=eldest_time, model="fixed_pedigree", ) @@ -1030,7 +982,7 @@ def test_multiple_models(self): end_time = 499 ts = msprime.sim_ancestry( 2, - population_size=10, + population_size=100, model=[ msprime.DiscreteTimeWrightFisher(duration=500), msprime.StandardCoalescent(), @@ -1044,10 +996,9 @@ def test_multiple_models(self): self.check_roots(ts, end_time, allow_multiple_roots=True) - end_time = 100000 ts = msprime.sim_ancestry( 2, - population_size=10, + population_size=100, model=[ msprime.DiscreteTimeWrightFisher(duration=500), msprime.StandardCoalescent(), @@ -1056,31 +1007,8 @@ def test_multiple_models(self): recombination_rate=0.1, sequence_length=10, stop_at_local_mrca=False, - end_time=end_time, ) - self.check_roots(ts, end_time, allow_multiple_roots=True) - - @pytest.mark.parametrize( - "model", - [ - "dtwf", - "hudson", - msprime.BetaCoalescent(alpha=1.5), - msprime.DiracCoalescent(psi=0.1, c=2), - ], - ) - def test_no_end_time(self, model): - with pytest.warns(msprime.PotentialInfiniteSimulationWarning): - ts = msprime.sim_ancestry( - 10, - stop_at_local_mrca=False, - population_size=10, - random_seed=1, - sequence_length=2, - recombination_rate=0, - model=model, - ) - self.check_roots(ts) + self.check_roots(ts) @pytest.mark.parametrize( "model", @@ -1090,7 +1018,6 @@ def test_no_end_time(self, model): ], ) def test_migration(self, model): - end_time = 10000 demography = msprime.Demography.stepping_stone_model( [100, 100], migration_rate=0.01 ) @@ -1098,12 +1025,11 @@ def test_migration(self, model): {0: 2, 1: 1}, demography=demography, stop_at_local_mrca=False, - end_time=end_time, recombination_rate=0.01, sequence_length=10, model=model, ) - self.check_roots(ts, end_time, allow_multiple_roots=True) + self.check_roots(ts) @pytest.mark.parametrize( "model", @@ -1114,7 +1040,6 @@ def test_migration(self, model): ) def test_additional_nodes(self, model): - end_time = 1000 ts = msprime.sim_ancestry( 2, population_size=10, @@ -1127,11 +1052,10 @@ def test_additional_nodes(self, model): coalescing_segments_only=False, stop_at_local_mrca=False, sequence_length=5, - end_time=end_time, model=model, ) - self.check_roots(ts, end_time) + self.check_roots(ts) class TestParseRandomSeed: From 7bfefb2b33c9113efbc7f65c7e59b526ae011a31 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Wed, 3 Dec 2025 14:14:39 +0000 Subject: [PATCH 3/4] Add tests for fixed pedigree and clarify semantics; document Closes #2157 Closes #1898 --- CHANGELOG.md | 7 +-- algorithms.py | 3 -- docs/ancestry.md | 58 +++++++++++++++------ docs/api.md | 16 +++--- lib/msprime.c | 52 +++++++------------ lib/tests/test_ancestry.c | 11 +--- lib/tests/test_pedigrees.c | 24 ++++++--- msprime/ancestry.py | 9 +++- tests/test_ancestry.py | 56 ++++++++++---------- tests/test_lowlevel.py | 4 +- tests/test_pedigree.py | 102 ++++++++++++++++++++++++++----------- 11 files changed, 197 insertions(+), 145 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2174db8f0..11e67327c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,21 +7,18 @@ - The `FixedPedigree` simulation model now supports internal samples ({issue}`1855`, {pr}`2321`, {pr}`2326`, {pr}`2388`, {user}`abureau`, {user}`jeromekelleher`). - +- Support for simulations after local mrca, + ({issue}`2157`, {pr}`2396`, {user}`jeromekelleher`, {user}`hossam26644`). - Add wheels on Windows ({pr}`2414`, {issue}`2200`,{user}`benjeffery`) **Maintenance** - Add support and wheels for Python3.13 -- Support for simulations after local mrca, - ({pr}`2396`, {user}`jeromekelleher`, {user}`hossam26644`). - - Drop Python 3.9 support, require Python >= 3.10 ({pr}`2418`, {user}`benjeffery`) **Breaking changes**: - Require tskit >= 1.0 ({issue}`2434`). - - The `.asdict()` methods for Demography, Population, and Event classes in the demography submodule now return a `__class__` key. This is also stored in their provenance entries, to help recreate demography objects from provenance. diff --git a/algorithms.py b/algorithms.py index 26369c309..6e28da22f 100644 --- a/algorithms.py +++ b/algorithms.py @@ -1541,10 +1541,8 @@ def dtwf_simulate(self, end_time): def dtwf_generation(self): """ Evolves one generation of a Wright Fisher population - Returns True if any events occurred, False otherwise. """ # Migration events happen at the rates in the matrix. - for j in range(len(self.P)): source_size = self.P[j].get_num_ancestors() for k in range(len(self.P)): @@ -2263,7 +2261,6 @@ def merge_ancestors(self, H, pop_id, label, new_node_id=-1): if r_max not in self.S: j = self.S.floor_key(r_max) self.S[r_max] = self.S[j] - min_overlap = len(X) if self.stop_at_local_mrca else 0 # Update the number of extant segments. if self.S[left] == min_overlap: diff --git a/docs/ancestry.md b/docs/ancestry.md index 07df54724..01b70c8af 100644 --- a/docs/ancestry.md +++ b/docs/ancestry.md @@ -1481,17 +1481,16 @@ Migrations nodes are also recorded in the ARG using the `NodeType.MIGRANT` flag. See the {ref}`sec_api_node_flags` section for more details. -(sec_simulation_after_mrca)= +(sec_ancestry_stop_at_local_mrca)= -### Simulations after local MRCA +### Simulating ancestral to local MRCAs By default, msprime stops simulating local trees when a local most recent common ancestor (MRCA) is found. This is because events that occur above the common ancestor are shared across all samples, and we are usually interested in -differences between samples. - -However, for some specialised applications, simulations after the local MRCA -might be needed. In this case, we set the parameter `stop_at_local_mrca` of +differences between samples. However, it is sometime useful to simulate ancestry +across the full span of the region until **all** local MRCAs have been reached. +To do this, we set the parameter `stop_at_local_mrca` of {func}`.sim_ancestry` to `False`. ```{code-cell} @@ -1508,15 +1507,10 @@ because 11 is the root of that tree. The stopping condition remains the same---we keep simulating until an MRCA has been found at *all* local trees. The interpretation of the the final "root" nodes -is a little subtle - -:::{note} -When using the ``stop_at_local_mrca=False`` option you usually need to specify an -``end_time`` argument (see {ref}`sec_ancestry_end_time`) or the simulation will -run indefinitely! It is in general not possible to know what the "right" end -time for a given simulation is, and you may need to specify a "large" value to be -sure that all trees have coalesced locally. -::: +across the trees is a little subtle, and follows the same logic as used when +{ref}`sec_ancestry_end_time`. All of the trees have a root node with time +equal to the oldest MRCA, and the nodes correspond to different lineages present +in the simulation. ## Manipulating simulation time @@ -2819,6 +2813,39 @@ Because the input pedigree fully describes the simulation many features such as combined with other {ref}`ancestry models`. ::: +(sec_ancestry_models_fixed_pedigree_tracing)= + +#### Tracing ancestry through a pedigree + +The previous example generated trees that are embedded within the pedigree, +but doesn't tell us directly the exact path that each sample's genetic +material took through the ancestors. To trace these paths exactly, we can +use a combination of the ``additional_nodes``, +``coalescing_segments_only`` +(see {ref}`sec_ancestry_additional_nodes`), +and ``stop_at_local_mrca`` (see {ref}`sec_ancestry_stop_at_local_mrca`). + +```{code-cell} +ped_ts = msprime.sim_ancestry( + initial_state=pedigree, + model="fixed_pedigree", + random_seed=41, + recombination_rate=0.001, + coalescing_segments_only=False, + additional_nodes=( + msprime.NodeType.COMMON_ANCESTOR | + msprime.NodeType.PASS_THROUGH), + stop_at_local_mrca=True +) +node_labels = {node.id: f"{node.individual}({node.id})" for node in ped_ts.nodes()} +SVG(ped_ts.draw_svg(y_axis=True, node_labels=node_labels, size=(600,200))) +``` + +We can now see that the ancestor of node 15 is 20, which we would have to +infer in the previous example. + + + #### Censoring pedigree information :::{warning} @@ -2922,6 +2949,7 @@ when calculating population sizes. Thus, the pedigree must be seen as entirely **external** to the population model simulated during recapitation. ::: + (sec_ancestry_models_fixed_pedigree_demography)= #### Pedigrees and demography diff --git a/docs/api.md b/docs/api.md index 656fd44c5..48f868bef 100644 --- a/docs/api.md +++ b/docs/api.md @@ -232,12 +232,12 @@ nodes in some situations: ```{eval-rst} .. note:: - These flags have been deprecated in favour of using - :class:`NodeType`. This is not a breaking change. The - constant associated with each flag remains the same. However, working - with flags should be more intuitive now that we are relying on the + These flags have been deprecated in favour of using + :class:`NodeType`. This is not a breaking change. The + constant associated with each flag remains the same. However, working + with flags should be more intuitive now that we are relying on the :class:`python:enum.Flag` functionality. -``` +``` ```{data} msprime.NODE_IS_RE_EVENT @@ -261,7 +261,7 @@ adding common ancestor events to the ``additional_nodes`` to be stored. The node is an ARG migration event identifying the individual that migrated. Can be used in combination with the ``record_migrations`` option. -Present either with the ``record_full_arg`` flag or when adding migration +Present either with the ``record_full_arg`` flag or when adding migration events to the ``additional_nodes`` to be stored. ``` @@ -282,9 +282,9 @@ The node was created by a gene conversion event. ```{data} msprime.NODE_IS_PASS_THROUGH The node identifies an ancestral genome/ploid through which the ancestral -material of only a single lineage passed (so no coalescence or common +material of only a single lineage passed (so no coalescence or common ancestor event). Can be used in combination with the ``record_migrations`` option. -Only present when storing pass through events using the +Only present when storing pass through events using the ``additional_nodes`` option. And only compatible with ``DTWF`` and ``FixedPedigree`` models. diff --git a/lib/msprime.c b/lib/msprime.c index bec53771d..f49d7732c 100644 --- a/lib/msprime.c +++ b/lib/msprime.c @@ -5189,7 +5189,6 @@ msp_run_coalescent(msp_t *self, double max_time, unsigned long max_events) GSL_MIN(gc_t_wait, GSL_MIN(gc_left_t_wait, GSL_MIN(re_t_wait, ca_t_wait)))); if (fixed_event_time == DBL_MAX && t_wait == DBL_MAX) { - msp_print_state(self, stdout); ret = MSP_ERR_INFINITE_WAITING_TIME; goto out; } @@ -5405,7 +5404,7 @@ msp_run_pedigree(msp_t *self, double max_time, unsigned long max_events) pedigree->next_individual++; num_events++; } - if (msp_get_num_ancestors(self) == 0) { + if (msp_is_completed(self) && self->stop_at_local_mrca) { ret = MSP_EXIT_COALESCENCE; } else if (pedigree->next_individual == num_individuals) { ret = MSP_EXIT_MODEL_COMPLETE; @@ -5534,7 +5533,6 @@ msp_dtwf_generation(msp_t *self) } } } - free(parents); free(segment_mem); segment_mem = NULL; @@ -5738,14 +5736,13 @@ msp_run_dtwf(msp_t *self, double max_time, unsigned long max_events) goto out; } ret = msp_apply_demographic_events(self, self->next_demographic_event->time); - // Demographic events can change population structure if (ret != 0) { goto out; } } self->time = cur_time; ret = msp_dtwf_generation(self); - if (ret < 0) { + if (ret != 0) { goto out; } @@ -6084,11 +6081,7 @@ msp_run(msp_t *self, double max_time, unsigned long max_events) goto out; } - if (msp_is_completed(self)) { - /* If the simulation is completed, run() is a no-op for - * all models. */ - ret = 0; - } else if (self->model.type == MSP_MODEL_DTWF) { + if (self->model.type == MSP_MODEL_DTWF) { ret = msp_run_dtwf(self, max_time, max_events); } else if (self->model.type == MSP_MODEL_WF_PED) { ret = msp_run_pedigree(self, max_time, max_events); @@ -6515,36 +6508,27 @@ msp_get_time(msp_t *self) int msp_is_completed(msp_t *self) { - /* int model = self->model.type; */ bool completed; avl_node_t *node; node_mapping_t *nm; - /* NOTE! This is a quick first implementation. This could be a performance - * bottleneck and maybe we keep track of whether there's any non-1 values - * instead - */ - completed = true; - for (node = self->overlap_counts.head; node->next != NULL; node = node->next) { - nm = (node_mapping_t *) node->item; - if (nm->value > 1) { - completed = false; - break; + if (self->stop_at_local_mrca) { + /* When we stop at the local MRCA we have a cheap way of testing for completion + */ + completed = msp_get_num_ancestors(self) == 0; + } else { + /* When we simulate ancestry after local roots, we have to look at the overlap + * counts to find when we've coalesced. + */ + completed = true; + for (node = self->overlap_counts.head; node->next != NULL; node = node->next) { + nm = (node_mapping_t *) node->item; + if (nm->value > 1) { + completed = false; + break; + } } } - - /* } */ - - /* if (model == MSP_MODEL_HUDSON || model == MSP_MODEL_SMC_K || model == - * MSP_MODEL_BETA */ - /* || model == MSP_MODEL_DIRAC || model == MSP_MODEL_SMC */ - /* || model == MSP_MODEL_SMC_PRIME */ - /* || model == MSP_MODEL_DTWF) { */ - /* completed = msp_coalescent_is_complete(self); */ - /* } else { */ - /* completed = msp_get_num_ancestors(self) == 0; */ - /* } */ - return self->state == MSP_STATE_SIMULATING && completed; } diff --git a/lib/tests/test_ancestry.c b/lib/tests/test_ancestry.c index cd51c54e8..069824197 100644 --- a/lib/tests/test_ancestry.c +++ b/lib/tests/test_ancestry.c @@ -817,22 +817,16 @@ test_single_locus_continue_after_local_mrca(void) tsk_treeseq_t ts; tsk_tree_t tree; tsk_id_t root; - gsl_rng *rng = safe_rng_alloc(); - gsl_rng_set(rng, seed); - for (j = 0; j < sizeof(models) / sizeof(int); j++) { - + gsl_rng_set(rng, seed); ret = build_sim(&msp, &tables, rng, m, 1, NULL, n); ret = msp_set_recombination_rate(&msp, 0); CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = msp_set_stop_at_local_mrca(&msp, false); CU_ASSERT_EQUAL_FATAL(ret, 0); - set_simulation_model(&msp, models[j]); - ret = msp_initialise(&msp); CU_ASSERT_EQUAL(ret, 0); ret = msp_run(&msp, DBL_MAX, ULONG_MAX); @@ -845,17 +839,14 @@ test_single_locus_continue_after_local_mrca(void) ret = tsk_tree_init(&tree, &ts, 0); CU_ASSERT_EQUAL_FATAL(ret, 0); CU_ASSERT_EQUAL_FATAL(tsk_treeseq_get_num_trees(&ts), 1); - ret = tsk_tree_first(&tree); CU_ASSERT_EQUAL_FATAL(ret, TSK_TREE_OK); CU_ASSERT_EQUAL_FATAL(tsk_tree_get_num_roots(&tree), 1); - root = tsk_tree_get_left_root(&tree); CU_ASSERT_EQUAL(msp.tables->nodes.time[root], msp_get_time(&msp)); model = msp_get_model(&msp)->type; CU_ASSERT_EQUAL(model, models[j]); - ret = msp_free(&msp); CU_ASSERT_EQUAL(ret, 0); tsk_table_collection_free(&tables); diff --git a/lib/tests/test_pedigrees.c b/lib/tests/test_pedigrees.c index 65022d982..e5889f16f 100644 --- a/lib/tests/test_pedigrees.c +++ b/lib/tests/test_pedigrees.c @@ -207,7 +207,7 @@ verify_pedigree(double recombination_rate, unsigned long seed, static void verify_pedigree_event_by_event(double recombination_rate, unsigned long seed, tsk_size_t num_individuals, tsk_id_t *parents, double *time, tsk_flags_t *is_sample, - tsk_id_t *population, uint32_t additional_nodes) + tsk_id_t *population, uint32_t additional_nodes, bool stop_at_local_mrca) { int ret, status1, status2; int ploidy = 2; @@ -223,6 +223,9 @@ verify_pedigree_event_by_event(double recombination_rate, unsigned long seed, parents, time, is_sample, population); CU_ASSERT_EQUAL_FATAL(ret, 0); + msp_set_stop_at_local_mrca(&msp1, stop_at_local_mrca); + msp_set_stop_at_local_mrca(&msp2, stop_at_local_mrca); + ret = msp_initialise(&msp1); CU_ASSERT_EQUAL_FATAL(ret, 0); ret = msp_initialise(&msp2); @@ -899,8 +902,9 @@ test_internal_samples(void) verify_pedigree(0, 1, 3, parents, time, is_sample, NULL, 0); verify_pedigree(0.1, 1, 3, parents, time, is_sample, NULL, 0); - verify_pedigree_event_by_event(0, 1, 3, parents, time, is_sample, NULL, 0); - verify_pedigree_event_by_event(0.1, 1, 3, parents, time, is_sample, NULL, 0); + verify_pedigree_event_by_event(0, 1, 3, parents, time, is_sample, NULL, 0, true); + verify_pedigree_event_by_event(0.1, 1, 3, parents, time, is_sample, NULL, 0, true); + verify_pedigree_event_by_event(0.1, 1, 3, parents, time, is_sample, NULL, 0, false); } static void @@ -919,8 +923,12 @@ test_no_leaf_samples(void) verify_pedigree(0, 1, num_inds, parents, time, is_sample, NULL, 0); verify_pedigree(0.1, 1, num_inds, parents, time, is_sample, NULL, 0); - verify_pedigree_event_by_event(0, 1, num_inds, parents, time, is_sample, NULL, 0); - verify_pedigree_event_by_event(0.1, 1, num_inds, parents, time, is_sample, NULL, 0); + verify_pedigree_event_by_event( + 0, 1, num_inds, parents, time, is_sample, NULL, 0, true); + verify_pedigree_event_by_event( + 0, 1, num_inds, parents, time, is_sample, NULL, 0, false); + verify_pedigree_event_by_event( + 0.1, 1, num_inds, parents, time, is_sample, NULL, 0, true); } static void @@ -966,7 +974,8 @@ test_trio_same_pop(void) tsk_id_t population[] = { 2, 2, 2 }; verify_pedigree(0, 1, 3, parents, time, NULL, population, 0); - verify_pedigree_event_by_event(0, 1, 3, parents, time, NULL, population, 0); + verify_pedigree_event_by_event(0, 1, 3, parents, time, NULL, population, 0, true); + verify_pedigree_event_by_event(0, 1, 3, parents, time, NULL, population, 0, false); } static void @@ -978,7 +987,8 @@ test_trio_child_different_pop(void) tsk_id_t population[] = { 2, 2, 1 }; verify_pedigree(0, 1, 3, parents, time, NULL, population, 0); - verify_pedigree_event_by_event(0, 1, 3, parents, time, NULL, population, 0); + verify_pedigree_event_by_event(0, 1, 3, parents, time, NULL, population, 0, true); + verify_pedigree_event_by_event(0, 1, 3, parents, time, NULL, population, 0, false); } int diff --git a/msprime/ancestry.py b/msprime/ancestry.py index 3de63a3ea..88415d9d7 100644 --- a/msprime/ancestry.py +++ b/msprime/ancestry.py @@ -1258,8 +1258,13 @@ def sim_ancestry( and the :ref:`sec_ancestry_models` section for the available models and examples. :param stop_at_local_mrca: If True (the default), the simulation will stop for a - tree when local mrca is reached. If False, simulations will continue on the path - to the grand mrca. + tree when local MRCA is reached. If False, simulations will continue across + the full sequence until an MRCA has been found on all segments (except + for the FixedPedigree model, which will continue simulation until all + individuals have been processed). With this option, all trees will have + a root at the same time as the oldest local MRCA, with nodes marking + distinct segments of ancestry. See :ref:`sec_ancestry_stop_at_local_mrca` + for more information. :type model: str or msprime.AncestryModel or list :return: The :class:`tskit.TreeSequence` object representing the results of the simulation if no replication is performed, or an diff --git a/tests/test_ancestry.py b/tests/test_ancestry.py index 1f7f5fdca..6bf217ed5 100644 --- a/tests/test_ancestry.py +++ b/tests/test_ancestry.py @@ -925,53 +925,51 @@ def test_single_models(self, model): for ts in tss: self.check_roots(ts) - def test_fixed_pedigree(self): - """ - To test pedigree model, we make diamond shaped pedigree and make - sure that simulations run until the eldest individual. - There is a chance that no coalescence happens until the eldest individual, - so, we use longer sequence with a high recombination rate and confirm that - all trees does not end before the eldest individual. - - The pedigree used has this shape: - O - | - O - / \ - O O - \\ // - O - / \ - O O - \\ // - O - """ - eldest_time = 10 + def test_fixed_pedigree_full_trace(self): + # The pedigree used has this shape: + # O + # | + # O + # / \ + # O O + # \\ // + # O + # / \ + # O O + # \\ // + # O + eldest_time = 5 pb = msprime.PedigreeBuilder() - gggpa = pb.add_individual(time=eldest_time) ggpa = pb.add_individual(time=4, parents=[gggpa, gggpa]) - # first diamound gmom = pb.add_individual(time=3, parents=[ggpa, ggpa]) gdad = pb.add_individual(time=3, parents=[ggpa, ggpa]) grandpa_id = pb.add_individual(time=2, parents=[gmom, gdad]) - # second diamond mom_id = pb.add_individual(time=1, parents=[grandpa_id, grandpa_id]) dad_id = pb.add_individual(time=1, parents=[grandpa_id, grandpa_id]) pb.add_individual(time=0, parents=[mom_id, dad_id], is_sample=True) - pedigree = pb.finalise(100) ts = msprime.sim_ancestry( initial_state=pedigree, stop_at_local_mrca=False, - recombination_rate=0.1, + coalescing_segments_only=False, + additional_nodes=( + msprime.NodeType.RECOMBINANT + | msprime.NodeType.PASS_THROUGH + | msprime.NodeType.COMMON_ANCESTOR + ), + recombination_rate=0.01, model="fixed_pedigree", ) - - self.check_roots(ts, eldest_time, allow_multiple_roots=True) + for tree in ts.trees(): + for sample in ts.samples(): + path = [sample] + list(tree.ancestors(sample)) + time = ts.nodes_time[path] + # Each sample should pass through each level of the pedigree. + np.testing.assert_array_equal(time, np.arange(eldest_time + 1)) def test_multiple_models(self): """ diff --git a/tests/test_lowlevel.py b/tests/test_lowlevel.py index 8a4f20cec..10cdddda2 100644 --- a/tests/test_lowlevel.py +++ b/tests/test_lowlevel.py @@ -1510,10 +1510,10 @@ def test_sweep_after_coalescence(self): status = sim.run() assert status == _msprime.EXIT_COALESCENCE t_before = sim.time - for _ in range(100): + for _ in range(10): model = get_sweep_genic_selection_model(position=5) sim.model = model - assert sim.run() == _msprime.EXIT_COALESCENCE + assert sim.run() == _msprime.EXIT_MODEL_COMPLETE assert t_before == sim.time def test_store_migrations(self): diff --git a/tests/test_pedigree.py b/tests/test_pedigree.py index bb6b74ace..e8991c501 100644 --- a/tests/test_pedigree.py +++ b/tests/test_pedigree.py @@ -484,9 +484,8 @@ def test_three_pedigrees(self): class TestSimulateThroughPedigree: - def verify(self, input_tables, recombination_rate=0): + def verify(self, input_tables, recombination_rate=0, stop_at_local_mrca=True): initial_state = input_tables.tree_sequence() - # print() # print(list(input_tables.individuals.parents)) # time = [ @@ -502,6 +501,7 @@ def verify(self, input_tables, recombination_rate=0): initial_state=initial_state, recombination_rate=recombination_rate, random_seed=1, + stop_at_local_mrca=stop_at_local_mrca, ) sim.run() # print(sim) @@ -542,22 +542,33 @@ def verify(self, input_tables, recombination_rate=0): assert parent_node.individual in ancestors[individual.id] | { tskit.NULL } + if not stop_at_local_mrca: + # All roots should be founders + for root in tree.roots: + node = ts.node(root) + individual = ts.individual(node.individual) + assert all(p == -1 for p in individual.parents) + return ts @pytest.mark.parametrize("num_founders", [2, 3, 5, 100]) @pytest.mark.parametrize("recombination_rate", [0, 0.01]) - def test_shallow(self, num_founders, recombination_rate): + @pytest.mark.parametrize("stop_at_local_mrca", [True, False]) + def test_shallow(self, num_founders, recombination_rate, stop_at_local_mrca): tables = simulate_pedigree( num_founders=num_founders, num_children_prob=[0, 0, 1], num_generations=2, sequence_length=100, ) - self.verify(tables, recombination_rate) + self.verify(tables, recombination_rate, stop_at_local_mrca) @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): + @pytest.mark.parametrize("stop_at_local_mrca", [True, False]) + def test_shallow_internal( + self, num_founders, recombination_rate, stop_at_local_mrca + ): tables = simulate_pedigree( num_founders=num_founders, num_children_prob=[0, 0, 1], @@ -565,11 +576,14 @@ def test_shallow_internal(self, num_founders, recombination_rate): sequence_length=100, sample_gen=[0, 1], ) - self.verify(tables, recombination_rate) + self.verify(tables, recombination_rate, stop_at_local_mrca) @pytest.mark.parametrize("num_founders", [2, 3, 5]) @pytest.mark.parametrize("recombination_rate", [0, 0.01]) - def test_internal_and_leaf_samples(self, num_founders, recombination_rate): + @pytest.mark.parametrize("stop_at_local_mrca", [True, False]) + def test_internal_and_leaf_samples( + self, num_founders, recombination_rate, stop_at_local_mrca + ): tables = simulate_pedigree( num_founders=num_founders, num_children_prob=[0, 0, 1], @@ -577,11 +591,14 @@ def test_internal_and_leaf_samples(self, num_founders, recombination_rate): sequence_length=100, sample_gen=[0, 2], ) - self.verify(tables, recombination_rate) + self.verify(tables, recombination_rate, stop_at_local_mrca) @pytest.mark.parametrize("num_founders", [2, 3, 5]) @pytest.mark.parametrize("recombination_rate", [0, 0.01]) - def test_no_leaf_samples(self, num_founders, recombination_rate): + @pytest.mark.parametrize("stop_at_local_mrca", [True, False]) + def test_no_leaf_samples( + self, num_founders, recombination_rate, stop_at_local_mrca + ): tables = simulate_pedigree( num_founders=num_founders, num_children_prob=[0, 0, 1], @@ -589,53 +606,58 @@ def test_no_leaf_samples(self, num_founders, recombination_rate): sequence_length=100, sample_gen=[1], ) - self.verify(tables, recombination_rate) + self.verify(tables, recombination_rate, stop_at_local_mrca) @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): + @pytest.mark.parametrize("stop_at_local_mrca", [True, False]) + def test_deep(self, num_founders, recombination_rate, stop_at_local_mrca): tables = simulate_pedigree( num_founders=num_founders, num_children_prob=[0, 0, 1], num_generations=6, # Takes a long time if this is increased sequence_length=100, ) - self.verify(tables, recombination_rate) + self.verify(tables, recombination_rate, stop_at_local_mrca) @pytest.mark.parametrize("num_founders", [2, 3]) @pytest.mark.parametrize("recombination_rate", [0, 0.01]) - def test_very_deep(self, num_founders, recombination_rate): + @pytest.mark.parametrize("stop_at_local_mrca", [True, False]) + def test_very_deep(self, num_founders, recombination_rate, stop_at_local_mrca): tables = simulate_pedigree( num_founders=num_founders, num_children_prob=[0, 0, 1], num_generations=16, # Takes a long time if this is increased sequence_length=10, ) - self.verify(tables, recombination_rate) + self.verify(tables, recombination_rate, stop_at_local_mrca) @pytest.mark.parametrize("num_founders", [2, 3, 10, 20]) @pytest.mark.parametrize("recombination_rate", [0, 0.01]) - def test_many_children(self, num_founders, recombination_rate): + @pytest.mark.parametrize("stop_at_local_mrca", [True, False]) + def test_many_children(self, num_founders, recombination_rate, stop_at_local_mrca): tables = simulate_pedigree( num_founders=num_founders, num_children_prob=[0] * 99 + [1], # Each pair of parents will always have 100 children num_generations=2, ) - self.verify(tables, recombination_rate) + self.verify(tables, recombination_rate, stop_at_local_mrca) @pytest.mark.parametrize("num_founders", [2, 3, 10, 20]) @pytest.mark.parametrize("recombination_rate", [0, 0.01]) - def test_unrelated(self, num_founders, recombination_rate): + @pytest.mark.parametrize("stop_at_local_mrca", [True, False]) + def test_unrelated(self, num_founders, recombination_rate, stop_at_local_mrca): tables = simulate_pedigree( num_founders=num_founders, num_children_prob=[0], num_generations=1, ) - self.verify(tables, recombination_rate) + self.verify(tables, recombination_rate, stop_at_local_mrca) @pytest.mark.parametrize("recombination_rate", [0, 0.01]) - def test_two_pedigrees(self, recombination_rate): + @pytest.mark.parametrize("stop_at_local_mrca", [True, False]) + def test_two_pedigrees(self, recombination_rate, stop_at_local_mrca): tables1 = simulate_pedigree( num_founders=5, num_generations=5, @@ -647,12 +669,15 @@ def test_two_pedigrees(self, recombination_rate): sequence_length=100, ) joined = join_pedigrees([tables1, tables2]) - ts = self.verify(joined.tables, recombination_rate) + ts = self.verify(joined.tables, recombination_rate, stop_at_local_mrca) for tree in ts.trees(): assert tree.num_roots >= 2 @pytest.mark.parametrize("recombination_rate", [0, 0.01]) - def test_two_pedigrees_different_times(self, recombination_rate): + @pytest.mark.parametrize("stop_at_local_mrca", [True, False]) + def test_two_pedigrees_different_times( + self, recombination_rate, stop_at_local_mrca + ): tables1 = simulate_pedigree( num_founders=5, num_generations=5, @@ -666,7 +691,7 @@ def test_two_pedigrees_different_times(self, recombination_rate): # The pedigree in tables2 starts at 1 generation ago tables2.nodes.time += 1 joined = join_pedigrees([tables1, tables2]) - ts = self.verify(joined.tables, recombination_rate) + ts = self.verify(joined.tables, recombination_rate, stop_at_local_mrca) for tree in ts.trees(): assert tree.num_roots >= 2 for root in tree.roots: @@ -847,24 +872,28 @@ def test_all_samples(self, num_founders): class TestSimulateThroughPedigreeEventByEvent(TestSimulateThroughPedigree): - def verify(self, input_tables, recombination_rate=0): + def verify(self, input_tables, recombination_rate=0, stop_at_local_mrca=True): ts1 = msprime.sim_ancestry( model="fixed_pedigree", initial_state=input_tables, recombination_rate=recombination_rate, random_seed=1, + stop_at_local_mrca=stop_at_local_mrca, ) sim = msprime.ancestry._parse_sim_ancestry( model="fixed_pedigree", initial_state=input_tables, recombination_rate=recombination_rate, random_seed=1, + stop_at_local_mrca=stop_at_local_mrca, ) # print(ts1.tables) # print(ts1.draw_text()) sim.run(event_chunk=1) output_tables = tskit.TableCollection.fromdict(sim.tables.asdict()) # print(output_tables) + # ts2 = output_tables.tree_sequence() + # print(ts2.draw_text()) output_tables.assert_equals(ts1.tables, ignore_provenance=True) return ts1 @@ -874,11 +903,12 @@ class TestContinueSimulateThroughPedigree(TestSimulateThroughPedigree): Can we extend the simulations of the pedigree as we'd expect? """ - def verify(self, input_tables, recombination_rate=0): + def verify(self, input_tables, recombination_rate=0, stop_at_local_mrca=True): ts1 = msprime.sim_ancestry( model="fixed_pedigree", initial_state=input_tables, recombination_rate=recombination_rate, + stop_at_local_mrca=stop_at_local_mrca, random_seed=42, ) # print(ts1.draw_text()) @@ -944,7 +974,7 @@ def test_family_different_times(self): class TestSimulateThroughPedigreeReplicates(TestSimulateThroughPedigree): - def verify(self, input_tables, recombination_rate=0): + def verify(self, input_tables, recombination_rate=0, stop_at_local_mrca=True): num_replicates = 5 replicates = list( msprime.sim_ancestry( @@ -953,6 +983,7 @@ def verify(self, input_tables, recombination_rate=0): recombination_rate=recombination_rate, random_seed=42, num_replicates=num_replicates, + stop_at_local_mrca=stop_at_local_mrca, ) ) @@ -960,6 +991,7 @@ def verify(self, input_tables, recombination_rate=0): model="fixed_pedigree", initial_state=input_tables, recombination_rate=recombination_rate, + stop_at_local_mrca=stop_at_local_mrca, random_seed=42, ) ts1.tables.assert_equals(replicates[0].tables, ignore_provenance=True) @@ -1160,7 +1192,7 @@ def test_trio_parents_different_pops_with_migrations(self): class TestSimulateAdditionalNodes: - def verify_pedigree_unary(self, ts, initial_state): + def verify_pedigree_unary(self, ts, initial_state, stop_at_local_mrca=True): # verifies whether for each step in the pedigree and for each # marginal tree there is an edge recorded for each ploid to # either one of the parent's ploids @@ -1184,6 +1216,12 @@ def verify_pedigree_unary(self, ts, initial_state): direct_ancestor = direct_ancestors[node_id] assert tree.parent(node_id) in direct_ancestor + if not stop_at_local_mrca: + for root in tree.roots: + node = ts.node(root) + individual = ts.individual(node.individual) + assert all(p == -1 for p in individual.parents) + def verify_pedigree_recombination(self, tables): children = collections.defaultdict(set) for edge in tables.edges: @@ -1269,7 +1307,8 @@ def test_pedigree_full_arg(self): record_full_arg=True, ) - def test_pedigree_unary(self): + @pytest.mark.parametrize("stop_at_local_mrca", [True, False]) + def test_pedigree_unary(self, stop_at_local_mrca): node_values = [ msprime.NODE_IS_RE_EVENT, msprime.NODE_IS_CA_EVENT | msprime.NODE_IS_PASS_THROUGH, @@ -1286,13 +1325,15 @@ def test_pedigree_unary(self): model="fixed_pedigree", additional_nodes=additional_nodes, coalescing_segments_only=False, + stop_at_local_mrca=stop_at_local_mrca, initial_state=initial_state, random_seed=1234, ) ts = self.verify_pedigree(sim, additional_nodes) - self.verify_pedigree_unary(ts, initial_state) + self.verify_pedigree_unary(ts, initial_state, stop_at_local_mrca) - def test_pedigree_unary_simple(self): + @pytest.mark.parametrize("stop_at_local_mrca", [True, False]) + def test_pedigree_unary_simple(self, stop_at_local_mrca): node_value = ( msprime.NODE_IS_RE_EVENT | msprime.NODE_IS_CA_EVENT @@ -1313,10 +1354,11 @@ def test_pedigree_unary_simple(self): recombination_rate=0.1, model="fixed_pedigree", additional_nodes=additional_nodes, + stop_at_local_mrca=stop_at_local_mrca, coalescing_segments_only=False, initial_state=pedigree, ) - self.verify_pedigree_unary(ts, pedigree) + self.verify_pedigree_unary(ts, pedigree, stop_at_local_mrca) def test_deep_pedigree(self): node_value = ( From 19dafa7546ed11a4fb091742fcfda4d0382df1b8 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Tue, 9 Dec 2025 12:08:49 +0000 Subject: [PATCH 4/4] Fix trailing tickmark in docs --- docs/ancestry.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/ancestry.md b/docs/ancestry.md index 01b70c8af..e80b007b2 100644 --- a/docs/ancestry.md +++ b/docs/ancestry.md @@ -1526,7 +1526,7 @@ for two diploid samples: ```{code-cell} ts = msprime.sim_ancestry(2, sequence_length=10, recombination_rate=0.1, random_seed=42) -SVG(ts.draw_svg(y_axis=True, y_ticks=[0, 1, 2, 3, 4, 5])) +SVG(ts.draw_svg(y_axis=True, y_ticks=[0, 1, 2, 3, 4])) ``` Sometimes we would like to stop the simulation early, **before** complete