ForwardBackwardSpatiallyExplicit.hpp

// Copyright 2021 Arnaud Becheler    <abechele@umich.edu>
 
/*********************************************************************** * This program 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 2 of the License, or    * (at your option) any later version. *
 *                                                                      *
 ***************************************************************************/
 
#ifndef __SPATIALLY_EXPLICIT_H_INCLUDED__
#define __SPATIALLY_EXPLICIT_H_INCLUDED__
 
#include "../coalescence/container/Forest.hpp"
#include "../coalescence/merger_policy.hpp"
#include "../coalescence/occupancy_spectrum/sampling_policy.hpp"
#include "../demography/History.hpp"
 
#include <boost/math/special_functions/binomial.hpp>
#include <map>
 
namespace quetzal::simulator
{
template <typename Space, typename DispersalPolicy, typename CoalescencePolicy,
          typename Memory = quetzal::demography::memory_policy::on_RAM>
class ForwardBackwardSpatiallyExplicit : public CoalescencePolicy
{
  public:
    using coord_type = Space;
    using dispersal_policy_type = DispersalPolicy;
    using N_value_type = typename DispersalPolicy::value_type;
    using memory_policy = Memory;
    template <typename tree_type> using forest_type = quetzal::coalescence::container::Forest<coord_type, tree_type>;
 
  private:
    using history_type = demography::History<coord_type, DispersalPolicy, memory_policy>;
    history_type m_history;
 
  public:
    ForwardBackwardSpatiallyExplicit(coord_type x_0, N_value_type N_0, unsigned int nb_generations)
        : m_history(x_0, N_0, nb_generations)
    {
    }
    auto distribution_area(unsigned int t) const
    {
        assert(t >= 0 && t < m_history.nb_generations());
        return m_history.distribution_area(t);
    }
    auto get_functor_N() const noexcept
    {
        return m_history.get_functor_N();
    }
    template <typename Generator, typename Growth, typename Dispersal>
    void simulate_forward_demography(Growth growth, Dispersal kernel, Generator &gen)
    {
        m_history.simulate_forward(growth, kernel, gen);
    }
    template <typename Merger, typename Generator>
    auto make_forest_and_coalesce_along_spatial_history(std::map<coord_type, unsigned int> sample, Generator &gen)
    {
        test_sample_consistency(sample);
        auto forest = this->make_forest(sample);
        return coalesce_along_spatial_history<Merger>(forest, this->branch(), gen, this->init());
    }
    template <typename Merger, typename Generator>
    auto make_forest_using_sampling_time_and_coalesce_along_spatial_history(std::map<coord_type, unsigned int> sample,
                                                                            unsigned int sampling_time, Generator &gen)
    {
        test_sample_consistency(sample);
        auto forest = this->make_forest(sample, sampling_time);
        return coalesce_along_spatial_history<Merger>(forest, this->branch(), gen, this->init());
    }
 
    template <typename Merger, typename Generator, typename F, typename Tree, typename U>
    auto coalesce_along_spatial_history(forest_type<Tree> forest, F binary_op, Generator &gen, U make_tree)
    {
        this->check_consistency(m_history, forest);
        auto t = m_history.nb_generations();
        while ((forest.nb_trees() > 1) && (t > 0))
        {
            may_coalesce_colocated<Merger>(forest, t, gen, binary_op, make_tree);
            migrate_backward(forest, t, gen);
            --t;
        }
        may_coalesce_colocated<Merger>(forest, t, gen, binary_op, make_tree);
        return forest;
    }
    template <typename Merger = quetzal::coalescence::merger_policy::SimultaneousMultipleMerger<
                  quetzal::coalescence::occupancy_spectrum::sampling_policy::on_the_fly>,
              typename Generator>
    auto coalesce_to_mrca(std::map<coord_type, unsigned int> sample, unsigned int sampling_time, Generator &gen)
    {
        test_sample_consistency(sample);
        auto forest = this->make_forest(sample, sampling_time);
        auto new_forest = coalesce_along_spatial_history<Merger>(forest, this->branch(), gen, this->init());
        auto tree = this->find_mrca(new_forest, 0, gen);
        return this->treat(tree);
    }
 
