allow option to not prune selected fixations during standard pop-gen …

…simulations

fwdpy11/src/wfevolve.cc

@@ -30,9 +30,72 @@
 #include <fwdpy11/samplers.hpp>
 #include <fwdpy11/fitness/fitness.hpp>
 #include <fwdpy11/rules/wf_rules.hpp>
+#include <fwdpy11/sim_functions.hpp>
 
 namespace py = pybind11;
 
+template <typename bound_mmodels, typename bound_recmodels>
+void
+evolve_and_prune_all(const fwdpy11::GSLrng_t& rng, fwdpy11::singlepop_t& pop,
+                     fwdpy11::wf_rules& rules,
+                     py::array_t<std::uint32_t> popsizes,
+                     const double mu_neutral, const double mu_selected,
+                     const bound_mmodels& mmodels,
+                     const bound_recmodels& recmap,
+                     fwdpy11::singlepop_fitness& fitness,
+                     fwdpy11::singlepop_temporal_sampler& recorder,
+                     const double selfing_rate)
+{
+    auto generations = popsizes.size();
+    auto fitness_callback = fitness.callback();
+    for (unsigned generation = 0; generation < generations;
+         ++generation, ++pop.generation)
+        {
+            fitness.update(pop);
+            const auto N_next = popsizes.at(generation);
+            double wbar = KTfwd::experimental::sample_diploid(
+                rng.get(), pop.gametes, pop.diploids, pop.mutations,
+                pop.mcounts, pop.N, N_next, mu_neutral + mu_selected, mmodels,
+                recmap, fitness_callback, pop.neutral, pop.selected,
+                selfing_rate, rules);
+            pop.N = N_next;
+            KTfwd::update_mutations(pop.mutations, pop.fixations,
+                                    pop.fixation_times, pop.mut_lookup,
+                                    pop.mcounts, generation, 2 * pop.N);
+            recorder(pop);
+        }
+}
+
+template <typename bound_mmodels, typename bound_recmodels>
+void
+evolve_and_prune_neutral_fixations(
+    const fwdpy11::GSLrng_t& rng, fwdpy11::singlepop_t& pop,
+    fwdpy11::wf_rules& rules, py::array_t<std::uint32_t> popsizes,
+    const double mu_neutral, const double mu_selected,
+    const bound_mmodels& mmodels, const bound_recmodels& recmap,
+    fwdpy11::singlepop_fitness& fitness,
+    fwdpy11::singlepop_temporal_sampler& recorder, const double selfing_rate)
+{
+    auto generations = popsizes.size();
+    auto fitness_callback = fitness.callback();
+    for (unsigned generation = 0; generation < generations;
+         ++generation, ++pop.generation)
+        {
+            fitness.update(pop);
+            const auto N_next = popsizes.at(generation);
+            double wbar = KTfwd::experimental::sample_diploid(
+                rng.get(), pop.gametes, pop.diploids, pop.mutations,
+                pop.mcounts, pop.N, N_next, mu_neutral + mu_selected, mmodels,
+                recmap, fitness_callback, pop.neutral, pop.selected,
+                selfing_rate, rules, KTfwd::remove_neutral());
+            pop.N = N_next;
+            KTfwd::update_mutations_n(pop.mutations, pop.fixations,
+                                      pop.fixation_times, pop.mut_lookup,
+                                      pop.mcounts, generation, 2 * pop.N);
+            recorder(pop);
+        }
+}
+
 // Evolve the population for some amount of time with mutation and
 // recombination
 void
