random/distribution/NormalDistribution.hh

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