Change stopping condition for stop_at_local_mrca=False

Author: jeromekelleher

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,38 +1526,24 @@ 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):
         """
         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()
@@ -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

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