noise.hpp

//=======================================================================
// 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 "etl/impl/egblas/logistic_noise.hpp"

namespace etl {

template <typename T>
struct uniform_noise_unary_op {
    static constexpr bool linear      = true;  
    static constexpr bool thread_safe = false; 

    template <vector_mode_t V>
    static constexpr bool vectorizable = false;

    template <typename E>
    static constexpr bool gpu_computable = false;

    static constexpr int complexity() {
        return 1;
    }

    static T apply(const T& x) {
        static random_engine rand_engine(std::time(nullptr));
        static std::uniform_real_distribution<double> real_distribution(0.0, 1.0);

        return x + real_distribution(rand_engine);
    }

    static std::string desc() noexcept {
        return "uniform_noise";
    }
};

template <typename G, typename T>
struct uniform_noise_unary_g_op {
    static constexpr bool linear      = true;  
    static constexpr bool thread_safe = false; 

private:
    G& rand_engine; 

public:
    explicit uniform_noise_unary_g_op(G& rand_engine) : rand_engine(rand_engine) {
        //Nothing else to init
    }

    template <vector_mode_t V>
    static constexpr bool vectorizable = false;

    template <typename E>
    static constexpr bool gpu_computable = false;

    static constexpr int complexity() {
        return 1;
    }

    T apply(const T& x) const {
        std::uniform_real_distribution<double> real_distribution(0.0, 1.0);

        return x + real_distribution(rand_engine);
    }

    static std::string desc() noexcept {
        return "uniform_noise";
    }
};

template <typename T>
struct normal_noise_unary_op {
    static constexpr bool linear      = true;  
    static constexpr bool thread_safe = false; 

    template <vector_mode_t V>
    static constexpr bool vectorizable = false;

    template <typename E>
    static constexpr bool gpu_computable = false;

    static constexpr int complexity() {
        return 1;
    }

    static T apply(const T& x) {
        static random_engine rand_engine(std::time(nullptr));
        static std::normal_distribution<T> normal_distribution(0.0, 1.0);

        return x + normal_distribution(rand_engine);
    }

    static std::string desc() noexcept {
        return "normal_noise";
    }
};

template <typename G, typename T>
struct normal_noise_unary_g_op {
    static constexpr bool linear      = true;  
    static constexpr bool thread_safe = false; 

private:
    G& rand_engine; 

public:
    explicit normal_noise_unary_g_op(G& rand_engine) : rand_engine(rand_engine) {
        //Nothing else to init
    }

    template <vector_mode_t V>
    static constexpr bool vectorizable = false;

    template <typename E>
    static constexpr bool gpu_computable = false;

    static constexpr int complexity() {
        return 1;
    }

    T apply(const T& x) const {
        std::normal_distribution<T> normal_distribution(0.0, 1.0);

        return x + normal_distribution(rand_engine);
    }

    static std::string desc() noexcept {
        return "normal_noise";
    }
};

template <typename T>
struct logistic_noise_unary_op {
    static constexpr bool linear      = true;  
    static constexpr bool thread_safe = false; 

    template <vector_mode_t V>
    static constexpr bool vectorizable = false;

    template <typename E>
    static constexpr bool gpu_computable = (is_single_precision_t<T> && impl::egblas::has_slogistic_noise_seed)
                                           || (is_double_precision_t<T> && impl::egblas::has_dlogistic_noise_seed);

    static constexpr int complexity() {
        return 1;
    }

    static T apply(const T& x) {
        static random_engine rand_engine(std::time(nullptr));

        std::normal_distribution<T> noise_distribution(0.0, math::logistic_sigmoid(x));

        return x + noise_distribution(rand_engine);
    }