@@ -43,7 +106,8 @@ evolve_singlepop_regions_cpp(
     const KTfwd::extensions::discrete_mut_model& mmodel,
     const KTfwd::extensions::discrete_rec_model& rmodel,
     fwdpy11::singlepop_fitness& fitness,
-    fwdpy11::singlepop_temporal_sampler recorder, const double selfing_rate)
+    fwdpy11::singlepop_temporal_sampler recorder, const double selfing_rate,
+    const bool remove_selected_fixations = false)
 {
     const auto generations = popsizes.size();
     if (!generations)
@@ -63,7 +127,6 @@ evolve_singlepop_regions_cpp(
             throw std::runtime_error("negative recombination rate: "
                                      + std::to_string(recrate));
         }
-    const auto fitness_callback = fitness.callback();
     pop.mutations.reserve(std::ceil(
         std::log(2 * pop.N)
         * (4. * double(pop.N) * (mu_neutral + mu_selected)
@@ -75,21 +138,17 @@ evolve_singlepop_regions_cpp(
         mu_selected, &pop.generation);
     ++pop.generation;
     auto rules = fwdpy11::wf_rules();
-    for (unsigned generation = 0; generation < generations;
-         ++generation, ++pop.generation)
+    if (remove_selected_fixations)
         {
-            fitness.update(pop);
-            const auto N_next = popsizes.at(generation);
-            double wbar = KTfwd::experimental::sample_diploid(
-                rng.get(), pop.gametes, pop.diploids, pop.mutations,
-                pop.mcounts, pop.N, N_next, mu_neutral + mu_selected, mmodels,
-                recmap, fitness_callback, pop.neutral, pop.selected,
-                selfing_rate, rules);
-            pop.N = N_next;
-            KTfwd::update_mutations(pop.mutations, pop.fixations,
-                                    pop.fixation_times, pop.mut_lookup,
-                                    pop.mcounts, generation, 2 * pop.N);
-            recorder(pop);
+            evolve_and_prune_all(rng, pop, rules, popsizes, mu_neutral,
+                                 mu_selected, mmodels, recmap, fitness,
+                                 recorder, selfing_rate);
+        }
+    else
+        {
+            evolve_and_prune_neutral_fixations(
+                rng, pop, rules, popsizes, mu_neutral, mu_selected, mmodels,
+                recmap, fitness, recorder, selfing_rate);
         }
     --pop.generation;
 }

fwdpy11/wright_fisher.py

@@ -157,7 +157,7 @@ def evolve_regions_sampler(rng,pop,popsizes,mu_neutral,
 
 def evolve_regions_sampler_fitness(rng,pop,popsizes,mu_neutral,
         mu_selected,recrate,nregions,sregions,recregions,fitness,
-        recorder,selfing_rate = 0.):
+        recorder,selfing_rate = 0.,prune_all_fixations=True):
     """
     Evolve a single deme according to a Wright-Fisher life cycle 
     with arbitrary changes in population size, a specified fitness model,
@@ -175,9 +175,11 @@ def evolve_regions_sampler_fitness(rng,pop,popsizes,mu_neutral,
     :param fitness: A :class:`fwdpy11.fitness.SpopFitness`.
     :param recorder: A callable to record data from the population.
     :param selfing_rate: (default 0.0) The probability than an individual selfs.
+    :param prune_all_fixations: (True) Remove fixations affecting fitness from population.
     """
     from .internal import makeMutationRegions,makeRecombinationRegions
     mm=makeMutationRegions(nregions,sregions)
     rm=makeRecombinationRegions(recregions)
     evolve_singlepop_regions_cpp(rng,pop,popsizes,mu_neutral,
-            mu_selected,recrate,mm,rm,fitness,recorder,selfing_rate)
+            mu_selected,recrate,mm,rm,fitness,recorder,selfing_rate,
+            prune_all_fixations)

tests/test_stateful_fitness.py

@@ -32,7 +32,7 @@ def evolve_snowdrift(args):
     nlist = np.array([N]*100,dtype=np.uint32)
     recorder = fp11ts.RecordNothing()
     fp11.wright_fisher.evolve_regions_sampler_fitness(rng,pop,nlist,0.0,0.0025,0.001,
-            [],sregions,recregions,fitness,recorder)
+            [],sregions,recregions,fitness,recorder,prune_all_fixations=False)
     #return our pop
     return pop