base_state.h

// @HEADER
// *******************************************************************************
//                                OpenRAND                                       *
//   A Performance Portable, Reproducible Random Number Generation Library       *
//                                                                               *
// Copyright (c) 2023, Michigan State University                                 *
//                                                                               *
// Permission is hereby granted, free of charge, to any person obtaining a copy  *
// of this software and associated documentation files (the "Software"), to deal *
// in the Software without restriction, including without limitation the rights  *
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell     *
// copies of the Software, and to permit persons to whom the Software is         *
// furnished to do so, subject to the following conditions:                      *
//                                                                               *
// The above copyright notice and this permission notice shall be included in    *
// all copies or substantial portions of the Software.                           *
//                                                                               *
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR    *
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,      *
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE   *
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER        *
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, *
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE *
// SOFTWARE.                                                                     *
//********************************************************************************
// @HEADER
 
#ifndef OPENRAND_BASE_STATE_H_
#define OPENRAND_BASE_STATE_H_
 
#include <openrand/util.h>
 
#include <cstdint>
#include <limits>
#include <type_traits>
 
namespace openrand {
 
template <typename RNG>
class BaseRNG {
 public:
  using result_type = uint32_t;
 
  static constexpr result_type min() {
    return 0u;
  }
  static constexpr result_type max() {
    return ~((result_type)0);
  }
 
  OPENRAND_DEVICE result_type operator()() {
    return gen().template draw<uint32_t>();
  }
 
  template <typename T = float>
  OPENRAND_DEVICE T rand() {
    if constexpr (sizeof(T) <= 4) {
      const uint32_t x = gen().template draw<uint32_t>();
      if constexpr (std::is_integral_v<T>)
        return static_cast<T>(x);
      else
        return u01<float, uint32_t>(x);
    } else {
      const uint64_t x = gen().template draw<uint64_t>();
      if constexpr (std::is_integral_v<T>)
        return static_cast<T>(x);
      else
        return u01<double, uint64_t>(x);
    }
  }
 
  template <typename T = float>
  OPENRAND_DEVICE T uniform(const T low, const T high) {
    // TODO: Allow 64 bit integers
    static_assert(!(std::is_integral_v<T> && sizeof(T) > sizeof(int32_t)),
                  "64 bit int not yet supported");
 
    T r = high - low;
 
    if constexpr (std::is_floating_point_v<T>) {
      return low + r * rand<T>();
    } else if constexpr (std::is_integral_v<T>) {
      return low + range<true, T>(r);
    }
  }
 
  /*
   * @brief Fills an array with random numbers from a uniform distribution [0, 1)
   *
   * @Note This array is filled serially. `N` ideally should not be large.
   */
  template <typename T = float>
  OPENRAND_DEVICE void fill_random(T *array, const int N) {
    for (int i = 0; i < N; i++) array[i] = rand<T>();
  }
 
  template <typename T = float>
  OPENRAND_DEVICE T randn() {
    static_assert(std::is_floating_point_v<T>);
    constexpr T M_PI2 = 2 * static_cast<T>(M_PI);
 
    T u = rand<T>();
    T v = rand<T>();
    T r = openrand::sqrt(T(-2.0) * openrand::log(u));
    T theta = v * M_PI2;
    return r * openrand::cos(theta);
  }
 
  template <typename T = float>
  OPENRAND_DEVICE vec2<T> randn2() {
    // Implements box-muller method
    static_assert(std::is_floating_point_v<T>);
    constexpr T M_PI2 = 2 * static_cast<T>(M_PI);
 
    T u = rand<T>();
    T v = rand<T>();
    T r = sqrt(T(-2.0) * log(u));
    T theta = v * M_PI2;
    return {r * cos(theta), r * sin(theta)};
  }
 
  template <typename T = float>
  OPENRAND_DEVICE T randn(const T mean, const T std_dev) {
    return mean + randn<T>() * std_dev;
  }
 
  template <bool biased = true, typename T = int>
  OPENRAND_DEVICE T range(const T N) {
    // static_assert(std::is_integral_v<T> && sizeof(T) <= sizeof(int32_t),
    //               "64 bit int not yet supported");
 
    uint32_t x = gen().template draw<uint32_t>();
    uint64_t res = static_cast<uint64_t>(x) * static_cast<uint64_t>(N);
 
    if constexpr (biased) {
      return static_cast<T>(res >> 32);
    } else {
      uint32_t leftover = static_cast<uint32_t>(res);
      if (leftover < N) {
        uint32_t threshold = -N % N;
        while (leftover < threshold) {
          x = gen().template draw<uint32_t>();
          res = static_cast<uint64_t>(x) * static_cast<uint64_t>(N);
          leftover = static_cast<uint32_t>(res);
        }
      }
      return static_cast<T>(res);
    }
  }
 
  template <typename T = float>
  OPENRAND_DEVICE inline T gamma(T alpha, T b) {
    T d = alpha - T((1. / 3.));
    T c = T(1.) / sqrt(9.f * d);
    T v, x;
    while (true) {
      do {
        x = randn<T>();
        v = T(1.0) + c * x;
      } while (v <= T(0.));
      v = v * v * v;
      T u = rand<T>();
 
      const T x2 = x * x;
      if (u < 1.0f - 0.0331f * x2 * x2) return (d * v * b);
 
      if (log(u) < 0.5f * x2 + d * (1.0f - v + log(v))) return (d * v * b);
    }
  }
 
  template <typename T = RNG>
  typename std::enable_if_t<has_counter<T>::value, RNG> forward_state(
      int n) const {
    RNG rng = *static_cast<const RNG *>(this);  // copy
    rng._ctr += n;
    return rng;
  }
 
 protected:
  /*
   * @brief Converts a random number integer to a floating point number between [0., 1.)
   */
  template <typename Ftype, typename Utype>
  inline OPENRAND_DEVICE Ftype u01(const Utype in) const {
    constexpr Ftype factor =
        Ftype(1.) / (Ftype(~static_cast<Utype>(0)) + Ftype(1.));
    constexpr Ftype halffactor = Ftype(0.5) * factor;
    return static_cast<Ftype>(in) * factor + halffactor;
  }
 
 private:
  OPENRAND_DEVICE __inline__ RNG &gen() {
    return *static_cast<RNG *>(this);
  }
 
};  // class BaseRNG
 
}  // namespace openrand
 
#endif  // OPENRAND_BASE_STATE_H_