random/distribution/RejectionSampler.hh

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