improve time checks for add mutation

doc/misc/changelog.md

@@ -27,13 +27,14 @@ Behavior changes
 New features
 
 * Add {func}`fwdpy11.DiploidPopulation.add_mutation`.
-  {pr}`764`
+  PR {pr}`764`
+  PR {pr}`799`
 * Add {class}`fwdpy11.NewMutationData`.
-  {pr}`764`
+  PR {pr}`764`
 * Add `__copy__` and `__deepcopy__` to {class}`fwdpy11.DiploidPopulation`.
-  {pr}`770`
+  PR {pr}`770`
 * Add `__deepcopy__` to {class}`fwdpy11.DiscreteDemography`.
-  {pr}`773`
+  PR {pr}`773`
 
 C++ back-end
 

fwdpy11/_types/diploid_population.py

@@ -389,7 +389,11 @@ def add_mutation(
           exactly `ndescendants` alive nodes.
         * From this list of candidate nodes, one is chosen randomly
           to be the node where we place the new mutation.
-        * This node's time is the origin time of the new mutation.
+        * If the node's parent's time is NULL, then the mutations' origin
+          time is the first 64 bit integer ancestral to the node time,
+          which could be the node time itself.
+          Otherwise, a 64 bit integer is chosen uniformly between the
+          node time and the node's parent's time.
         * If `deme` is None, any set of `ndescendants` will be considered.
         * If `deme` is :math:`\geq 0`, all `ndescendants` alive nodes must be
           from that deme.

fwdpy11/src/functions/add_mutation.cc

@@ -21,6 +21,7 @@
 #include <vector>
 #include <stdexcept>
 #include <algorithm>
+#include <sstream>
 #include <fwdpy11/rng.hpp>
 #include <fwdpy11/types/DiploidPopulation.hpp>
 #include <fwdpy11/policies/mutation.hpp>
@@ -81,6 +82,48 @@ namespace
         }
     };
 
+    bool
+    node_has_valid_time(const std::vector<fwdpp::ts::node>& node_table,
+                        const std::int32_t mutation_node,
+                        const std::int32_t mutation_node_parent)
+    {
+        if (mutation_node_parent == fwdpp::ts::NULL_INDEX)
+            {
+                return true;
+            }
+        auto mutation_node_time = node_table[mutation_node].time;
+        // Check for easy case where we can assign mutation time == node time
+        if (mutation_node_time < 0.0)
+            {
+                if (mutation_node_time - std::floor(mutation_node_time) == 0.0)
+                    {
+                        return true;
+                    }
+            }
+        else if (mutation_node_time - std::ceil(mutation_node_time) == 0.0)
+            {
+                return true;
+            }
+
+        auto mutation_node_parent_time = node_table[mutation_node_parent].time;
+        bool valid = false;
+        if (mutation_node_parent_time < 0.0)
+            {
+                if (std::ceil(mutation_node_parent_time) <= mutation_node_time)
+                    {
+                        valid = true;
+                    }
+            }
+        else
+            {
+                if (std::floor(mutation_node_parent_time) <= mutation_node_time)
+                    {
+                        valid = true;
+                    }
+            }
+        return valid;
+    }
+
     std::vector<candidate_node_map>
     generate_canidate_list(const double left, const double right,
                            const unsigned ndescendants,
@@ -139,11 +182,16 @@ namespace
                                                                == deme;
                                                     })))
                                             {
-                                                candidates.emplace_back(
-                                                    std::max(tree.left, left),
-                                                    std::min(tree.right, right), n,
-                                                    tree.parents[n],
-                                                    std::move(descendants));
+                                                if (node_has_valid_time(
+                                                        pop.tables->nodes, n,
+                                                        tree.parents[n]))
+                                                    {
+                                                        candidates.emplace_back(
+                                                            std::max(tree.left, left),
+                                                            std::min(tree.right, right),
+                                                            n, tree.parents[n],
+                                                            std::move(descendants));
+                                                    }
                                             }
                                     }
                                 n = ni();
@@ -156,8 +204,48 @@ namespace
             }
         return candidates;
     }
