Move sweep API and extended events from stdpopsim.ext to stdpopsim

.github/workflows/tests.yml

@@ -41,7 +41,6 @@ jobs:
         if: runner.os == 'macOS'
         run: |
           # Fix incorrect conda permissions on mac that prevent cache restore.
-          echo "$USER:staff ${{ steps.find-conda.outputs.CONDA }}"
           sudo chown -R $USER:staff ${{ steps.find-conda.outputs.CONDA }}
 
       - name: cache conda

CHANGELOG.rst

@@ -11,6 +11,9 @@
     instead of a single value. This will not affect anyone who is not
     parsing the metadata related to DFEs.
 
+- SLiM extended events and selective sweep infrastructure have been
+  moved from the `stdpopsim.ext` namespace into `stdpopsim` proper
+
 **New features**:
 
 - *Relationship between dominance and selection coefficient:*

docs/api.rst

@@ -169,6 +169,30 @@ that applies to the entire Contig.
 
 .. _sec_api_generic_models:
 
+
+.. _sec_api_sweeps:
+
+********************
+Selection and sweeps
+********************
+
+:class:`ExtendedEvent` and subclasses may be used to condition on sequences of
+events at particular loci using the SLiM engine, by passing lists of events to
+the `extended_events` argument in `Engine.simulate`. A simplified API is provided
+to construct the necessary events for selective sweeps.
+
+.. autoclass:: stdpopsim.DrawMutation()
+    :members:
+
+.. autoclass:: stdpopsim.ChangeMutationFitness()
+    :members:
+
+.. autoclass:: stdpopsim.ConditionOnAlleleFrequency()
+    :members:
+
+.. autofunction:: stdpopsim.selective_sweep
+
+
 **************
 Generic models
 **************

docs/selection_example.py