    template <typename X, typename Y>
    static auto gpu_compute_hint(const X & x, Y& y) noexcept {
        static random_engine rand_engine(std::time(nullptr));

        std::uniform_int_distribution<long> seed_dist;

        decltype(auto) t1 = smart_gpu_compute_hint(x, y);

        auto t2 = force_temporary_gpu_dim_only(t1);

        T alpha(1);
        impl::egblas::logistic_noise_seed(etl::size(y), alpha, t1.gpu_memory(), 1, t2.gpu_memory(), 1, seed_dist(rand_engine));

        return t2;
    }
    template <typename X, typename Y>
    static Y& gpu_compute(const X & x, Y& y) noexcept {
        static random_engine rand_engine(std::time(nullptr));

        std::uniform_int_distribution<long> seed_dist;

        decltype(auto) t1 = select_smart_gpu_compute(x, y);

        T alpha(1);
        impl::egblas::logistic_noise_seed(etl::size(y), alpha, t1.gpu_memory(), 1, y.gpu_memory(), 1, seed_dist(rand_engine));

        y.validate_gpu();
        y.invalidate_cpu();

        return y;
    }

    static std::string desc() noexcept {
        return "logistic_noise";
    }
};

template <typename G, typename T>
struct logistic_noise_unary_g_op {
    static constexpr bool linear      = true;  
    static constexpr bool thread_safe = false; 

private:
    G& rand_engine; 

public:
    explicit logistic_noise_unary_g_op(G& rand_engine) : rand_engine(rand_engine) {
        //Nothing else to init
    }

    template <vector_mode_t V>
    static constexpr bool vectorizable = false;

    template <typename E>
    static constexpr bool gpu_computable = (is_single_precision_t<T> && impl::egblas::has_slogistic_noise_seed)
                                           || (is_double_precision_t<T> && impl::egblas::has_dlogistic_noise_seed);

    static constexpr int complexity() {
        return 1;
    }

    T apply(const T& x) const {
        std::normal_distribution<T> noise_distribution(0.0, math::logistic_sigmoid(x));

        return x + noise_distribution(rand_engine);
    }

    template <typename X, typename Y>
    auto gpu_compute_hint(const X & x, Y& y) const noexcept {
        std::uniform_int_distribution<long> seed_dist;

        decltype(auto) t1 = smart_gpu_compute_hint(x, y);

        auto t2 = force_temporary_gpu_dim_only(t1);

        T alpha(1);
        impl::egblas::logistic_noise_seed(etl::size(y), alpha, t1.gpu_memory(), 1, t2.gpu_memory(), 1, seed_dist(rand_engine));

        return t2;
    }
    template <typename X, typename Y>
    Y& gpu_compute(const X & x, Y& y) const noexcept {
        std::uniform_int_distribution<long> seed_dist;

        decltype(auto) t1 = select_smart_gpu_compute(x, y);

        T alpha(1);
        impl::egblas::logistic_noise_seed(etl::size(y), alpha, t1.gpu_memory(), 1, y.gpu_memory(), 1, seed_dist(rand_engine));

        y.validate_gpu();
        y.invalidate_cpu();

        return y;
    }

    static std::string desc() noexcept {
        return "logistic_noise";
    }
};

template <typename T>
struct state_logistic_noise_unary_op {
    static constexpr bool linear      = true;  
    static constexpr bool thread_safe = false; 

    template <vector_mode_t V>
    static constexpr bool vectorizable = false;

    template <typename E>
    static constexpr bool gpu_computable = (is_single_precision_t<T> && impl::egblas::has_slogistic_noise_seed)
                                           || (is_double_precision_t<T> && impl::egblas::has_dlogistic_noise_seed);

    static constexpr int complexity() {
        return 1;
    }

    mutable random_engine  rand_engine; 
    std::shared_ptr<void*> states;      

    state_logistic_noise_unary_op() : rand_engine(std::time(nullptr)) {
        if constexpr (impl::egblas::has_logistic_noise_prepare) {
            states  = std::make_shared<void*>();
            *states = impl::egblas::logistic_noise_prepare();
        }
    }

