Repository: popsim-consortium/stdpopsim (language: python)

Issue:
Fix checkpointing behavior in SLiM rejection sampling
As described in #1325, checkpointed rejection sampling can easily produce incorrect simulations. The only situation where checkpointing is "safe" is when a mutation is drawn and there are no other extant mutations that are being conditioned on.

I think the following would fix the issue:

- Remove `save` option from python API (extended event python classes)
- Remove save mechanics from SLiM engine for `ConditionOnAlleleFrequency` events
- In `slim_makescript`, sort `DrawMutation` events by time and set `save` flag only at the first event


Discussion / hints:
I agree, this sounds like the right thing to do!

Fix patch:
diff --git a/docs/selection_example.py b/docs/selection_example.py
index 4a14decca..16590b39b 100644
--- a/docs/selection_example.py
+++ b/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
diff --git a/docs/tutorial.rst b/docs/tutorial.rst
index bab510ead..f5c0d782f 100644
--- a/docs/tutorial.rst
+++ b/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
 
diff --git a/stdpopsim/ext/selection.py b/stdpopsim/ext/selection.py
index 4bd68f97a..60fb6d597 100644
--- a/stdpopsim/ext/selection.py
+++ b/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 = ("<", "<=", ">", ">=")
 
diff --git a/stdpopsim/genomes.py b/stdpopsim/genomes.py
index 93cb56dc3..fc14a1708 100644
--- a/stdpopsim/genomes.py
+++ b/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),
diff --git a/stdpopsim/slim_engine.py b/stdpopsim/slim_engine.py
index ffd686dc3..d529822ac 100644
--- a/stdpopsim/slim_engine.py
+++ b/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",
         )
         + ");"
     )


Test patch:
diff --git a/tests/test_slim_engine.py b/tests/test_slim_engine.py
index 3cb81b113..e7ceebd39 100644
--- a/tests/test_slim_engine.py
+++ b/tests/test_slim_engine.py
@@ -1356,36 +1356,45 @@ def test_logfile(self, tmp_path):
 
 @pytest.mark.skipif(IS_WINDOWS, reason="SLiM not available on windows")
 class TestDrawMutation(PiecewiseConstantSizeMixin):
-    def test_draw_mutation_save(self):
+    def test_draw_mutation(self):
+        contig = get_test_contig()
+        contig.add_single_site(id=self.mut_id, coordinate=100)
         extended_events = [
             stdpopsim.ext.DrawMutation(
                 time=self.T_mut,
                 single_site_id=self.mut_id,
                 population_id=0,
-                save=True,
             ),
         ]
         engine = stdpopsim.get_engine("slim")
         engine.simulate(
             demographic_model=self.model,
-            contig=self.contig,
+            contig=contig,
             samples=self.samples,
             extended_events=extended_events,
             dry_run=True,
         )
 
