DiscreteDistribution.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 __DISCRETE_DISTRIBUTION_H_INCLUDED__
#define __DISCRETE_DISTRIBUTION_H_INCLUDED__
 
#include <assert.h>
#include <random>
#include <vector>
 
namespace quetzal::utils::random
{
template <typename State> class DiscreteDistribution
{
  private:
    using SelfType = DiscreteDistribution<State>;
    class Param
    {
      private:
        std::vector<State> m_support;
        std::vector<double> m_weights;
        using distribution_type = DiscreteDistribution<State>;
 
      public:
        Param() : m_support(std::vector<State>(1)), m_weights(std::vector<double>(1, 1.0))
        {
        }
        Param(Param const &other) = default;
        Param(std::vector<State> const &support, std::vector<double> const &weights)
            : m_support(support), m_weights(weights)
        {
        }
        Param(std::vector<State> &&support, std::vector<double> &&weights)
            : m_support(std::move(support)), m_weights(std::move(weights))
        {
        }
        Param &operator=(Param const &other) = default;
        bool operator==(Param const &other)
        {
            return (m_support == other.m_support && m_weights == other.m_weights);
        }
        std::vector<State> const &support() const
        {
            return m_support;
        }
        std::vector<double> const &weights() const
        {
            return m_weights;
        }
    }; // end inner class Param
    Param m_param;
 
  public:
    using result_type = State;
    using param_type = Param;
    DiscreteDistribution() = default;
    explicit DiscreteDistribution(const param_type &p) : m_param(p)
    {
    }
    DiscreteDistribution(DiscreteDistribution<State> const &other) = default;
    DiscreteDistribution(DiscreteDistribution<State> &&other) = default;
    DiscreteDistribution(std::vector<result_type> const &support, std::vector<double> const &weights)
        : m_param(support, weights)
    {
    }
    DiscreteDistribution(std::vector<result_type> &&support, std::vector<double> &&weights) noexcept
        : m_param(std::move(support), std::move(weights))
    {
    }
    DiscreteDistribution<State> &operator=(DiscreteDistribution<State> const &other) = default;
    DiscreteDistribution<State> &operator=(DiscreteDistribution<State> &&other) = default;
    template <typename Generator> result_type operator()(Generator &g) const
    {
        std::discrete_distribution<int> d(m_param.weights().cbegin(), m_param.weights().cend());
        return m_param.support()[d(g)];
    }
    template <typename Generator> result_type operator()(Generator &g, const param_type &params) const
    {
        std::discrete_distribution<int> d(params.weights().cbegin(), params.weights().cend());
        return params.support()[d(g)];
    }
}; // end DiscreteDistribution
} // namespace quetzal::utils::random
 
#endif