diff --git a/CHANGELOG.md b/CHANGELOG.md index adf47bb10..11e67327c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,19 +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 - - 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 e17a8ed19..6e28da22f 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)) @@ -1046,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) @@ -1253,7 +1249,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": @@ -1311,38 +1307,43 @@ 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. """ - 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() # 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 # 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 +1351,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 +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 + 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 @@ -1400,6 +1401,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, @@ -1441,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: @@ -1524,13 +1526,16 @@ def pedigree_simulate(self): self.pedigree = Pedigree(self.tables) self.dtwf_climb_pedigree() - def dtwf_simulate(self): + def dtwf_simulate(self, end_time): """ Simulates the algorithm until all loci have coalesced. """ - while self.ancestors_remain(): - self.t += 1 + while not self.is_completed(): self.verify() + if self.t >= end_time: + break + self.t += 1 + # print("DTWF", self.t) self.dtwf_generation() def dtwf_generation(self): @@ -2256,13 +2261,14 @@ 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): + while right < r_max and self.S[right] != min_overlap: self.S[right] -= len(X) - 1 right = self.S.succ_key(right) alpha = self.alloc_segment(left, right, new_node_id, pop_id) @@ -2394,6 +2400,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 +2442,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( @@ -2732,10 +2739,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 @@ -2860,6 +2867,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 +2954,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..e80b007b2 100644 --- a/docs/ancestry.md +++ b/docs/ancestry.md @@ -1129,6 +1129,7 @@ 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_additional_nodes_ca)= ### Common ancestor events @@ -1480,6 +1481,38 @@ 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_ancestry_stop_at_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, 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} +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 +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 (sec_ancestry_end_time)= @@ -1493,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 @@ -2780,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} @@ -2883,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 fc84f3cc3..f49d7732c 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,7 @@ 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->avl_node_block_size = 1024; self->node_mapping_block_size = 1024; self->segment_block_size = 1024; @@ -1098,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) { @@ -3777,6 +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 = self->stop_at_local_mrca ? 2 : 0; avl_node_t *node; node_mapping_t *nm, search; segment_t *x, *y, *z, *alpha, *beta; @@ -3862,7 +3860,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 +3868,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); @@ -3981,7 +3979,7 @@ 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; avl_node_t *node; node_mapping_t *nm, search; @@ -4072,7 +4070,8 @@ 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) { + min_overlap = self->stop_at_local_mrca ? h : 0; + if (nm->value == min_overlap) { nm->value = 0; node = node->next; tsk_bug_assert(node != NULL); @@ -4080,7 +4079,7 @@ 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) { + while (nm->value != min_overlap && r < r_max) { nm->value -= h - 1; node = node->next; tsk_bug_assert(node != NULL); @@ -5119,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; @@ -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; @@ -5638,7 +5637,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; @@ -5942,7 +5941,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++) { @@ -6082,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); @@ -6206,7 +6201,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; @@ -6507,6 +6502,36 @@ msp_get_time(msp_t *self) return self->time; } +/* + * Returns true if the simulation has completed. + */ +int +msp_is_completed(msp_t *self) +{ + bool completed; + avl_node_t *node; + node_mapping_t *nm; + + 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; + } + } + } + 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. */ @@ -8042,9 +8067,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 c352ef539..25c5da47b 100644 --- a/lib/msprime.h +++ b/lib/msprime.h @@ -256,6 +256,7 @@ typedef struct _msp_t { bool store_full_arg; uint32_t additional_nodes; bool coalescing_segments_only; + bool stop_at_local_mrca; double sequence_length; bool discrete_genome; rate_map_t recomb_map; @@ -497,6 +498,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..069824197 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) { @@ -666,6 +738,124 @@ test_multi_locus_simulation(void) gsl_rng_free(rng); } +static void +test_multi_locus_continue_after_local_mrca(void) +{ + int ret; + uint32_t n = 10; + uint32_t m = 100; + long seed = 11; + double t; + int model; + 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(); + + 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); + + 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(ret == MSP_EXIT_COALESCENCE); + t = msp_get_time(&msp); + ret = msp_finalise_tables(&msp); + CU_ASSERT_EQUAL(ret, 0); + + /* 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)) { + 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); + } + 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); + 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 seed = 10; + int model; + int models[] = { MSP_MODEL_HUDSON, MSP_MODEL_SMC, MSP_MODEL_SMC_K, MSP_MODEL_DTWF, + MSP_MODEL_BETA }; + 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(); + + 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); + 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); + 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); + tsk_treeseq_free(&ts); + tsk_tree_free(&tree); + } + gsl_rng_free(rng); +} + static void test_multi_locus_bottleneck_arg(void) { @@ -4471,6 +4661,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 }, @@ -4484,6 +4675,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_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/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/_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..88415d9d7 100644 --- a/msprime/ancestry.py +++ b/msprime/ancestry.py @@ -815,6 +815,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 +1088,9 @@ 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 + return Simulator( tables=initial_state, recombination_map=recombination_map, @@ -1103,6 +1107,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 +1135,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 +1257,14 @@ 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 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 @@ -1289,6 +1303,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 +1328,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 +1414,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 +1465,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..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,6 +276,23 @@ 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 + @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 -r {r} --model={model}") + assert ts.num_trees >= 2 + # All roots should have the same time + root_times = [] + for tree in ts.trees(): + 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 + 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..6bf217ed5 100644 --- a/tests/test_ancestry.py +++ b/tests/test_ancestry.py @@ -891,6 +891,171 @@ 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 + + @pytest.mark.parametrize( + "model", + [ + "dtwf", + "hudson", + msprime.BetaCoalescent(alpha=1.5), + msprime.DiracCoalescent(psi=0.1, c=2), + ], + ) + def test_single_models(self, model): + tss = msprime.sim_ancestry( + 10, + population_size=10, + stop_at_local_mrca=False, + random_seed=1, + sequence_length=10, + recombination_rate=0.1, + model=model, + num_replicates=10, + ) + for ts in tss: + self.check_roots(ts) + + 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, + coalescing_segments_only=False, + additional_nodes=( + msprime.NodeType.RECOMBINANT + | msprime.NodeType.PASS_THROUGH + | msprime.NodeType.COMMON_ANCESTOR + ), + recombination_rate=0.01, + model="fixed_pedigree", + ) + 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): + """ + 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=100, + 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) + + ts = msprime.sim_ancestry( + 2, + population_size=100, + model=[ + msprime.DiscreteTimeWrightFisher(duration=500), + msprime.StandardCoalescent(), + ], + random_seed=2, + recombination_rate=0.1, + sequence_length=10, + stop_at_local_mrca=False, + ) + self.check_roots(ts) + + @pytest.mark.parametrize( + "model", + [ + "dtwf", + "hudson", + ], + ) + def test_migration(self, model): + 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, + recombination_rate=0.01, + sequence_length=10, + model=model, + ) + self.check_roots(ts) + + @pytest.mark.parametrize( + "model", + [ + "dtwf", + "hudson", + ], + ) + def test_additional_nodes(self, model): + + 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, + model=model, + ) + + self.check_roots(ts) + + class TestParseRandomSeed: """ Tests for parsing the random seed values. 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 = (