Set single site mutation types to non neutral

docs/tutorial.rst

@@ -1216,7 +1216,9 @@ First, we need to add the site at which the selected mutation will occur.  This
 is like adding a DFE, except to a single site -- we're saying that there is a
 potential mutation at a particular site with defined fitness consequences. So
 that we can refer to the single site later, we give it a unique string ID.
-Here, we'll add the site in the middle of the contig, with ID "hard sweep".
+Here, we'll add the site in the middle of the contig with ID "hard sweep",
+so named because we will imagine this beneficial mutation originates at
+frequency :math:`1 / 2N`.
 
 .. code-block:: python
 
@@ -1235,13 +1237,13 @@ Here, we'll add the site in the middle of the contig, with ID "hard sweep".
     contains a single site mutation, the single site mutation will be
     "overwritten" and an error will be raised in simulation.
 
-Next, we will set up the "extended events" which will modify the demography.  This
-is done through :func:`stdpopsim.ext.selective_sweep`, which represents a general model
-for a mutation that is beneficial over a time interval within a single
-population.  We specify that the mutation should originate 1000 generations ago
-in a random individual from the first population (id 0); that the selection
-coefficient for the mutation should be 0.5; and that the frequency of the mutation
-in the present day (e.g. at the end of the sweep) should be 0.8.
+Next, we will set up the "extended events" which will modify the demography.
+This is done through :func:`stdpopsim.ext.selective_sweep`, which represents a
+general model for a mutation that is beneficial within a single population.  We
+specify that the mutation should originate 1000 generations ago in a random
+individual from the first population (id 0); that the selection coefficient for
+the mutation should be 0.5; and that the frequency of the mutation in the
+present day (e.g. at the end of the sweep) should be greater than 0.8.
 
 .. code-block:: python
 

stdpopsim/slim_engine.py

@@ -1534,7 +1534,7 @@ def script_file_f():
         ts = _add_dfes_to_metadata(ts, contig)
         if _recap_and_rescale:
             ts = self._recap_and_rescale(
-                ts, seed, recap_epoch, contig, slim_scaling_factor
+                ts, seed, recap_epoch, contig, slim_scaling_factor, extended_events
             )
 
         if contig.inclusion_mask is not None:
@@ -1613,7 +1613,9 @@ def _simplify_remembered(self, ts):
         )
         return ts.simplify(samples=list(nodes), filter_populations=False)
 