@@ -62,7 +62,7 @@ def adaptive_introgression(seed):
     # generations before present.
     extended_events = [
         # Draw mutation in DenA.
-        stdpopsim.ext.DrawMutation(
+        stdpopsim.DrawMutation(
             time=T_mut,
             single_site_id=locus_id,
             population="DenA",
@@ -74,11 +74,11 @@ def adaptive_introgression(seed):
         # 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(
+        stdpopsim.ConditionOnAlleleFrequency(
             # Note: if T_mut ~= T_Den_split, then we end up with:
             #       GenerationAfter(T_mut) < T_Den_split,
             #       which will give an error due to "start_time < end_time".
-            start_time=stdpopsim.ext.GenerationAfter(T_mut),
+            start_time=stdpopsim.GenerationAfter(T_mut),
             end_time=T_Den_split,
             single_site_id=locus_id,
             population="DenA",
@@ -87,8 +87,8 @@ def adaptive_introgression(seed):
         ),
         # Denisovans split into DenA and Den1 at time T_Den_split,
         # so now we condition on having AF > 0 in Den1.
-        stdpopsim.ext.ConditionOnAlleleFrequency(
-            start_time=stdpopsim.ext.GenerationAfter(T_Den_split),
+        stdpopsim.ConditionOnAlleleFrequency(
+            start_time=stdpopsim.GenerationAfter(T_Den_split),
             end_time=T_mig,
             single_site_id=locus_id,
             population="Den1",
@@ -97,8 +97,8 @@ def adaptive_introgression(seed):
         ),
         # The Den1 lineage has migrants entering the Papaun lineage at T_mig,
         # so condition on AF > 0 in Papuans.
-        stdpopsim.ext.ConditionOnAlleleFrequency(
-            start_time=stdpopsim.ext.GenerationAfter(T_mig),
+        stdpopsim.ConditionOnAlleleFrequency(
+            start_time=stdpopsim.GenerationAfter(T_mig),
             end_time=0,
             single_site_id=locus_id,
             population="Papuan",
@@ -108,7 +108,7 @@ def adaptive_introgression(seed):
         # The mutation is positively selected in Papuans at T_sel.
         # Note that this will have no effect, unless/until a mutation with the
         # specified mutation_type_id is found in the population.
-        stdpopsim.ext.ChangeMutationFitness(
+        stdpopsim.ChangeMutationFitness(
             start_time=T_sel,
             end_time=0,
             single_site_id=locus_id,
@@ -117,7 +117,7 @@ def adaptive_introgression(seed):
             dominance_coeff=0.5,
         ),
         # Condition on AF > 0.05 in Papuans at the end of the simulation.
-        stdpopsim.ext.ConditionOnAlleleFrequency(
+        stdpopsim.ConditionOnAlleleFrequency(
             start_time=0,
             end_time=0,
             single_site_id=locus_id,

docs/tutorial.rst

@@ -1070,7 +1070,7 @@ resulting simulation:
 .. code-block:: python
 
     selection_coeffs = [
-        stdpopsim.ext.selection_coeff_from_mutation(ts, mut) for mut in ts.mutations()
+        stdpopsim.selection_coeff_from_mutation(ts, mut) for mut in ts.mutations()
     ]
     num_neutral = sum([s == 0 for s in selection_coeffs])
     print(
@@ -1131,7 +1131,7 @@ We'll count up the number of neutral and deleterious mutations in the three regi
         region = np.digitize(site.position, gene_interval.flatten())
         for mut in site.mutations:
             selection_coeffs[region].append(
-                stdpopsim.ext.selection_coeff_from_mutation(ts, mut)
+                stdpopsim.selection_coeff_from_mutation(ts, mut)
             )
 
     for region, coeffs in enumerate(selection_coeffs):
@@ -1292,7 +1292,7 @@ frequency :math:`1 / 2N`.
     "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
+This is done through :func:`stdpopsim.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 (named "pop_0" by default); that the
@@ -1302,7 +1302,7 @@ greater than 0.8.
 
 .. code-block:: python
 
-    extended_events = stdpopsim.ext.selective_sweep(
+    extended_events = stdpopsim.selective_sweep(
         single_site_id=locus_id,
         population="pop_0",
         selection_coeff=0.5,

stdpopsim/__init__.py

@@ -20,13 +20,11 @@
 from .engines import *  # NOQA
 from .warning_categories import *  # NOQA
 
-# Extensions.
-from . import ext  # NOQA
-
 # We import catalog here, but the internal functions
 # defined are not part of the external API.
 from .catalog import *  # NOQA
 
 from . import qc  # NOQA
 
+from .selection import *  # NOQA
 from .slim_engine import *  # NOQA

stdpopsim/slim_engine.py

@@ -908,7 +908,7 @@
                 continue
             t = getattr(event, attr)
             t_rounded = round(float(t) / scaling_factor) * scaling_factor
-            if isinstance(t, stdpopsim.ext.GenerationAfter):
+            if isinstance(t, stdpopsim.GenerationAfter):
                 t_rounded -= scaling_factor
             if t_rounded < 0:
                 raise ValueError(f"Bad {attr}: {getattr(event, attr)}")
@@ -1000,7 +1000,7 @@
     drawn_mutations = []
     fitness_callbacks = []
     condition_on_allele_frequency = []
-    op_id = stdpopsim.ext.ConditionOnAlleleFrequency.op_id
+    op_id = stdpopsim.ConditionOnAlleleFrequency.op_id
     slim_mutation_ids = _dfe_to_mtypes(contig)
     drawn_single_site_ids = collections.defaultdict(int)
     referenced_single_site_ids = set()
@@ -1061,7 +1061,7 @@
         if hasattr(ee, "start_time") and hasattr(ee, "end_time"):
             # Now that GenerationAfter times have been accounted for, we can
             # properly catch invalid start/end times.
-            stdpopsim.ext.validate_time_range(ee.start_time, ee.end_time)
+            stdpopsim.validate_time_range(ee.start_time, ee.end_time)
         if hasattr(ee, "population"):
             # Convert population name to integer index. "-1" is used to encode
             # all populations (currently only valid for ChangeMutationFitness).
@@ -1078,7 +1078,7 @@
                 population_id = population_name_to_index[ee.population]
 
         # Append attributes to lists per event type
-        if isinstance(ee, stdpopsim.ext.DrawMutation):
+        if isinstance(ee, stdpopsim.DrawMutation):
             assert population_id >= 0
             drawn_mutations.append(
                 (
@@ -1090,7 +1090,7 @@
                 )
             )
             drawn_single_site_ids[ee.single_site_id] += 1
-        elif isinstance(ee, stdpopsim.ext.ChangeMutationFitness):
+        elif isinstance(ee, stdpopsim.ChangeMutationFitness):
             fitness_callbacks.append(
                 (
                     ee.start_time * demographic_model.generation_time,
@@ -1102,7 +1102,7 @@
                 )
             )
             referenced_single_site_ids.add(ee.single_site_id)
-        elif isinstance(ee, stdpopsim.ext.ConditionOnAlleleFrequency):
+        elif isinstance(ee, stdpopsim.ConditionOnAlleleFrequency):
             assert population_id >= 0
             condition_on_allele_frequency.append(
                 (
@@ -1432,7 +1432,7 @@
 
     # Allele frequency conditioning
     op_types = ", ".join(
-        f'"{op}"' for op in stdpopsim.ext.ConditionOnAlleleFrequency.op_types
+        f'"{op}"' for op in stdpopsim.ConditionOnAlleleFrequency.op_types
     )
     printsc(f'    defineConstant("op_types", c({op_types}));')
     printsc("    // Allele frequency conditioning, one row for each.")
@@ -1560,8 +1560,8 @@
 
         :param seed: The seed for the random number generator.
         :type seed: int
-        :param extended_events: A list of :class:`ext.ExtendedEvents` to be
-            passed to SLiM, e.g. produced by :func:`ext.selective_sweep()`.
+        :param extended_events: A list of :class:`ExtendedEvents` to be
+            passed to SLiM, e.g. produced by :func:`selective_sweep()`.
         :type extended_events: list
         :param slim_path: The full path to the slim executable, or the name of
             a command in the current PATH.