+
+    std::int64_t
+    generate_mutation_time(const fwdpy11::GSLrng_t& rng,
+                           const std::vector<fwdpp::ts::node>& node_table,
+                           const std::int32_t mutation_node,
+                           const std::int32_t mutation_node_parent)
+    {
+        auto mutation_node_time = node_table[mutation_node].time;
+        // Check for easy case where we can assign mutation time == node time
+        if (mutation_node_time < 0.0)
+            {
+                if (mutation_node_time - std::floor(mutation_node_time) == 0.0)
+                    {
+                        return static_cast<std::int64_t>(std::floor(mutation_node_time));
+                    }
+            }
+        else if (mutation_node_time - std::ceil(mutation_node_time) == 0.0)
+            {
+                return static_cast<std::int64_t>(std::ceil(mutation_node_time));
+            }
+
+        if (mutation_node_parent == fwdpp::ts::NULL_INDEX)
+            {
+                // Time is the closest int64_t to the node time
+                if (mutation_node_time < 0.0)
+                    {
+                        return static_cast<std::int64_t>(std::floor(mutation_node_time));
+                    }
+                return static_cast<std::int64_t>(std::ceil(mutation_node_time));
+            }
+        // choose a random time
+        auto candidate_time = gsl_ran_flat(rng.get(), mutation_node_time,
+                                           node_table[mutation_node_parent].time);
+        if (candidate_time < 0.0)
+            {
+                return static_cast<std::int64_t>(std::floor(mutation_node_time));
+            }
+        return static_cast<std::int64_t>(std::ceil(mutation_node_time));
+    }
 }
 