-    def _recap_and_rescale(self, ts, seed, recap_epoch, contig, slim_scaling_factor):
+    def _recap_and_rescale(
+        self, ts, seed, recap_epoch, contig, slim_scaling_factor, extended_events=None
+    ):
         """
         Apply post-SLiM transformations to ``ts``. This rescales node times,
         does recapitation, simplification, and adds neutral mutations.
@@ -1625,6 +1627,17 @@ def _recap_and_rescale(self, ts, seed, recap_epoch, contig, slim_scaling_factor)
             table.time *= slim_scaling_factor
         metadata = tables.metadata
         metadata["SLiM"]["tick"] *= slim_scaling_factor
+        # If a single site DFE is referenced by an extended event, mark the DFE
+        # as non-neutral so that mutations aren't simulated over it.
+        if extended_events is not None:
+            single_sites = set()
+            for ee in extended_events:
+                assert hasattr(ee, "single_site_id")
+                single_sites.add(ee.single_site_id)
+            for dfe in metadata["stdpopsim"]["DFEs"]:
+                if dfe["id"] in single_sites:
+                    assert len(dfe["mutation_types"]) == 1
+                    dfe["mutation_types"][0]["is_neutral"] = False
         # Finding what slim id to use in recap DFE
         max_id = -1
         for dfe in metadata["stdpopsim"]["DFEs"]:
@@ -1777,7 +1790,9 @@ def recap_and_rescale(
                 slim_rate_map,
             )
 
-        ts = self._recap_and_rescale(ts, seed, recap_epoch, contig, slim_scaling_factor)
+        ts = self._recap_and_rescale(
+            ts, seed, recap_epoch, contig, slim_scaling_factor, extended_events
+        )
         return ts
 
 

tests/test_slim_engine.py

@@ -1930,6 +1930,32 @@ def test_positive_mutation_meets_AF_threshold(self):
             assert ts.num_mutations == 1
             assert self.allele_frequency(ts) >= af_threshold
 
+    @pytest.mark.filterwarnings("ignore::stdpopsim.SLiMScalingFactorWarning")
+    def test_referenced_single_site_is_nonneutral(self):
+        engine = stdpopsim.get_engine("slim")
+        contig = get_test_contig()
+        contig.add_single_site("one", coordinate=100)
+        contig.add_single_site("two", coordinate=101)
+        extended_events = [
+            stdpopsim.ext.DrawMutation(
+                time=self.T_mut,
+                single_site_id="one",
+                population_id=0,
+            ),
+        ]
+        ts = engine.simulate(
+            demographic_model=self.model,
+            contig=contig,
+            samples=self.samples,
+            extended_events=extended_events,
+        )
+        referenced_dfe = ts.metadata["stdpopsim"]["DFEs"][1]
+        assert referenced_dfe["id"] == "one"
+        assert referenced_dfe["mutation_types"][0]["is_neutral"] is False
+        unreferenced_dfe = ts.metadata["stdpopsim"]["DFEs"][2]
+        assert unreferenced_dfe["id"] == "two"
+        assert unreferenced_dfe["mutation_types"][0]["is_neutral"] is True
+
     def test_no_drawn_mutation(self):
         extended_events = [
             stdpopsim.ext.ChangeMutationFitness(
@@ -2150,21 +2176,25 @@ def test_sweep_meets_min_freq_at_start(self, tmp_path):
             selection_coeff=s,
             dominance_coeff=0.5,
         )
-        engine.simulate(
-            demographic_model=self.model,
-            contig=contig,
-            samples=self.samples,
-            extended_events=extended_events,
-            slim_burn_in=1,
-            logfile=logfile,
-            logfile_interval=1,
-        )
-        in_sweep, outside_sweep, rejections = self._fitness_per_generation(
-            logfile=logfile,
-            start_generation_ago=start_generation_ago,
-            end_generation_ago=0,
-        )
-        assert start_generation_ago in rejections.keys()
+        while True:
+            engine.simulate(
+                demographic_model=self.model,
+                contig=contig,
+                samples=self.samples,
+                extended_events=extended_events,
+                slim_burn_in=1,
+                logfile=logfile,
+                logfile_interval=1,
+            )
+            in_sweep, outside_sweep, rejections = self._fitness_per_generation(
+                logfile=logfile,
+                start_generation_ago=start_generation_ago,
+                end_generation_ago=0,
+            )
+            # ensure that rejections are occuring in the generation of the AF
+            # condition
+            if start_generation_ago in rejections.keys():
+                break
         assert np.all(outside_sweep == 1.0)
         assert in_sweep[0] >= s * min_freq + 1
 
@@ -2192,21 +2222,25 @@ def test_sweep_meets_min_freq_at_end(self, tmp_path):
             selection_coeff=s,
             dominance_coeff=0.5,
         )
-        engine.simulate(
-            demographic_model=self.model,
-            contig=contig,
-            samples=self.samples,
-            extended_events=extended_events,
-            slim_burn_in=1,
-            logfile=logfile,
-            logfile_interval=1,
-        )
-        in_sweep, outside_sweep, rejections = self._fitness_per_generation(
-            logfile=logfile,
-            start_generation_ago=mutation_generation_ago,
-            end_generation_ago=end_generation_ago,
-        )
-        assert end_generation_ago in rejections.keys()
+        while True:
+            engine.simulate(
+                demographic_model=self.model,
+                contig=contig,
+                samples=self.samples,
+                extended_events=extended_events,
+                slim_burn_in=1,
+                logfile=logfile,
+                logfile_interval=1,
+            )
+            in_sweep, outside_sweep, rejections = self._fitness_per_generation(
+                logfile=logfile,
+                start_generation_ago=mutation_generation_ago,
+                end_generation_ago=end_generation_ago,
+            )
+            # ensure that rejections are occuring in the generation of the AF
+            # condition
+            if end_generation_ago in rejections.keys():
+                break
         assert np.all(outside_sweep == 1.0)
         assert in_sweep[-1] >= s * min_freq + 1