random/distribution/RejectionSampler.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/RejectionSampler.hh
0006 //---------------------------------------------------------------------------//
0007 #pragma once
0008
0009 #include <type_traits>
0010
0011 #include "corecel/Assert.hh"
0012 #include "corecel/Macros.hh"
0013 #include "corecel/Types.hh"
0014
0015 #include "GenerateCanonical.hh"
0016
0017 namespace celeritas
0018 {
0019 //---------------------------------------------------------------------------//
0020 /*!
0021  * Return whether a rejection loop needs to continue trying.
0022  *
0023  * A common implementation of sampling from a "difficult" (non-analytically
0024  * invertible) probability distribution function is to bound the difficult
0025  * distribution \em f(x) with another easily sampled function \em g(x) . Given
0026  a
0027  * maximum value \em M over the \em x interval being sampled, it is equivalent
0028  * to sampling \em f(x) by instead sampling from \em g(x) and rejecting with
0029  * probability \f[
0030    \frac{f(x)}{M g(x)}
0031  * \f]
0032  *
0033  * These invocations generate statistically equivalent results:
0034  *  - `BernoulliDistribution(1 - p / pmax)(rng);`
0035  *  - `!BernoulliDistribution(p / pmax)(rng);`
0036  *  - `p < pmax * UniformDistribution{}(rng)`
0037  *  - `RejectionSampler(p, pmax)(rng);`
0038  *
0039  * This is meant for rejection sampling, e.g., on cross section:
0040  * \code
0041     do {
0042       xs = sample_xs(rng);
0043     } while (RejectionSampler{xs, xs_max}(rng));
0044    \endcode
0045  */
0046 template<class RealType = ::celeritas::real_type>
0047 class RejectionSampler
0048 {
0049     static_assert(std::is_floating_point_v<RealType>);
0050
0051   public:
0052     //!@{
0053     //! \name Type aliases
0054     using real_type = RealType;
0055     using result_type = real_type;
0056     //!@}
0057
0058   public:
0059     // Construct with acceptance probability
0060     inline CELER_FUNCTION RejectionSampler(real_type f, real_type fmax);
0061
0062     //! Construct when the distribution's maximum is normalized
0063     explicit CELER_FUNCTION RejectionSampler(real_type f)
0064         : RejectionSampler{f, 1}
0065     {
0066     }
0067
0068     // Sample a random number according to the distribution
0069     template<class Generator>
0070     inline CELER_FUNCTION result_type operator()(Generator& rng);
0071
0072   private:
0073     RealType f_;
0074     RealType fmax_;
0075 };
0076
0077 //---------------------------------------------------------------------------//
0078 // INLINE DEFINITIONS
0079 //---------------------------------------------------------------------------//
0080 /*!
0081  * Construct with acceptance probability and maximum probability.
0082  */
0083 template<class RealType>
0084 CELER_FUNCTION
0085 RejectionSampler<RealType>::RejectionSampler(real_type f, real_type fmax)
0086     : f_{f}, fmax_{fmax}
0087 {
0088     CELER_EXPECT(f_ >= 0);
0089     CELER_EXPECT(fmax_ >= f_);
0090 }
0091
0092 //---------------------------------------------------------------------------//
0093 /*!
0094  * Sample a random number according to the distribution.
0095  */
0096 template<class RealType>
0097 template<class Generator>
0098 CELER_FUNCTION auto
0099 RejectionSampler<RealType>::operator()(Generator& rng) -> result_type
0100 {
0101     return f_ < fmax_ * generate_canonical<RealType>(rng);
0102 }
0103
0104 //---------------------------------------------------------------------------//
0105 }  // namespace celeritas