TransitionKernel.hpp

// Copyright 2021 Arnaud Becheler    <abechele@umich.edu> Florence Jornod <florence@jornod.com>
 
/*********************************************************************** * 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 __TRANSITION_KERNEL_H_INCLUDED__
#define __TRANSITION_KERNEL_H_INCLUDED__
 
#include <assert.h>
#include <unordered_map>
 
namespace quetzal::utils::random
{
 
template <class... T> class TransitionKernel;
 
template <typename Distribution> class TransitionKernel<Distribution>
{
 
  private:
    std::unordered_map<typename Distribution::result_type, Distribution> m_distributions;
 
  public:
    using state_type = typename Distribution::result_type;
 
    TransitionKernel()
    {
    }
 
    TransitionKernel(state_type const &x, Distribution const &d)
    {
        m_distributions[x] = d;
    }
 
    TransitionKernel(state_type const &x, Distribution &&d)
    {
        m_distributions[x] = std::move(d);
    }
 
    TransitionKernel(const TransitionKernel &) = default;
 
    TransitionKernel(TransitionKernel &&) = default;
 
    TransitionKernel<Distribution> &operator=(const TransitionKernel<Distribution> &) = default;
 
    TransitionKernel<Distribution> &operator=(TransitionKernel<Distribution> &&) = default;
 
    TransitionKernel<Distribution> &set(state_type const &x, Distribution const &d)
    {
        m_distributions[x] = d;
        return *this;
    }
 
    TransitionKernel<Distribution> &set(state_type const &x, Distribution &&d)
    {
        m_distributions[x] = std::move(d);
        return *this;
    }
 
    bool has_distribution(state_type const &x) const
    {
        return m_distributions.find(x) != m_distributions.end();
    }
 
    template <typename Generator> state_type operator()(Generator &g, state_type const &x)
    {
        assert(has_distribution(x));
        return m_distributions.at(x).operator()(g);
    }
};
 
template <typename Time, typename Distribution> class TransitionKernel<Time, Distribution>
{
 
  private:
    using InsiderType = TransitionKernel<Distribution>;
 
    std::unordered_map<Time, InsiderType> m_kernels;
 
  public:
    using state_type = typename Distribution::result_type;
 
    TransitionKernel()
    {
    }
 
    TransitionKernel(state_type const &x, Time const &t, Distribution const &d)
    {
        m_kernels[t] = InsiderType(x, d);
    }
 
    TransitionKernel(state_type const &x, Time const &t, Distribution &&d)
    {
        m_kernels[t] = InsiderType(x, std::move(d));
    }
 
    TransitionKernel(const TransitionKernel &) = default;
 
    TransitionKernel(TransitionKernel &&) = default;
 
    TransitionKernel<Time, Distribution> &operator=(const TransitionKernel<Time, Distribution> &) = default;
 
    TransitionKernel<Time, Distribution> &operator=(TransitionKernel<Time, Distribution> &&) = default;
 
    TransitionKernel<Time, Distribution> &set(state_type const &x, Time const &t, Distribution const &d)
    {
        m_kernels[t] = InsiderType(x, d);
        return *this;
    }
 
    TransitionKernel<Time, Distribution> &set(state_type const &x, Time const &t, Distribution &&d)
    {
        m_kernels[t] = InsiderType(x, std::move(d));
        return *this;
    }
 
    bool has_distribution(state_type const &x, Time const &t) const
    {
        bool answer = false;
        if (m_kernels.count(t) == 1)
        {
            if (m_kernels.at(t).has_distribution(x))
            {
                answer = true;
            }
        }
        return answer;
    }
 
    template <typename Generator> state_type operator()(Generator &g, state_type const &x, Time const &t)
    {
        assert(this->has_distribution(x, t));
        return m_kernels.at(t).operator()(g, x);
    }
};
 
} // namespace quetzal::utils::random
 
#endif