Draft of selective_sweep API

doc/_toc.yml

@@ -31,6 +31,10 @@ parts:
   - file: short_vignettes/demography_vignettes_intro
   - file: short_vignettes/demes_vignette
   - file: short_vignettes/demography_debugger_vignette
+- caption: Short vignettes - special case models
+  chapters:
+  - file: short_vignettes/trackmutationfates
+  - file: short_vignettes/selective_sweep
 - caption: Longer vignettes - exporting data to tskit
   chapters:
   - file: long_vignettes/tskitconvert_vignette
@@ -73,6 +77,7 @@ parts:
   - file: pages/model_params
   - file: pages/regiontypes
   - file: pages/tskit_tools
+  - file: pages/conditional_models
   - file: pages/types
 - caption: Examples
   chapters:

doc/misc/changelog.md

@@ -17,7 +17,12 @@ Bug fixes
   PR {pr}`845`
   PR {pr}`847`
 
-Dependencies
+New features
+
+* Add module {mod}`fwdpy11.conditional_models`.
+  PR {pr}`828`
+
+Deprecations
 
 * Deprecate `fwdpy11.tskit_tools.WrappedTreeSequence`
   PR {pr}`841`

doc/pages/conditional_models.md

@@ -0,0 +1,15 @@
+(conditional_models)=
+
+# Module `fwdpy11.conditional_models`
+
+```{eval-rst}
+.. automodule:: fwdpy11.conditional_models
+   :members:
+   :undoc-members:
+
+.. autofunction:: fwdpy11.conditional_models.track_mutation
+
+.. autofunction:: fwdpy11.conditional_models.selective_sweep
+```
+
+

doc/short_vignettes/selective_sweep.md

@@ -0,0 +1,152 @@
+---
+jupytext:
+  formats: md:myst
+  text_representation:
+    extension: .md
+    format_name: myst
+kernelspec:
+  display_name: Python 3
+  language: python
+  name: python3
+---
+
+(selective_sweeps)=
+
+# Selective sweeps
+
+```{code-cell} python
+---
+tags: ['hide-input']
+---
+
+import fwdpy11
+import numpy as np
+import msprime
+```
+
+```{code-cell} python
+import fwdpy11.conditional_models
+import fwdpy11.tskit_tools
+```
+
+
+```{code-cell} python
+---
+tags: ['hide-input']
+---
+
+def setup(prune_selected=False):
+    # Dropping mutations requires existing
+    # ancestry, which we can get either
+    # from a burn-in or from msprime.
+    initial_ts = msprime.sim_ancestry(
+        samples=500,
+        population_size=500,
+        recombination_rate=1e-1,
+        random_seed=43215,
+        sequence_length=1.0,
+    )
+
+    # Build the pop from msprime output
+    pop = fwdpy11.DiploidPopulation.create_from_tskit(initial_ts)
+
+    # Set up basic model parameters
+    pdict = {
+        "recregions": [fwdpy11.PoissonInterval(0, 1, 1e-1)],
+        # Here, 2 means that fitness is multiplicative
+        # over 1, 1+hs, 1+2s.
+        "gvalue": fwdpy11.Multiplicative(2.0),
+        "rates": (0, 0, None),
+        "prune_selected": False,
+        "simlen": 200,
+    }
+    params = fwdpy11.ModelParams(**pdict)
+
+    return pop, params
+```
+
+## From a new mutation
+
+```{code-cell} python
+ALPHA = 1000.0
+```
+
+
+```{code-cell} python
+rng = fwdpy11.GSLrng(12345)
+pop, params = setup()
+```
+
+```{code-cell} python
+mutation_data = fwdpy11.conditional_models.NewMutationParameters(
+    frequency=fwdpy11.conditional_models.AlleleCount(1),
+    data=fwdpy11.NewMutationData(effect_size=ALPHA / 2 / pop.N, dominance=1),
+    position=fwdpy11.conditional_models.PositionRange(left=0.49, right=0.51),
+)
+```
+
+
+```{code-cell} python
+output = fwdpy11.conditional_models.selective_sweep(
+    rng, 
+    pop,
+    params,
+    mutation_data,
+    fwdpy11.conditional_models.GlobalFixation()
+)
+```
+
+```{code-cell} python
+assert output.pop.generation == params.simlen
+assert pop.generation == 0
+```
+
+```{code-cell} python
+print(output.pop.mutations[output.index])
+```
+
+```{code-cell}
+for fixation, time in zip(output.pop.fixations, output.pop.fixation_times):
+    print(fixation, time)
+```
+
+```{code-cell}
+FIXATION_TIME = output.pop.fixation_times[0]
+```
+
+
+
+### Recording the generation when fixation happened
+
+```{code-cell} python
+rng = fwdpy11.GSLrng(12345)
+```
+
+```{code-cell} python
+output = fwdpy11.conditional_models.selective_sweep(
+    rng, 
+    pop,
+    params,
+    mutation_data,
+    fwdpy11.conditional_models.GlobalFixation(),
+    sampling_policy=fwdpy11.conditional_models.AncientSamplePolicy.COMPLETION,
+)
+```
+
+```{code-cell} python
+assert len(output.pop.ancient_sample_nodes) == 2 * output.pop.N
+assert output.pop.fixation_times[output.index] == FIXATION_TIME
+```
+
+
+```{code-cell} python
+node_array = np.array(output.pop.tables.nodes, copy=False)
+ancient_sample_node_times = \
+    node_array["time"][output.pop.ancient_sample_nodes]
+assert np.all([ancient_sample_node_times == \
+    output.pop.fixation_times[output.index]])
+```
+
+## From a standing variant
+
+The recipes for a standing variant are identical to those show above, except that one uses {class}`fwdpy11.conditional_models.AlleleCountRange` or {class}`fwdpy11.conditional_models.FrequencyRange` to specify the starting frequencies.

