Remove SLiM checkpointing except for at first DrawMutation (#1335)

* Save only at first DrawMutation event

* Remove mention of save argument from docs

* Small fixups for SLiM engine errors

docs/selection_example.py

@@ -72,14 +72,12 @@ def adaptive_introgression(seed):
             mutation_type_id=mut_id,
             population_id=pop["DenA"],
             coordinate=coordinate,
-            # Save state before the mutation is introduced.
-            save=True,
         ),
         # Because the drawn mutation is neutral at the time of introduction,
         # it's likely to be lost due to drift. To avoid this, we condition on
         # the mutation having AF > 0 in DenA. If this condition is false at any
         # time between start_time and end_time, the simulation will be
-        # restored to the most recent save point.
+        # returned to the point where the mutation was introduced.
         # Conditioning should start one generation after T_mut (not at T_mut!),
         # to avoid checking for the mutation before SLiM can introduce it.
         stdpopsim.ext.ConditionOnAlleleFrequency(
@@ -102,8 +100,6 @@ def adaptive_introgression(seed):
             population_id=pop["Den1"],
             op=">",
             allele_frequency=0,
-            # Update save point at start_time (if the condition is met).
-            save=True,
         ),
         # The Den1 lineage has migrants entering the Papaun lineage at T_mig,
         # so condition on AF > 0 in Papuans.
@@ -114,10 +110,6 @@ def adaptive_introgression(seed):
             population_id=pop["Papuan"],
             op=">",
             allele_frequency=0,
-            # Update save point at start_time (if the condition is met).
-            # If the Den1 migrants didn't carry the mutation, we will
-            # instead restore to the previous save point.
-            save=True,
         ),
         # The mutation is positively selected in Papuans at T_sel.
         # Note that this will have no effect, unless/until a mutation with the

docs/tutorial.rst

@@ -1253,8 +1253,7 @@ unique string ID. Here, we'll add the site in the middle of the contig, with ID
 Next, we will set up the "extended events" which will modify the demography.
 The first extended event is the origination of the selected mutation, which
 will occur in a random individual from the first population (id 0), 1000
-generations ago.  We specify ``save=True`` to ``stdpopsim.ext.DrawMutation`` so
-that the simulation can restart from that point if the mutation is lost.
+generations ago.
 
 .. code-block:: python
 

stdpopsim/ext/selection.py

@@ -63,9 +63,6 @@ class DrawMutation(ExtendedEvent):
     Draw a new mutation with the given mutation type in the given population
     at the given genomic coordinate. The mutation will be added to one randomly
     chosen chromosome in the given population.
-    If save=True, the simulation state is saved before the mutation is
-    introduced, to facilitate the save/restore mechanism used for allele
-    frequency conditioning.
 
     FIXME: Drawing multiple mutations using the same mutation type causes
            erroneous allele frequency calculation in the SLiM code!
@@ -78,7 +75,6 @@ class DrawMutation(ExtendedEvent):
     time = attr.ib(type=float)
     single_site_id = attr.ib(type=str)
     population_id = attr.ib(type=int)
-    save = attr.ib(type=bool, default=False)
 
     def __attrs_post_init__(self):
         validate_time(self.time)
@@ -116,8 +112,6 @@ class ConditionOnAlleleFrequency(ExtendedEvent):
     Condition on the allele frequency of a drawn mutation with the given
     mutation type in the given population. The mutation need not have been
     drawn in the population being conditioned upon.
-    If save=True, the simulation state will be saved if the condition is
-    met at start_time.
 
     This uses a save/restore mechanism in the SLiM code to do rejection
     sampling on a simulation (when the condition is not met) without throwing
@@ -130,7 +124,6 @@ class ConditionOnAlleleFrequency(ExtendedEvent):
     population_id = attr.ib(type=int)
     op = attr.ib(type=str, default=None)
     allele_frequency = attr.ib(type=float, default=None)
-    save = attr.ib(type=bool, default=False)
 
     op_types = ("<", "<=", ">", ">=")
 

stdpopsim/genomes.py

