random_number_generator.h

Go to the documentation of this file.

/* 
* Copyright (C) 2020-2026 MEmilio
*
* Authors: Daniel Abele
*
* Contact: Martin J. Kuehn <Martin.Kuehn@DLR.de>
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
*     http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
 
#ifndef MIO_RANDOM_NUMBER_GENERATOR_H
#define MIO_RANDOM_NUMBER_GENERATOR_H
 
#include "memilio/io/default_serialize.h"
#include "memilio/utils/compiler_diagnostics.h"
#include "memilio/utils/logging.h"
#include "memilio/utils/miompi.h"
#include "memilio/utils/span.h"
#include "memilio/utils/type_safe.h"
 
MSVC_WARNING_DISABLE_PUSH(4127) //conditional expression is constant
GCC_CLANG_DIAGNOSTIC(push)
GCC_CLANG_DIAGNOSTIC(ignored "-Wexpansion-to-defined") //Random123 handles the portability of this warning internally
#include "Random123/array.h"
#include "Random123/threefry.h"
GCC_CLANG_DIAGNOSTIC(pop)
MSVC_WARNING_POP()
 
#include <cassert>
#include <cstdint>
#include <functional>
#include <limits>
#include <numeric>
#include <random>
#include <type_traits>
 
namespace mio
{
 
template <class Derived>
class RandomNumberGeneratorBase
{
public:
    using result_type = uint64_t;
 
    static constexpr result_type min()
    {
        return std::numeric_limits<result_type>::min();
    }
 
    static constexpr result_type max()
    {
        return std::numeric_limits<result_type>::max();
    }
 
    result_type operator()();
};
 
namespace details
{
inline uint64_t to_uint64(r123array2x32 tf_array)
{
    uint64_t i;
    std::memcpy(&i, tf_array.data(), sizeof(uint64_t));
    return i;
}
 
inline r123array2x32 to_r123_array(uint64_t i)
{
    threefry2x32_ctr_t c;
    std::memcpy(c.data(), &i, sizeof(uint64_t));
    return c;
}
} // namespace details
 
template <class T>
struct MEMILIO_ENABLE_EBO Key : TypeSafe<T, Key<T>>, OperatorComparison<Key<T>> {
    static_assert(std::is_unsigned<T>::value, "Underlying Integer type must be unsigned.");
    using TypeSafe<T, Key<T>>::TypeSafe;
};
static_assert(sizeof(Key<uint32_t>) == sizeof(uint32_t), "Empty Base Optimization isn't working.");
 
template <class T>
struct MEMILIO_ENABLE_EBO Counter : TypeSafe<T, Counter<T>>,
                                    OperatorComparison<Counter<T>>,
                                    OperatorAdditionSubtraction<Counter<T>> {
    static_assert(std::is_unsigned<T>::value, "Underlying Integer type must be unsigned.");
    using TypeSafe<T, Counter<T>>::TypeSafe;
};
static_assert(sizeof(Counter<uint32_t>) == sizeof(uint32_t), "Empty Base Optimization isn't working.");
 
template <class Derived>
auto RandomNumberGeneratorBase<Derived>::operator()() -> result_type
{
    //generate a random sample using the Random123 library.
    //Use another threefryNxR algorithm if larger or more random samples are needed than 64 bit.
    auto self = static_cast<Derived*>(this);
    auto c    = static_cast<uint64_t>(self->get_counter().get());
    auto k    = static_cast<uint64_t>(self->get_key().get());
    auto r    = details::to_uint64(threefry2x32(details::to_r123_array(k), details::to_r123_array(c)));
    self->increment_counter();
    return r;
}
 
template <class SeedSeq>
Key<uint64_t> seed_rng_key(SeedSeq& seed_seq)
{
    auto tf_key = threefry2x32_key_t::seed(seed_seq);
    return Key<uint64_t>(details::to_uint64(tf_key));
}
 
template <class UIntC, class UIntN, class CounterS>
Counter<UIntC> rng_totalsequence_counter(UIntN subsequence_idx, CounterS counter)
{
    //use UIntC for variables because it's the biggest integer type in this function
    static const UIntC BITS_PER_BYTE = 8;
    static const UIntC C_BITS        = sizeof(UIntC) * BITS_PER_BYTE;
    static const UIntC S_BITS        = sizeof(CounterS) * BITS_PER_BYTE;
    static const UIntC N_BITS        = C_BITS - S_BITS;
 
    static_assert(S_BITS < C_BITS, "Subsequence counter must be smaller than total sequence counter.");
    static_assert(N_BITS <= C_BITS, "Subsequence index must not be bigger than total sequence counter.");
    static_assert(N_BITS <= sizeof(UIntN) * BITS_PER_BYTE, "Subsequence index must be at least N bits");
 
    assert(UIntC(subsequence_idx) <= (UIntC(1) << N_BITS) &&
           "Subsequence index is too large."); //(1 << N) is the same as (2^N)
 
    //N high bits: subsequence idx
    //S low bits: subsequence counter
    //=> C = N + S bits: total sequence counter
    //example:
    //subsequence index uint32_t(181) = 0x000000B5
    //subsequence counter uint32_t(41309) = 0x0000A15D
    //total sequence counter = 0x000000B50000A15D
    const auto i = static_cast<UIntC>(subsequence_idx);
    const auto s = static_cast<UIntC>(counter.get());
    const auto c = (i << S_BITS) + s; //shift subsequence index to the high bits, add subsequence counter into low bits
    return Counter<UIntC>{c};
}
 
template <class UIntS, class CounterC>
Counter<UIntS> rng_subsequence_counter(CounterC counter)
{
    using UIntC                = typename CounterC::ValueType;
    static const UIntC C_BYTES = sizeof(UIntC);
    static const UIntC S_BYTES = sizeof(UIntS);
 
    static_assert(S_BYTES < C_BYTES, "Subsequence counter must be smaller than total sequence counter.");
 
    //the subsequence counter is in the lower bits of total sequence counter
    //see rng_totalsequence_counter above
    //so we just need to truncate the counter to get the subsequence counter
    return Counter<UIntS>(static_cast<UIntS>(counter.get()));
}
 
class RandomNumberGenerator : public RandomNumberGeneratorBase<RandomNumberGenerator>
{
public:
    RandomNumberGenerator()
        : m_counter(0)
    {
        seed(generate_seeds());
    }
 
    Key<uint64_t> get_key() const
    {
        return m_key;
    }
    Counter<uint64_t> get_counter() const
    {
        return m_counter;
    }
    void set_counter(Counter<uint64_t> counter)
    {
        m_counter = counter;
    }
    void increment_counter()
    {
        ++m_counter;
    }
    static std::vector<uint32_t> generate_seeds()
    {
        std::random_device rd;
        return {rd(), rd(), rd(), rd(), rd(), rd()};
    }
 
    void seed(const std::vector<uint32_t>& seeds)
    {
        m_seeds = seeds;
        std::seed_seq seed_seq(m_seeds.begin(), m_seeds.end());
        m_key = seed_rng_key(seed_seq);
    }
 
    const std::vector<uint32_t> get_seeds() const
    {
        return m_seeds;
    }
 
    void synchronize()
    {
#ifdef MEMILIO_ENABLE_MPI
        int rank;
        MPI_Comm_rank(mpi::get_world(), &rank);
        int num_seeds;
        if (rank == 0) {
            num_seeds = int(m_seeds.size());
        }
        MPI_Bcast(&num_seeds, 1, MPI_INT, 0, mpi::get_world());
        if (rank != 0) {
            m_seeds.assign(num_seeds, 0);
        }
        MPI_Bcast(m_seeds.data(), num_seeds, MPI_UNSIGNED, 0, mpi::get_world());
        if (rank != 0) {
            seed(m_seeds);
        }
#endif
    }
 
    auto default_serialize()
    {
        return Members("RandomNumberGenerator").add("key", m_key).add("counter", m_counter).add("seeds", m_seeds);
    }
 
private:
    Key<uint64_t> m_key;
    Counter<uint64_t> m_counter;
    std::vector<uint32_t> m_seeds;
};
 
RandomNumberGenerator& thread_local_rng();
 
inline void log_rng_seeds(const RandomNumberGenerator& rng, LogLevel level)
{
    const auto& seeds = rng.get_seeds();
    std::stringstream ss;
    bool first = true;
    for (auto& s : seeds) {
        if (!first) {
            ss << ", ";
        }
        first = false;
        ss << s;
    }
    log(level, "Using RNG with seeds: {0}.", ss.str());
}
 
inline void log_thread_local_rng_seeds(LogLevel level)
{
    log_rng_seeds(thread_local_rng(), level);
}
 
template <class DistT>
class DistributionAdapter
{
public:
    using ResultType = typename DistT::result_type;
 
    struct ParamType {
        using DistType = DistributionAdapter<DistT>;
 
        template <typename... Ps>
            requires std::is_constructible_v<typename DistT::param_type, Ps...>
        ParamType(Ps&&... ps)
            : params(std::forward<Ps>(ps)...)
        {
        }
 
        static DistributionAdapter& get_distribution_instance()
        {
            return DistType::get_instance();
        }
 
        typename DistT::param_type params;
    };
    using GeneratorFunction = std::function<ResultType(const typename DistT::param_type& p)>;
 
private:
    DistributionAdapter()                                      = default;
    DistributionAdapter(const DistributionAdapter&)            = default;
    DistributionAdapter& operator=(const DistributionAdapter&) = default;
    DistributionAdapter(DistributionAdapter&&)                 = default;
    DistributionAdapter& operator=(DistributionAdapter&&)      = default;
 
public:
    template <class RNG, class... T>
    ResultType operator()(RNG& rng, T&&... params)
    {
        if (m_generator) {
            //unlikely outside of tests
            return m_generator(typename DistT::param_type{std::forward<T>(params)...});
        }
        else {
            return DistT(std::forward<T>(params)...)(rng);
        }
    }
 
    GeneratorFunction get_generator() const
    {
        return m_generator;
    }
 
    void set_generator(GeneratorFunction g)
    {
        m_generator = g;
    }
 
    static DistributionAdapter& get_instance()
    {
        static DistributionAdapter instance;
        return instance;
    }
 
private:
    GeneratorFunction m_generator;
};
 
template <class Int>
class DiscreteDistributionInPlace
{
public:
    using result_type = Int;
 
    class param_type
    {
    public:
        using distribution = DiscreteDistributionInPlace;
 
        param_type() = default;
 
        param_type(Span<ScalarType> weights)
            : m_weights(weights)
        {
        }
 
        Span<ScalarType> weights() const
        {
            return m_weights;
        }
 
    private:
        Span<ScalarType> m_weights;
    };
 
    DiscreteDistributionInPlace() = default;
 
    DiscreteDistributionInPlace(Span<ScalarType> weights)
        : m_params(weights)
    {
    }
 
    DiscreteDistributionInPlace(param_type params)
        : m_params(params)
    {
    }
 
    void reset()
    {
    }
 
    param_type param()
    {
        return m_params;
    }
 
    void param(param_type p)
    {
        m_params = p;
    }
 
    Span<ScalarType> weights()
    {
        return m_params.weights();
    }
 
    template <class RNG>
    result_type operator()(RNG& rng)
    {
        return (*this)(rng, m_params);
    }
 
    template <class RNG>
    result_type operator()(RNG& rng, param_type p)
    {
        auto weights = p.weights();
        if (weights.size() <= 1) {
            return 0;
        }
        auto sum = std::accumulate(weights.begin(), weights.end(), 0.0);
        auto u   = std::uniform_real_distribution<ScalarType>()(
            rng, std::uniform_real_distribution<ScalarType>::param_type{0.0, sum});
        auto intermediate_sum = 0.0;
        for (size_t i = 0; i < weights.size(); ++i) {
            intermediate_sum += weights.get_ptr()[i];
            if (u < intermediate_sum) {
                return i;
            }
        }
        assert(false && "this should never happen.");
        return result_type(-1);
    }
 
private:
    param_type m_params;
};
 
template <class Int>
using DiscreteDistribution = DistributionAdapter<DiscreteDistributionInPlace<Int>>;
 
template <class Real>
using ExponentialDistribution = DistributionAdapter<std::exponential_distribution<Real>>;
 
template <class Real>
using NormalDistribution = DistributionAdapter<std::normal_distribution<Real>>;
 
template <class Int>
using UniformIntDistribution = DistributionAdapter<std::uniform_int_distribution<Int>>;
 
template <class Real>
using UniformDistribution = DistributionAdapter<std::uniform_real_distribution<Real>>;
 
template <class IOContext, class UniformDistributionParams,
          class Real = typename UniformDistributionParams::DistType::ResultType>
    requires std::is_same_v<UniformDistributionParams, typename UniformDistribution<Real>::ParamType>
void serialize_internal(IOContext& io, const UniformDistributionParams& p)
{
    auto obj = io.create_object("UniformDistributionParams");
    obj.add_element("a", p.params.a());
    obj.add_element("b", p.params.b());
}
 
template <class IOContext, class UniformDistributionParams,
          class Real = typename UniformDistributionParams::DistType::ResultType>
    requires std::is_same_v<UniformDistributionParams, typename UniformDistribution<Real>::ParamType>
IOResult<UniformDistributionParams> deserialize_internal(IOContext& io, Tag<UniformDistributionParams>)
{
    auto obj = io.expect_object("UniformDistributionParams");
    auto a   = obj.expect_element("a", Tag<Real>{});
    auto b   = obj.expect_element("b", Tag<Real>{});
    return apply(
        io,
        [](auto&& a_, auto&& b_) {
            return UniformDistributionParams{a_, b_};
        },
        a, b);
}
 
template <class Int>
using PoissonDistribution = DistributionAdapter<std::poisson_distribution<Int>>;
 
template <class Real>
using LogNormalDistribution = DistributionAdapter<std::lognormal_distribution<Real>>;
 
template <class Real>
using GammaDistribution = DistributionAdapter<std::gamma_distribution<Real>>;
 
template <class Real>
using NormalDistribution = DistributionAdapter<std::normal_distribution<Real>>;
 
} // namespace mio
 
#endif