truncated_normal.hpp

Go to the documentation of this file.

//=======================================================================
// Copyright (c) 2014-2023 Baptiste Wicht
// Distributed under the terms of the MIT License.
// (See accompanying file LICENSE or copy at
//  http://opensource.org/licenses/MIT)
//=======================================================================

#pragma once

#include <chrono> //for std::time

namespace etl {

template <typename T = double>
struct truncated_normal_generator_op {
    using value_type = T; 

    random_engine rand_engine;                         
    std::normal_distribution<value_type> distribution; 

    static constexpr bool gpu_computable = false; 

    truncated_normal_generator_op(T mean, T stddev) : rand_engine(std::time(nullptr)), distribution(mean, stddev) {}

    value_type operator()() {
        auto x = distribution(rand_engine);

        while (std::abs(x - distribution.mean()) > 2.0 * distribution.stddev()) {
            x = distribution(rand_engine);
        }

        return x;
    }

    friend std::ostream& operator<<(std::ostream& os, [[maybe_unused]] const truncated_normal_generator_op& s) {
        return os << "TN(0,1)";
    }
};

template <typename G, typename T = double>
struct truncated_normal_generator_g_op {
    using value_type = T; 

    G& rand_engine;                                    
    std::normal_distribution<value_type> distribution; 

    static constexpr bool gpu_computable = false; 

    truncated_normal_generator_g_op(G& g, T mean, T stddev) : rand_engine(g), distribution(mean, stddev) {}

    value_type operator()() {
        auto x = distribution(rand_engine);

        while (std::abs(x - distribution.mean()) > 2.0 * distribution.stddev()) {
            x = distribution(rand_engine);
        }

        return x;
    }

    friend std::ostream& operator<<(std::ostream& os, [[maybe_unused]] const truncated_normal_generator_g_op& s) {
        return os << "TN(0,1)";
    }
};

} //end of namespace etl