    explicit state_logistic_noise_unary_op(const std::shared_ptr<void*> states) : rand_engine(std::time(nullptr)) {
        if constexpr (impl::egblas::has_logistic_noise_prepare) {
            this->states = states;

            if (!*this->states) {
                std::uniform_int_distribution<long> seed_dist;
                *this->states = impl::egblas::logistic_noise_prepare_seed(seed_dist(rand_engine));
            }
        }
    }

    T apply(const T& x) const {
        std::normal_distribution<double> noise_distribution(0.0, math::logistic_sigmoid(x));

        return x + noise_distribution(rand_engine);
    }

    template <typename X, typename Y>
    auto gpu_compute_hint(const X & x, Y& y) const noexcept {
        decltype(auto) t1 = smart_gpu_compute_hint(x, y);

        auto t2 = force_temporary_gpu_dim_only(t1);

        T alpha(1);
        impl::egblas::logistic_noise_states(etl::size(y), alpha, t1.gpu_memory(), 1, t2.gpu_memory(), 1, *states);

        return t2;
    }
    template <typename X, typename Y>
    Y& gpu_compute(const X & x, Y& y) const noexcept {
        decltype(auto) t1 = select_smart_gpu_compute(x, y);

        T alpha(1);
        impl::egblas::logistic_noise_states(etl::size(y), alpha, t1.gpu_memory(), 1, y.gpu_memory(), 1, *states);

        y.validate_gpu();
        y.invalidate_cpu();

        return y;
    }

    static std::string desc() noexcept {
        return "state_logistic_noise";
    }
};

template <typename G, typename T>
struct state_logistic_noise_unary_g_op {
    static constexpr bool linear      = true;  
    static constexpr bool thread_safe = false; 

private:
    G&                     rand_engine; 
    std::shared_ptr<void*> states;      

public:
    explicit state_logistic_noise_unary_g_op(G& rand_engine) : rand_engine(rand_engine) {
        if constexpr (impl::egblas::has_logistic_noise_prepare) {
            std::uniform_int_distribution<long> seed_dist;

            states  = std::make_shared<void*>();
            *states = impl::egblas::logistic_noise_prepare_seed(seed_dist(rand_engine));
        }
    }

    state_logistic_noise_unary_g_op(G& rand_engine, const std::shared_ptr<void*> & states) : rand_engine(rand_engine) {
        if constexpr (impl::egblas::has_logistic_noise_prepare) {
            this->states = states;

            if (!*this->states) {
                std::uniform_int_distribution<long> seed_dist;
                *this->states = impl::egblas::logistic_noise_prepare_seed(seed_dist(rand_engine));
            }
        }
    }

    template <vector_mode_t V>
    static constexpr bool vectorizable = false;

    template <typename E>
    static constexpr bool gpu_computable = (is_single_precision_t<T> && impl::egblas::has_slogistic_noise_states)
                                           || (is_double_precision_t<T> && impl::egblas::has_dlogistic_noise_states);

    static constexpr int complexity() {
        return 1;
    }

    T apply(const T& x) const {
        std::normal_distribution<double> noise_distribution(0.0, math::logistic_sigmoid(x));

        return x + noise_distribution(rand_engine);
    }

    template <typename X, typename Y>
    auto gpu_compute_hint(const X & x, Y& y) const noexcept {
        decltype(auto) t1 = smart_gpu_compute_hint(x, y);

        auto t2 = force_temporary_gpu_dim_only(t1);

        T alpha(1);
        impl::egblas::logistic_noise_states(etl::size(y), alpha, t1.gpu_memory(), 1, t2.gpu_memory(), 1, *states);

        return t2;
    }
    template <typename X, typename Y>
    Y& gpu_compute(const X & x, Y& y) const noexcept {
        decltype(auto) t1 = select_smart_gpu_compute(x, y);

        T alpha(1);
        impl::egblas::logistic_noise_states(etl::size(y), alpha, t1.gpu_memory(), 1, y.gpu_memory(), 1, *states);

        y.validate_gpu();
        y.invalidate_cpu();

        return y;
    }

    static std::string desc() noexcept {
        return "state_logistic_noise";
    }
};

} //end of namespace etl