random/distribution/NormalDistribution.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/NormalDistribution.hh
0007 //---------------------------------------------------------------------------//
0008 #pragma once
0009
0010 #include <cmath>
0011 #include <type_traits>
0012
0013 #include "corecel/Assert.hh"
0014 #include "corecel/Macros.hh"
0015 #include "corecel/Types.hh"
0016 #include "corecel/math/Algorithms.hh"
0017 #include "celeritas/Constants.hh"
0018
0019 #include "GenerateCanonical.hh"
0020
0021 namespace celeritas
0022 {
0023 //---------------------------------------------------------------------------//
0024 /*!
0025  * Sample from a normal distribution.
0026  *
0027  * This uses the Box-Muller transform to generate pairs of independent,
0028  * normally distributed random numbers, returning them one at a time. Two
0029  * random numbers uniformly distributed on [0, 1] are mapped to two
0030  * independent, standard, normally distributed samples using the relations:
0031  * \f[
0032   x_1 = \sqrt{-2 \ln \xi_1} \cos(2 \pi \xi_2)
0033   x_2 = \sqrt{-2 \ln \xi_1} \sin(2 \pi \xi_2)
0034    \f]
0035  */
0036 template<class RealType = ::celeritas::real_type>
0037 class NormalDistribution
0038 {
0039     static_assert(std::is_floating_point_v<RealType>);
0040
0041   public:
0042     //!@{
0043     //! \name Type aliases
0044     using real_type = RealType;
0045     using result_type = real_type;
0046     //!@}
0047
0048   public:
0049     // Construct with mean and standard deviation
0050     inline CELER_FUNCTION NormalDistribution(real_type mean, real_type stddev);
0051
0052     //! Construct with unit deviation
0053     explicit CELER_FUNCTION NormalDistribution(real_type mean)
0054         : NormalDistribution{mean, 1}
0055     {
0056     }
0057
0058     //! Construct with unit deviation and zero mean
0059     CELER_FUNCTION NormalDistribution() : NormalDistribution{0, 1} {}
0060
0061     // Initialize with parameters but not spare values
0062     inline CELER_FUNCTION NormalDistribution(NormalDistribution const& other);
0063
0064     // Reset spare value of other distribution
0065     inline CELER_FUNCTION NormalDistribution(NormalDistribution&& other);
0066
0067     // Keep spare value but change distribution
0068     inline CELER_FUNCTION NormalDistribution&
0069     operator=(NormalDistribution const&);
0070
0071     // Possibly use spare value, change distribution
0072     inline CELER_FUNCTION NormalDistribution& operator=(NormalDistribution&&);
0073
0074     // Default destructor (rule of 5)
0075     ~NormalDistribution() = default;
0076
0077     // Sample a random number according to the distribution
0078     template<class Generator>
0079     inline CELER_FUNCTION result_type operator()(Generator& rng);
0080
0081   private:
0082     // Distribution properties
0083     real_type mean_;
0084     real_type stddev_;
0085
0086     // Intermediate samples
0087     real_type spare_{};
0088     bool has_spare_{false};
0089 };
0090
0091 //---------------------------------------------------------------------------//
0092 // INLINE DEFINITIONS
0093 //---------------------------------------------------------------------------//
0094 /*!
0095  * Construct with mean and standard deviation.
0096  */
0097 template<class RealType>
0098 CELER_FUNCTION
0099 NormalDistribution<RealType>::NormalDistribution(real_type mean,
0100                                                  real_type stddev)
0101     : mean_(mean), stddev_(stddev)
0102 {
0103     CELER_EXPECT(stddev > 0);
0104 }
0105
0106 //---------------------------------------------------------------------------//
0107 /*!
0108  * Initialize with parameters but not spare values.
0109  */
0110 template<class RealType>
0111 CELER_FUNCTION
0112 NormalDistribution<RealType>::NormalDistribution(NormalDistribution const& other)
0113     : mean_{other.mean}, stddev_{other.stddev}
0114 {
0115 }
0116
0117 //---------------------------------------------------------------------------//
0118 /*!
0119  * Reset spare value of other distribution.
0120  */
0121 template<class RealType>
0122 CELER_FUNCTION
0123 NormalDistribution<RealType>::NormalDistribution(NormalDistribution&& other)
0124     : mean_{other.mean_}
0125     , stddev_{other.stddev_}
0126     , spare_{other.spare_}
0127     , has_spare_{other.has_spare_}
0128 {
0129     other.has_spare_ = false;
0130 }
0131
0132 //---------------------------------------------------------------------------//
0133 /*!
0134  * Keep spare value but change distribution.
0135  */
0136 template<class RealType>
0137 CELER_FUNCTION NormalDistribution<RealType>&
0138 NormalDistribution<RealType>::operator=(NormalDistribution const& other)
0139 {
0140     mean_ = other.mean_;
0141     stddev_ = other.stddev_;
0142     return *this;
0143 }
0144
0145 //---------------------------------------------------------------------------//
0146 /*!
0147  * Possibly use spare value, change distribution.
0148  */
0149 template<class RealType>
0150 CELER_FUNCTION NormalDistribution<RealType>&
0151 NormalDistribution<RealType>::operator=(NormalDistribution&& other)
0152 {
0153     mean_ = other.mean_;
0154     stddev_ = other.stddev_;
0155     if (!has_spare_ && other.has_spare_)
0156     {
0157         spare_ = other.spare_;
0158         has_spare_ = other.has_spare_;
0159         other.has_spare_ = false;
0160     }
0161     return *this;
0162 }
0163
0164 //---------------------------------------------------------------------------//
0165 /*!
0166  * Sample a random number according to the distribution.
0167  */
0168 template<class RealType>
0169 template<class Generator>
0170 CELER_FUNCTION auto
0171 NormalDistribution<RealType>::operator()(Generator& rng) -> result_type
0172 {
0173     if (has_spare_)
0174     {
0175         has_spare_ = false;
0176         return std::fma(spare_, stddev_, mean_);
0177     }
0178
0179     constexpr auto twopi = static_cast<RealType>(2 * m_pi);
0180     real_type theta = twopi * generate_canonical<RealType>(rng);
0181     real_type r = std::sqrt(-2 * std::log(generate_canonical<RealType>(rng)));
0182     spare_ = r * std::cos(theta);
0183     has_spare_ = true;
0184     return std::fma(r * std::sin(theta), stddev_, mean_);
0185 }
0186
0187 //---------------------------------------------------------------------------//
0188 }  // namespace celeritas