doc/short_vignettes/trackmutationfates.md

@@ -0,0 +1,150 @@
+---
+jupytext:
+  formats: md:myst
+  text_representation:
+    extension: .md
+    format_name: myst
+kernelspec:
+  display_name: Python 3
+  language: python
+  name: python3
+---
+
+(tracking_mutation_fates)=
+
+# Need title
+
+```{code-cell} python
+---
+tags: ['hide-input']
+---
+
+import fwdpy11
+import numpy as np
+import msprime
+```
+
+```{code-cell} python
+import fwdpy11.conditional_models
+```
+
+
+```{code-cell} python
+---
+tags: ['hide-input']
+---
+
+def setup(prune_selected=False):
+    # Dropping mutations requires existing
+    # ancestry, which we can get either
+    # from a burn-in or from msprime.
+    initial_ts = msprime.sim_ancestry(
+        samples=500,
+        population_size=500,
+        recombination_rate=1e-1,
+        random_seed=43215,
+        sequence_length=1.0,
+    )
+
+    # Build the pop from msprime output
+    pop = fwdpy11.DiploidPopulation.create_from_tskit(initial_ts)
+
+    # Set up basic model parameters
+    pdict = {
+        "recregions": [fwdpy11.PoissonInterval(0, 1, 1e-1)],
+        # Here, 2 means that fitness is multiplicative
+        # over 1, 1+hs, 1+2s.
+        "gvalue": fwdpy11.Multiplicative(2.0),
+        "rates": (0, 0, None),
+        "prune_selected": False,
+        "simlen": 200,
+    }
+    params = fwdpy11.ModelParams(**pdict)
+
+    return pop, params
+```
+
+## Need title
+
+```{code-cell} python
+ALPHA = -10.0
+```
+
+
+```{code-cell} python
+rng = fwdpy11.GSLrng(12345)
+pop, params = setup()
+```
+
+```{code-cell} python
+mutation_data = fwdpy11.conditional_models.NewMutationParameters(
+    frequency=fwdpy11.conditional_models.AlleleCount(1),
+    data=fwdpy11.NewMutationData(effect_size=ALPHA / 2 / pop.N, dominance=1),
+    position=fwdpy11.conditional_models.PositionRange(left=0.49, right=0.51),
+)
+```
+
+
+```{code-cell} python
+output = fwdpy11.conditional_models.track_mutation(
+    rng, 
+    pop,
+    params,
+    mutation_data,
+    when=3,
+    until=7,
+)
+```
+
+When tracking deleterious variants, it is unlikely that they will be around at the end of the simulation:
+
+```{code-cell} python
+try:
+    print(output.pop.mutations[output.index])
+except IndexError as _:
+    print(f"mutation {output.index} is no longer in the population!") 
+```
+
+### Recording all generations of the mutation's sojourn
+
+So, we've lost all the information about this variant.
+That's not so useful.
+Let's record all generations of its existence as ancient samples:
+
+```{code-cell} python
+rng = fwdpy11.GSLrng(12345)
+```
+
+```{code-cell} python
+output = fwdpy11.conditional_models.track_mutation(
+    rng, 
+    pop,
+    params,
+    mutation_data,
+    when=3,
+    until=7,
+    sampling_policy=fwdpy11.conditional_models.AncientSamplePolicy.DURATION,
+)
+```
+
+Now, our mutation is present in nodes in our tree sequence.
+Let's try to print it again:
+
+```{code-cell} python
+try:
+    print(output.pop.mutations[output.index])
+except IndexError as _:
+    output.index(f"mutation {output.index} is no longer in the population!") 
+```
+
+Let's track this variant's frequency at each time point:
+
+```{code-cell}
+for time, nodes, _ in output.pop.sample_timepoints(include_alive=False):
+    print(time, len(nodes))
+    tree_itr = fwdpy11.TreeIterator(output.pop.tables, nodes)
+    for tree in tree_itr:
+        for mutation in tree.mutations():
+            print(time, tree.leaf_counts(mutation.node), mutation.key)
+``` 
+