+// Returns size_t max value if no candidates are found.
 std::size_t
 add_mutation(const fwdpy11::GSLrng_t& rng, const double left, const double right,
              const fwdpp::ts::table_index_t ndescendants,
@@ -226,18 +314,35 @@ add_mutation(const fwdpy11::GSLrng_t& rng, const double left, const double right
         = static_cast<std::size_t>(gsl_ran_flat(rng.get(), 0, candidates.size()));
 
     auto candidate_data = std::move(candidates[candidate]);
-    // TODO: what to do if candidates.empty() == true?
+    auto mutation_node_time = generate_mutation_time(
+        rng, pop.tables->nodes, candidate_data.node, candidate_data.parent);
+    if (mutation_node_time > pop.tables->nodes[candidate_data.node].time)
+        {
+            std::ostringstream o;
+            o << "origin time of " << mutation_node_time
+              << " for mutation is invalid: it must be <= mutation node time of "
+              << pop.tables->nodes[candidate_data.node].time;
+            throw std::runtime_error(o.str());
+        }
+    if (candidate_data.parent != fwdpp::ts::NULL_INDEX)
+        {
+            if (mutation_node_time <= pop.tables->nodes[candidate_data.parent].time)
+                {
+                    std::ostringstream o;
+                    o << "origin time of " << mutation_node_time
+                      << " for mutation is invalid: it must be < parent node time of "
+                      << pop.tables->nodes[candidate_data.parent].time;
+                    throw std::runtime_error(o.str());
+                }
+        }
 
     // NOTE: can we do all of what we need using existing
     // mutate/recombine functions?  Gotta check the back-end.
 
     // The new variant gets a derived state of 1
     auto empty = fwdpp::empty_mutation_queue();
     new_mutation_key = fwdpy11::infsites_Mutation(
-        empty, pop.mutations, pop.mut_lookup, false,
-        // The new mutation's origin time == the node time
-        static_cast<fwdpy11::mutation_origin_time>(
-            pop.tables->nodes[candidate_data.node].time),
+        empty, pop.mutations, pop.mut_lookup, false, mutation_node_time,
         // The lambdas will simply transfer input parameters to the new object.
         [&rng, &candidate_data]() {
             return gsl_ran_flat(rng.get(), candidate_data.left, candidate_data.right);

tests/conftest.py

@@ -18,9 +18,18 @@
 #
 import demes
 import fwdpy11
+import numpy as np
 import pytest
 
 
+@pytest.fixture(scope="function")
+def numpy_generator(request):
+    try:
+        seed = request.param.get("seed", None)
+        return np.random.Generator(np.random.MT19937(seed=seed))
+    except AttributeError as a:  # NOQA
+        raise a
+
 @pytest.fixture(scope="function")
 def rng(request):
     try:

tests/test_add_mutation.py

@@ -22,6 +22,7 @@
 import numpy as np
 import pytest
 
+import tests.test_utilities as test_utilities
 
 # NOTE: this is copied from test/test_tree_sequences.py
 # FIXME: this should be a more general fixture?
@@ -256,3 +257,24 @@ def test_attempt_to_add_to_pop_without_ancestry():
     data = fwdpy11.NewMutationData(effect_size=0.3, dominance=1.0)
     with pytest.raises(ValueError):
         pop.add_mutation(rng, ndescendants=1, data=data)
+
+
+@pytest.mark.parametrize("msprime_seed", test_utilities.seed_list(135123, 10))
+@pytest.mark.parametrize("fp11_seed", test_utilities.seed_list(5130125, 10))
+@pytest.mark.parametrize("ndescendants", [2, 7, 10, 23, 100, 257])
+def test_add_mutation_to_random_msprime_output(msprime_seed, fp11_seed, ndescendants):
+    initial_ts = msprime.sim_ancestry(
+        samples=250,
+        population_size=500,
+        recombination_rate=1e-3,
+        random_seed=msprime_seed,
+        sequence_length=1.0,
+    )
+    pop = fwdpy11.DiploidPopulation.create_from_tskit(initial_ts)
+    rng = fwdpy11.GSLrng(fp11_seed)
+    data = fwdpy11.NewMutationData(effect_size=1e-3, dominance=1.0)
+    idx = pop.add_mutation(
+        rng, ndescendants=ndescendants, data=data, window=(0.49, 0.51)
+    )
+    if idx is not None:
+        _ = pop.dump_tables_to_tskit()

tests/test_selective_sweeps.py

@@ -0,0 +1,125 @@
+#
+# Copyright (C) 2021 Kevin Thornton <krthornt@uci.edu>
+#
+# This file is part of fwdpy11.
+#
+# fwdpy11 is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# fwdpy11 is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with fwdpy11.  If not, see <http://www.gnu.org/licenses/>.
+#
+
+import copy
+import fwdpy11
+import msprime
+import numpy as np
+
+import pytest
+import tests.test_utilities as test_utilities
+
+
+class MutMonitor(object):
+    def __init__(self, index, key):
+        self.index = index
+        self.key = key
+
+    def __call__(self, pop, _):
+        if pop.mcounts[self.index] == 0 or pop.mcounts[self.index] == 2 * pop.N:
+            return True
+        return False
+
+
+def run_selective_sweep(
+    msprime_seed,
+    fp11_seed,
+    ndescendants,
+    N=500,
+    alpha=1000,  # 2Ns
+    rho=100,
+    L=1.0,
+    max_attempts=100,
+):
+    for initial_ts in msprime.sim_ancestry(
+        samples=N,
+        population_size=1000,
+        recombination_rate=rho / 4 / N,
+        random_seed=msprime_seed,
+        sequence_length=1.0,
+        num_replicates=max_attempts,
+    ):
+        pop = fwdpy11.DiploidPopulation.create_from_tskit(initial_ts)
+        pdict = {
+            "recregions": [fwdpy11.PoissonInterval(0, 1, 5e-2)],
+            "gvalue": fwdpy11.Multiplicative(2.0),
+            "rates": (0, 0, None),
+            "prune_selected": False,
+            "simlen": 10 * pop.N,
+        }
+        params = fwdpy11.ModelParams(**pdict)
+
+        rng = fwdpy11.GSLrng(fp11_seed)
+        data = fwdpy11.NewMutationData(effect_size=alpha / 2 / N, dominance=1.0)
+        idx = pop.add_mutation(
+            rng, ndescendants=ndescendants, data=data, window=(0.49, 0.51)
+        )
+        if idx is None:
+            continue
+
+        # Make sure we've chosen a valid time
+        assert (
+            pop.mutations[idx].g <= pop.tables.nodes[pop.tables.mutations[0].node].time
+        ), f"{pop.mutations[idx].g} {pop.tables.nodes[pop.tables.mutations[0].node].time}"
+        monitor = MutMonitor(idx, pop.mutations[idx].key)
+
+        fixed = False
+        pop_with_fixation = None
+
+        while fixed is not True:
+            pcopy = copy.deepcopy(pop)
+            assert pcopy.generation == 0
+            assert len(pcopy.mutations) == 1
+
+            fwdpy11.evolvets(
+                rng,
+                pcopy,
+                params,
+                100,
+                track_mutation_counts=True,
+                stopping_criterion=monitor,
+                suppress_table_indexing=True,
+            )
+
+            for f in pcopy.fixations:
+                if f.key == monitor.key:
+                    pop_with_fixation = pcopy
+                    fixed = True
+
+        return pop_with_fixation, idx
+    return None, None
+
+
+@pytest.mark.parametrize("msprime_seed", test_utilities.seed_list(135123, 5))
+@pytest.mark.parametrize("fp11_seed", test_utilities.seed_list(5130125, 5))
+def test_sweep_from_new_mutation(msprime_seed, fp11_seed):
+    pop_with_fixation, idx = run_selective_sweep(msprime_seed, fp11_seed, 1)
+    if pop_with_fixation is not None:
+        assert pop_with_fixation.mcounts[idx] == 2 * pop_with_fixation.N
+        _ = pop_with_fixation.dump_tables_to_tskit()
+
+
+@pytest.mark.parametrize("msprime_seed", test_utilities.seed_list(135123, 5))
+@pytest.mark.parametrize("fp11_seed", test_utilities.seed_list(5130125, 5))
+@pytest.mark.parametrize("ndescendants", [2, 7, 10, 23, 100, 257])
+def test_sweep_from_standing_variation(msprime_seed, fp11_seed, ndescendants):
+    pop_with_fixation, idx = run_selective_sweep(msprime_seed, fp11_seed, ndescendants)
+    if pop_with_fixation is not None:
+        assert pop_with_fixation.mcounts[idx] == 2 * pop_with_fixation.N
+        _ = pop_with_fixation.dump_tables_to_tskit()