1#ifndef __PRNG_EXTENSION_HPP__
2#define __PRNG_EXTENSION_HPP__
4#include "Kokkos_Macros.hpp"
6#include <Kokkos_Core.hpp>
7#include <Kokkos_MathematicalConstants.hpp>
8#include <Kokkos_Random.hpp>
10#include <common/traits.hpp>
47 template <
typename T,
typename F,
class DeviceType>
50 { obj.draw(gen) } -> std::same_as<F>;
51 { obj.mean() } -> std::same_as<F>;
52 { obj.var() } -> std::same_as<F>;
53 { obj.skewness() } -> std::same_as<F>;
74 template <FloatingPo
intType F> KOKKOS_INLINE_FUNCTION F
erfinv(F x)
79 constexpr F a = 0.147;
80 constexpr F inv_a = 1. / a;
81 constexpr F tmp = (2 / (M_PI * a));
82 const double ln1mx2 = Kokkos::log((1. - x) * (1. + x));
83 const F term1 = tmp + (0.5 * ln1mx2);
84 const F term2 = inv_a * ln1mx2;
85 return Kokkos::copysign(Kokkos::sqrt(Kokkos::sqrt(term1 * term1 - term2) - term1), x);
123 template <FloatingPo
intType F> KOKKOS_INLINE_FUNCTION F
norminv(F p, F mean, F stddev)
125 constexpr F erfinv_lb = -5;
126 constexpr F erfinv_up = 5;
127 auto clamped_inverse = Kokkos::clamp(
erfinv(2 * p - 1), erfinv_lb, erfinv_up);
128 KOKKOS_ASSERT(Kokkos::isfinite(stddev * Kokkos::numbers::sqrt2 * clamped_inverse));
129 return mean + stddev * Kokkos::numbers::sqrt2 * clamped_inverse;
149 constexpr double inv_sqrt_2_pi = 0.3989422804014327;
150 return inv_sqrt_2_pi * Kokkos::exp(-0.5 * x * x);
169 return 0.5 * (1 + Kokkos::erf(x / Kokkos::numbers::sqrt2));
183 template <FloatingPo
intType F>
struct Normal
194 template <
class DeviceType>
195 KOKKOS_INLINE_FUNCTION F
draw(Kokkos::Random_XorShift1024<DeviceType>& gen)
const
208 template <
class DeviceType>
209 static KOKKOS_INLINE_FUNCTION F
draw_from(Kokkos::Random_XorShift1024<DeviceType>& gen,
220 KOKKOS_INLINE_FUNCTION F
mean()
const
229 KOKKOS_INLINE_FUNCTION F
var()
const
279 template <
class DeviceType>
280 KOKKOS_INLINE_FUNCTION F
draw(Kokkos::Random_XorShift1024<DeviceType>& gen)
const
285 template <
class DeviceType>
286 static KOKKOS_INLINE_FUNCTION F
290 const F rand =
static_cast<F
>(gen.drand());
296 F zl = Kokkos::clamp((
lower -
mu) /
sigma, F(-5e3), F(0));
297 F zu = Kokkos::clamp((
upper -
mu) /
sigma, F(0), F(5e3));
299 F pl = 0.5 * Kokkos::erfc(-zl / Kokkos::numbers::sqrt2);
300 KOKKOS_ASSERT(Kokkos::isfinite(pl) &&
301 "Truncated normal draw leads to Nan of Inf with given parameters");
302 F pu = 0.5 * Kokkos::erfc(-zu / Kokkos::numbers::sqrt2);
303 KOKKOS_ASSERT(Kokkos::isfinite(pu) &&
304 "Truncated normal draw leads to Nan of Inf with given parameters");
305 F p = rand * (pu - pl) + pl;
307 KOKKOS_ASSERT(Kokkos::isfinite(x) &&
308 "Truncated normal draw leads to Nan of Inf with given parameters");
314 KOKKOS_INLINE_FUNCTION F
mean()
const
322 KOKKOS_INLINE_FUNCTION F
var()
const
329 const auto tmp = (phi_a - phi_b) / Z;
330 const auto tmp2 = (alpha * phi_a - beta * phi_b) / Z;
331 return sigma *
sigma * (1 - tmp2 - tmp * tmp);
363 template <
class DeviceType>
365 Kokkos::Random_XorShift1024<DeviceType>& gen, F factor, F mu, F sigma, F lower, F upper)
369 gen, factor * mu, factor * sigma, factor * lower, factor * upper);
372 template <
class DeviceType>
373 KOKKOS_INLINE_FUNCTION F
draw(Kokkos::Random_XorShift1024<DeviceType>& gen)
const
378 KOKKOS_INLINE_FUNCTION F
mean()
const
384 KOKKOS_INLINE_FUNCTION F
var()
const
406 template <
class DeviceType>
407 KOKKOS_INLINE_FUNCTION F
draw(Kokkos::Random_XorShift1024<DeviceType>& gen)
const
409 return Kokkos::exp(gen.normal(
mu,
sigma));
412 KOKKOS_INLINE_FUNCTION F
mean()
const
416 KOKKOS_INLINE_FUNCTION F
var()
const
419 return (Kokkos::exp(sigma2) - 1) * Kokkos::exp(2 *
mu + sigma2);
424 return (Kokkos::exp(sigma2) + 2) * Kokkos::sqrt(Kokkos::exp(sigma2) - 1);
441 template <
class DeviceType>
442 KOKKOS_INLINE_FUNCTION F
draw(Kokkos::Random_XorShift1024<DeviceType>& gen)
const
445 const double Z0 = gen.normal();
446 const double Z1 = gen.normal();
448 const double scale_factor = Kokkos::sqrt(1. + delta * delta);
449 const double X = (Z0 + delta * Kokkos::abs(Z1)) / scale_factor;
452 KOKKOS_INLINE_FUNCTION F
mean()
const
456 KOKKOS_INLINE_FUNCTION F
var()
const
459 return omega *
omega * (1 - 2 * delta * delta / M_PI);
464 return ((4 - M_PI) * Kokkos::pow(delta * std::sqrt(2 / M_PI), 3)) /
465 Kokkos::pow(1 - 2 * delta * delta / M_PI, 1.5);
469 static_assert(ProbabilityLaw<SkewNormal<float>, float, ComputeSpace>);
Concept for probability distribution laws.
Definition prng_extension.hpp:48
Kokkos compatible method to draw from specific probability distribution.
Definition prng_extension.hpp:16
KOKKOS_INLINE_FUNCTION F erfinv(F x)
Computes an approximation of the inverse error function.
Definition prng_extension.hpp:74
KOKKOS_INLINE_FUNCTION F norminv(F p, F mean, F stddev)
Computes the inverse CDF (probit function) of a normal distribution.
Definition prng_extension.hpp:123
KOKKOS_INLINE_FUNCTION F std_normal_pdf(F x)
Computes the standard normal probability density function (PDF).
Definition prng_extension.hpp:146
KOKKOS_INLINE_FUNCTION F std_normal_cdf(F x)
Computes the standard normal cumulative distribution function (CDF).
Definition prng_extension.hpp:166
Represents a LogNormal (Gaussian) probability distribution.
Definition prng_extension.hpp:402
KOKKOS_INLINE_FUNCTION F draw(Kokkos::Random_XorShift1024< DeviceType > &gen) const
Definition prng_extension.hpp:407
F sigma
Definition prng_extension.hpp:404
KOKKOS_INLINE_FUNCTION F mean() const
Definition prng_extension.hpp:412
KOKKOS_INLINE_FUNCTION F skewness() const
Definition prng_extension.hpp:421
KOKKOS_INLINE_FUNCTION F var() const
Definition prng_extension.hpp:416
F mu
Definition prng_extension.hpp:403
Represents a normal (Gaussian) probability distribution.
Definition prng_extension.hpp:184
F mu
Mean.
Definition prng_extension.hpp:185
KOKKOS_INLINE_FUNCTION F draw(Kokkos::Random_XorShift1024< DeviceType > &gen) const
Draws a random sample from the distribution.
Definition prng_extension.hpp:195
static KOKKOS_INLINE_FUNCTION F draw_from(Kokkos::Random_XorShift1024< DeviceType > &gen, F mu, F sigma)
Static method to draw a sample from N(μ, σ).
Definition prng_extension.hpp:209
KOKKOS_INLINE_FUNCTION F mean() const
Returns the mean of the distribution.
Definition prng_extension.hpp:220
KOKKOS_INLINE_FUNCTION F skewness() const
Returns the skewness of the distribution.
Definition prng_extension.hpp:247
KOKKOS_INLINE_FUNCTION F stddev() const
Returns the standard deviation of the distribution.
Definition prng_extension.hpp:238
KOKKOS_INLINE_FUNCTION F var() const
Returns the variance of the distribution.
Definition prng_extension.hpp:229
F sigma
Standard deviation.
Definition prng_extension.hpp:186
Represents a Scaled TruncatedNormal (Gaussian) probability distribution.
Definition prng_extension.hpp:351
F inverse_factor
Definition prng_extension.hpp:354
F scale_factor
Definition prng_extension.hpp:353
TruncatedNormal< F > dist
Definition prng_extension.hpp:355
constexpr ScaledTruncatedNormal(F factor, F m, F s, F l, F u)
Definition prng_extension.hpp:357
KOKKOS_INLINE_FUNCTION F skewness() const
Definition prng_extension.hpp:388
KOKKOS_INLINE_FUNCTION F var() const
Definition prng_extension.hpp:384
static KOKKOS_INLINE_FUNCTION F draw_from(Kokkos::Random_XorShift1024< DeviceType > &gen, F factor, F mu, F sigma, F lower, F upper)
Definition prng_extension.hpp:364
KOKKOS_INLINE_FUNCTION F draw(Kokkos::Random_XorShift1024< DeviceType > &gen) const
Definition prng_extension.hpp:373
KOKKOS_INLINE_FUNCTION F mean() const
Definition prng_extension.hpp:378
Represents a SkewNormal (Gaussian) probability distribution.
Definition prng_extension.hpp:436
F alpha
Definition prng_extension.hpp:439
KOKKOS_INLINE_FUNCTION F skewness() const
Definition prng_extension.hpp:461
F omega
Definition prng_extension.hpp:438
KOKKOS_INLINE_FUNCTION F var() const
Definition prng_extension.hpp:456
KOKKOS_INLINE_FUNCTION F mean() const
Definition prng_extension.hpp:452
F xi
Definition prng_extension.hpp:437
KOKKOS_INLINE_FUNCTION F draw(Kokkos::Random_XorShift1024< DeviceType > &gen) const
Definition prng_extension.hpp:442
Represents a TruncatedNormal (Gaussian) probability distribution.
Definition prng_extension.hpp:263
KOKKOS_INLINE_FUNCTION F mean() const
Definition prng_extension.hpp:314
KOKKOS_INLINE_FUNCTION F skewness() const
Definition prng_extension.hpp:333
KOKKOS_INLINE_FUNCTION F draw(Kokkos::Random_XorShift1024< DeviceType > &gen) const
Definition prng_extension.hpp:280
KOKKOS_INLINE_FUNCTION constexpr TruncatedNormal(F m, F s, F l, F u)
Definition prng_extension.hpp:271
KOKKOS_INLINE_FUNCTION F var() const
Definition prng_extension.hpp:322
F upper
Definition prng_extension.hpp:268
F lower
Definition prng_extension.hpp:267
F sigma
Definition prng_extension.hpp:266
F mu
Definition prng_extension.hpp:265
static KOKKOS_INLINE_FUNCTION F draw_from(Kokkos::Random_XorShift1024< DeviceType > &gen, F mu, F sigma, F lower, F upper)
Definition prng_extension.hpp:287