random/distribution/GammaDistribution.hh

0001 //----------------------------------*-C++-*----------------------------------//
0002 // Copyright 2021-2024 UT-Battelle, LLC, and other Celeritas developers.
0003 // See the top-level COPYRIGHT file for details.
0004 // SPDX-License-Identifier: (Apache-2.0 OR MIT)
0005 //---------------------------------------------------------------------------//
0006 //! \file celeritas/random/distribution/GammaDistribution.hh
0007 //---------------------------------------------------------------------------//
0008 #pragma once
0009
0010 #include <cmath>
0011
0012 #include "corecel/Assert.hh"
0013 #include "corecel/Macros.hh"
0014 #include "corecel/Types.hh"
0015 #include "corecel/math/Algorithms.hh"
0016 #include "celeritas/Constants.hh"
0017
0018 #include "NormalDistribution.hh"
0019
0020 namespace celeritas
0021 {
0022 //---------------------------------------------------------------------------//
0023 /*!
0024  * Sample from a gamma distribution.
0025  *
0026  * The gamma distribution can be parameterized with a shape parameter \f$
0027  * \alpha \f$ and a scale parameter \f$ \beta \f$ and has the PDF:
0028  * \f[
0029    f(x; \alpha, \beta) = \frac{x^{\alpha - 1} e^{-x / \beta}}{\beta^\alpha
0030  \Gamma(\alpha)} \quad \mathrm{for}\  x > 0, \quad \alpha, \beta > 0
0031    \f]
0032  * The algorithm described in Marsaglia and Tsang [MT00] is used here to
0033  * generate gamma-distributed random variables. The steps are:
0034  *  1. Set \f$ d = \alpha - 1/3 \f$ and \f$ c = 1 / \sqrt{9d} \f$
0035  *  2. Generate random variates \f$ Z \sim N(0,1) \f$ and \f$ U \sim U(0,1) \f$
0036  *  3. Set \f$ v = (1 + cZ)^3 \f$
0037  *  4. If \f$ \log U < Z^2 / 2 + d(1 - v + \log v) \f$ return \f$ dv \f$.
0038  *     Otherwise, go to step 2.
0039  * A squeeze function can be used to avoid the two logarithms in most cases by
0040  * accepting early if \f$ U < 1 - 0.0331 Z^4 \f$.
0041  *
0042  * Though this method is valid for \f$ \alpha \ge 1 \f$, it can easily be
0043  * extended for \f$ \alpha < 1 \f$: if \f$ X \sim \Gamma(\alpha + 1) \f$
0044  * and \f$ U \sim U(0,1) \f$, then \f$ X U^{1/\alpha} \sim \Gamma(\alpha) \f$.
0045  */
0046 template<class RealType = ::celeritas::real_type>
0047 class GammaDistribution
0048 {
0049   public:
0050     //!@{
0051     //! \name Type aliases
0052     using real_type = RealType;
0053     using result_type = real_type;
0054     //!@}
0055
0056   public:
0057     // Construct with shape and scale parameters
0058     explicit inline CELER_FUNCTION
0059     GammaDistribution(real_type alpha = 1, real_type beta = 1);
0060
0061     // Sample a random number according to the distribution
0062     template<class Generator>
0063     inline CELER_FUNCTION result_type operator()(Generator& rng);
0064
0065   private:
0066     real_type const alpha_;
0067     real_type const beta_;
0068     real_type const alpha_p_;
0069     real_type const d_;
0070     real_type const c_;
0071     NormalDistribution<real_type> sample_normal_;
0072 };
0073
0074 //---------------------------------------------------------------------------//
0075 // INLINE DEFINITIONS
0076 //---------------------------------------------------------------------------//
0077 /*!
0078  * Construct from the shape and scale parameters.
0079  */
0080 template<class RealType>
0081 CELER_FUNCTION
0082 GammaDistribution<RealType>::GammaDistribution(real_type alpha, real_type beta)
0083     : alpha_(alpha)
0084     , beta_(beta)
0085     , alpha_p_(alpha < 1 ? alpha + 1 : alpha)
0086     , d_(alpha_p_ - real_type(1) / 3)
0087     , c_(celeritas::rsqrt(9 * d_))
0088 {
0089     CELER_EXPECT(alpha_ > 0);
0090     CELER_EXPECT(beta_ > 0);
0091 }
0092
0093 //---------------------------------------------------------------------------//
0094 /*!
0095  * Sample a random number according to the distribution.
0096  */
0097 template<class RealType>
0098 template<class Generator>
0099 CELER_FUNCTION auto
0100 GammaDistribution<RealType>::operator()(Generator& rng) -> result_type
0101 {
0102     real_type u, v, z;
0103     do
0104     {
0105         do
0106         {
0107             z = sample_normal_(rng);
0108             v = 1 + c_ * z;
0109         } while (v <= 0);
0110         v = ipow<3>(v);
0111         u = generate_canonical<real_type>(rng);
0112     } while (u > 1 - real_type(0.0331) * ipow<4>(z)
0113              && std::log(u) > real_type(0.5) * ipow<2>(z)
0114                                   + d_ * (1 - v + std::log(v)));
0115
0116     result_type result = d_ * v * beta_;
0117     if (alpha_ != alpha_p_)
0118         result *= fastpow(generate_canonical<real_type>(rng), 1 / alpha_);
0119     return result;
0120 }
0121
0122 //---------------------------------------------------------------------------//
0123 }  // namespace celeritas