-    def test_draw_mutation_no_save(self):
+    def test_draw_multiple_mutations(self):
+        contig = get_test_contig()
+        contig.add_single_site(id="recent", coordinate=100)
+        contig.add_single_site(id="older", coordinate=101)
         extended_events = [
+            stdpopsim.ext.DrawMutation(
+                time=int(self.T_mut / 2),
+                single_site_id="recent",
+                population_id=0,
+            ),
             stdpopsim.ext.DrawMutation(
                 time=self.T_mut,
-                single_site_id=self.mut_id,
+                single_site_id="older",
                 population_id=0,
             ),
         ]
         engine = stdpopsim.get_engine("slim")
         engine.simulate(
             demographic_model=self.model,
-            contig=self.contig,
+            contig=contig,
             samples=self.samples,
             extended_events=extended_events,
             dry_run=True,
@@ -1627,42 +1636,6 @@ def test_bad_time(self):
 
 @pytest.mark.skipif(IS_WINDOWS, reason="SLiM not available on windows")
 class TestAlleleFrequencyConditioning(PiecewiseConstantSizeMixin):
-    def test_save_point_creation(self):
-        extended_events = [
-            stdpopsim.ext.DrawMutation(
-                time=self.T_mut,
-                single_site_id=self.mut_id,
-                population_id=0,
-                save=True,
-            ),
-            stdpopsim.ext.ConditionOnAlleleFrequency(
-                start_time=stdpopsim.ext.GenerationAfter(self.T_mut),
-                end_time=0,
-                single_site_id=self.mut_id,
-                population_id=0,
-                op=">",
-                allele_frequency=0,
-                save=True,
-            ),
-            stdpopsim.ext.ConditionOnAlleleFrequency(
-                start_time=self.T_mut // 2,
-                end_time=self.T_mut // 2,
-                single_site_id=self.mut_id,
-                population_id=0,
-                op=">",
-                allele_frequency=0,
-                save=True,
-            ),
-        ]
-        engine = stdpopsim.get_engine("slim")
-        engine.simulate(
-            demographic_model=self.model,
-            contig=self.contig,
-            samples=self.samples,
-            extended_events=extended_events,
-            dry_run=True,
-        )
-
     @pytest.mark.filterwarnings("ignore::stdpopsim.SLiMScalingFactorWarning")
     def test_drawn_mutation_not_lost(self):
         engine = stdpopsim.get_engine("slim")
@@ -1674,7 +1647,6 @@ def test_drawn_mutation_not_lost(self):
                 time=self.T_mut,
                 single_site_id="test",
                 population_id=0,
-                save=True,
             ),
             stdpopsim.ext.ConditionOnAlleleFrequency(
                 start_time=0,
@@ -1707,7 +1679,6 @@ def test_drawn_mutation_is_lost(self):
                 time=self.T_mut,
                 single_site_id="test",
                 population_id=0,
-                save=True,
             ),
             stdpopsim.ext.ConditionOnAlleleFrequency(
                 start_time=0,
@@ -1741,7 +1712,6 @@ def test_drawn_mutation_meets_AF_threshold(self):
                     time=self.T_mut,
                     single_site_id="test",
                     population_id=0,
-                    save=True,
                 ),
                 # Condition on desired AF at end of simulation.
                 stdpopsim.ext.ConditionOnAlleleFrequency(
@@ -1825,7 +1795,6 @@ def test_bad_GenerationAfter_times(self):
                         time=self.T_mut,
                         single_site_id=self.mut_id,
                         population_id=0,
-                        save=True,
                     ),
                     stdpopsim.ext.ConditionOnAlleleFrequency(
                         start_time=stdpopsim.ext.GenerationAfter(start_time),
@@ -1853,35 +1822,6 @@ def test_op_id(self):
             with pytest.raises(ValueError):
                 id = stdpopsim.ext.ConditionOnAlleleFrequency.op_id(op)
 
-    @pytest.mark.filterwarnings("ignore::stdpopsim.SLiMScalingFactorWarning")
-    def test_conditioning_without_save(self):
-        extended_events = [
-            stdpopsim.ext.DrawMutation(
-                time=self.T_mut,
-                single_site_id=self.mut_id,
-                population_id=0,
-            ),
-            stdpopsim.ext.ConditionOnAlleleFrequency(
-                start_time=stdpopsim.ext.GenerationAfter(self.T_mut),
-                end_time=0,
-                single_site_id=self.mut_id,
-                population_id=0,
-                op=">=",
-                allele_frequency=1,
-            ),
-        ]
-        engine = stdpopsim.get_engine("slim")
-        with pytest.raises(stdpopsim.SLiMException):
-            # TODO: get this to fail using dry_run=True
-            engine.simulate(
-                demographic_model=self.model,
-                contig=self.contig,
-                samples=self.samples,
-                extended_events=extended_events,
-                slim_scaling_factor=10,
-                slim_burn_in=0.1,
-            )
-
     def test_no_drawn_mutation(self):
         extended_events = [
             stdpopsim.ext.ConditionOnAlleleFrequency(
@@ -1891,11 +1831,10 @@ def test_no_drawn_mutation(self):
                 population_id=0,
                 op=">",
                 allele_frequency=0,
-                save=True,
             ),
         ]
         engine = stdpopsim.get_engine("slim")
-        with pytest.raises(ValueError, match="No drawn mutation"):
+        with pytest.raises(ValueError, match="no mutation is drawn at this site"):
             engine.simulate(
                 demographic_model=self.model,
                 contig=self.contig,
@@ -1923,7 +1862,6 @@ def test_drawn_mutation_has_correct_selection_coeff(self):
                 time=self.T_mut,
                 single_site_id=mut_id,
                 population_id=0,
-                save=True,
             ),
             stdpopsim.ext.ConditionOnAlleleFrequency(
                 start_time=0,
@@ -1969,7 +1907,6 @@ def test_positive_mutation_meets_AF_threshold(self):
                     time=self.T_mut,
                     single_site_id=self.mut_id,
                     population_id=0,
-                    save=True,
                 ),
                 stdpopsim.ext.ChangeMutationFitness(
                     start_time=stdpopsim.ext.GenerationAfter(self.T_mut),
@@ -2024,7 +1961,7 @@ def test_no_drawn_mutation(self):
             ),
         ]
         engine = stdpopsim.get_engine("slim")
-        with pytest.raises(ValueError, match="No drawn mutation"):
+        with pytest.raises(ValueError, match="no mutation is drawn at this site"):
             engine.simulate(
                 demographic_model=self.model,
                 contig=self.contig,
@@ -2064,7 +2001,6 @@ def test_bad_GenerationAfter_times(self):
                         time=self.T_mut,
                         single_site_id=self.mut_id,
                         population_id=0,
-                        save=True,
                     ),
                     stdpopsim.ext.ChangeMutationFitness(
                         start_time=stdpopsim.ext.GenerationAfter(start_time),