    template <typename Merger = quetzal::coalescence::merger_policy::SimultaneousMultipleMerger<
                  quetzal::coalescence::occupancy_spectrum::sampling_policy::on_the_fly>,
              typename Generator>
    auto coalesce_to_mrca(std::map<coord_type, unsigned int> sample, Generator &gen)
    {
        test_sample_consistency(sample);
        auto forest = this->make_forest(sample);
        auto new_forest = coalesce_along_spatial_history<Merger>(forest, this->branch(), gen, this->init());
        auto tree = this->find_mrca(new_forest, 0, gen);
        return this->treat(tree);
    }
 
    template <typename Merger = quetzal::coalescence::merger_policy::SimultaneousMultipleMerger<
                  quetzal::coalescence::occupancy_spectrum::sampling_policy::on_the_fly>,
              typename Generator, typename F>
    auto coalesce_to_mrca(std::map<coord_type, unsigned int> sample, unsigned int sampling_time, F leaf_name,
                          Generator &gen)
    {
        test_sample_consistency(sample);
        auto forest = this->make_forest(sample, sampling_time, leaf_name);
        auto new_forest = coalesce_along_spatial_history<Merger>(forest, this->branch(), gen, this->init());
        auto tree = this->find_mrca(new_forest, 0, gen);
        return this->treat(tree);
    }
 
    template <typename Merger = quetzal::coalescence::merger_policy::SimultaneousMultipleMerger<
                  quetzal::coalescence::occupancy_spectrum::sampling_policy::on_the_fly>,
              typename T, typename F1, typename F2, typename Generator>
    auto coalesce_to_mrca(std::vector<T> sample, unsigned int sampling_time, F1 get_location, F2 get_name,
                          Generator &gen)
    {
        auto forest = this->make_forest(sample, sampling_time, get_location, get_name);
        auto new_forest = coalesce_along_spatial_history<Merger>(forest, this->branch(), gen, this->init());
        auto tree = this->find_mrca(new_forest, 0, gen);
        return this->treat(tree);
    }
 
  private:
    template <typename T, typename U> void check_consistency(T &history, U const &forest)
    {
        auto t = history.nb_generations();
        for (auto const &it : forest.positions())
        {
            if (history.pop_size(it, t) < forest.nb_trees(it))
            {
                throw std::domain_error("Simulated population size inferior to sampling size");
            }
        }
    }
 
    void test_sample_consistency(std::map<coord_type, unsigned int> const &sample) const
    {
        if (sample.size() == 0)
        {
            throw std::logic_error("Sample to coalesce is empty");
        }
        else if (sample.size() == 1)
        {
            if (sample.begin()->second < 2)
            {
                throw std::logic_error("Sample to coalesce only contains one gene copy");
            }
        }
    }
 
    template <typename Generator, typename Forest>
    void migrate_backward(Forest &forest, unsigned int t, Generator &gen) const noexcept
    {
        Forest new_forest;
        for (auto const it : forest)
        {
            coord_type x = it.first;
            new_forest.insert(m_history.backward_kernel(x, t, gen), it.second);
        }
        assert(forest.nb_trees() == new_forest.nb_trees());
        forest = new_forest;
    }
 
    template <typename Merger, typename Generator, typename Tree, typename F, typename U>
    void may_coalesce_colocated(forest_type<Tree> &forest, unsigned int t, Generator &gen, F binop,
                                U make_tree) const noexcept
    {
        auto N = get_functor_N();
        for (auto const &x : forest.positions())
        {
            auto range = forest.trees_at_same_position(x);
            std::vector<Tree> v;
            for (auto it = range.first; it != range.second; ++it)
            {
                v.push_back(it->second);
            }
            if (v.size() >= 2)
            {
                assert(N(x, t) >= 1 && "Population size should be positive for evaluating coalescence probability");
                auto last = Merger::merge(v.begin(), v.end(), N(x, t), make_tree(x, t), binop, gen);
                forest.erase(x);
                for (auto it = v.begin(); it != last; ++it)
                {
                    forest.insert(x, *it);
                }
            }
        }
    }
}; // end ForwardBackwardSpatiallyExplicit
} // namespace quetzal::simulator
 
#endif