@@ -508,7 +508,6 @@ def add_single_site(
                     time=mutation_time,
                     single_site_id="hard_sweep",
                     population_id=0,
-                    save=True,
                 ),
                 stdpopsim.ext.ConditionOnAlleleFrequency(
                     start_time=stdpopsim.ext.GenerationAfter(mutation_time),

stdpopsim/slim_engine.py

@@ -293,14 +293,21 @@
 
     // Save/restore events. These should come before all other events.
     if (length(drawn_mutations) > 0) {
+        n_checkpoints = 0;
         for (i in 0:(ncol(drawn_mutations)-1)) {
             save = drawn_mutations[4,i] == 1;
-            if (!save) {
-                next;
+            if (save) {
+                // Saving the state at more than one timepoint can can cause
+                // incorrect conditioning in the rejection samples.
+                if (n_checkpoints > 0) {
+                    err("Attempt to save state at more than one checkpoint");
+                }
+                n_checkpoints = n_checkpoints + 1;
+
+                // Unconditionally save the state before the mutation is drawn.
+                g = G_start + gdiff(T_start, drawn_mutations[0,i]);
+                community.registerLateEvent(NULL, "{save();}", g, g);
             }
-            g = G_start + gdiff(T_start, drawn_mutations[0,i]);
-            // Unconditionally save the state before the mutation is drawn.
-            community.registerLateEvent(NULL, "{save();}", g, g);
         }
     }
     if (length(condition_on_allele_frequency) > 0) {
@@ -311,31 +318,17 @@
             pop_id = asInteger(condition_on_allele_frequency[3,i]);
             op = op_types[asInteger(drop(condition_on_allele_frequency[4,i]))];
             af = condition_on_allele_frequency[5,i];
-            save = condition_on_allele_frequency[6,i] == 1;
 
             if (g_start > g_end) {
                 err("Attempt to register AF conditioning callback with g_start="+
                     g_start+" > g_end="+g_end);
             }
 
-            // Save the state conditional on the allele frequency.
-            // If the condition isn't met, we restore.
-            if (save) {
-                // Save the state conditional on the allele frequency.
-                // If the condition isn't met, we restore.
-                community.registerLateEvent(NULL,
-                    "{if (af(m"+mut_type+", p"+pop_id+") "+op+" "+af+")" +
-                    " save(); else restore();}",
-                    g_start, g_start);
-                g_start = g_start + 1;
-            }
-
-            if (g_start <= g_end) {
-                community.registerLateEvent(NULL,
-                    "{if (!(af(m"+mut_type+", p"+pop_id+") "+op+" "+af+"))" +
-                    " restore();}",
-                    g_start, g_end);
-            }
+            // Restore state if AF condition not met.
+            community.registerLateEvent(NULL,
+                "{if (!(af(m"+mut_type+", p"+pop_id+") "+op+" "+af+"))" +
+                " restore();}",
+                g_start, g_end);
         }
     }
 
@@ -1017,7 +1010,7 @@ def fix_time(event):
                     mutation_type_id,
                     ee.population_id,
                     coordinate,
-                    int(ee.save),
+                    0,  # flag to save state in mutation generation
                 )
             )
             drawn_single_site_ids.add(ee.single_site_id)
@@ -1042,7 +1035,6 @@ def fix_time(event):
                     ee.population_id,
                     op_id(ee.op),
                     ee.allele_frequency,
-                    int(ee.save),
                 )
             )
             referenced_single_site_ids.add(ee.single_site_id)
@@ -1052,7 +1044,15 @@ def fix_time(event):
     # Check that drawn mutations exist for extended events that need them.
     for single_site_id in referenced_single_site_ids:
         if single_site_id not in drawn_single_site_ids:
-            raise ValueError(f"No drawn mutation for single site id {single_site_id}.")
+            raise ValueError(
+                f"An extended event requires a mutation at single site id "
+                f"{single_site_id}, but no mutation is drawn at this site."
+            )
+
+    # Set "save state" flag for the oldest drawn mutation
+    drawn_mutations = sorted(drawn_mutations, key=lambda x: x[0], reverse=True)
+    if len(drawn_mutations) > 0:
+        drawn_mutations[0] = drawn_mutations[0][:-1] + (1,)
 
     printsc = functools.partial(print, file=script_file)
 
@@ -1306,7 +1306,7 @@ def matrix2str(
         + matrix2str(
             condition_on_allele_frequency,
             col_comment="start_time, end_time, mut_type, pop_id, "
-            "op, allele_frequency, save",
+            "op, allele_frequency",
         )
         + ");"
     )