Add tests for fixed pedigree and clarify semantics; document

Closes #2157
Closes #1898

Author: jeromekelleher

CHANGELOG.md

@@ -7,13 +7,10 @@
 
 - The `FixedPedigree` simulation model now supports internal samples ({issue}`1855`,
   {pr}`2321`, {pr}`2326`, {pr}`2388`, {user}`abureau`, {user}`jeromekelleher`).
-
 - Add support and wheels for Python3.13
 - Support for simulations after local mrca,
-  ({pr}`2396`, {user}`jeromekelleher`, {user}`hossam26644`).
-
+  ({issue}`2157`, {pr}`2396`, {user}`jeromekelleher`, {user}`hossam26644`).
 - Drop Python 3.9 support, require Python >= 3.10 ({pr}`2418`, {user}`benjeffery`)
-
 - Add wheels on Windows ({pr}`2414`, {issue}`2200`,{user}`benjeffery`)
 
 **Breaking changes**:

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:

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
@@ -1532,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
@@ -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>`.
 :::
 
+(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

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<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<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.
 

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;
 }
 

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);

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