random/distribution/TruncatedDistribution.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/TruncatedDistribution.hh
0006 //---------------------------------------------------------------------------//
0007 #pragma once
0008
0009 #include "corecel/Config.hh"
0010
0011 #include "corecel/Assert.hh"
0012 #include "corecel/Macros.hh"
0013 #include "corecel/random/data/DistributionData.hh"
0014
0015 namespace celeritas
0016 {
0017 //---------------------------------------------------------------------------//
0018 /*!
0019  * Sample from a truncated distribution.
0020  *
0021  * Sample from an arbitrary distribution truncated to a finite interval. Values
0022  * are drawn from the underlying distribution using rejection sampling until
0023  * they fall within the truncation bounds.
0024  *
0025  * \warning Because of the rejection sampling, it is possible to create a
0026  * distribution that never (or almost never) accepts a value from the
0027  * underlying distribution. A loop checker is active in debug mode that
0028  * prevents more than 32 failed samples. Since GPU performance is so sensitive
0029  * to rejection failures, consider transforming the underlying distribution if
0030  * this limit is hit. When using this distribution for physics distributions,
0031  * it is \em strongly recommended to use \c test::HistogramSampler to verify
0032  * the number of samples being taken.
0033  */
0034 template<class Distribution>
0035 class TruncatedDistribution
0036 {
0037   public:
0038     //!@{
0039     //! \name Type aliases
0040     using result_type = typename Distribution::result_type;
0041     using RecordT = TruncatedDistributionRecord<typename Distribution::RecordT>;
0042     //!@}
0043
0044   public:
0045     // Construct with distribution and truncation bounds
0046     template<class T, class U, class... Args>
0047     inline CELER_FUNCTION TruncatedDistribution(
0048         T lower, U upper, Args&&... args);
0049
0050     //! Construct from a device-friendly variant record
0051     explicit CELER_FUNCTION TruncatedDistribution(RecordT const& record)
0052         : TruncatedDistribution(
0053               record.lower, record.upper, Distribution(record.distribution))
0054     {
0055     }
0056
0057     // Sample a random number according to the truncated distribution
0058     template<class Generator>
0059     inline CELER_FUNCTION result_type operator()(Generator& rng);
0060
0061     //! Assert fewer than this number of samples is tried when CELERITAS_DEBUG
0062     static constexpr inline int max_debug_samples = 32;
0063
0064   private:
0065     Distribution sample_;
0066     result_type lower_;
0067     result_type upper_;
0068 };
0069
0070 //---------------------------------------------------------------------------//
0071 // INLINE DEFINITIONS
0072 //---------------------------------------------------------------------------//
0073 /*!
0074  * Construct with distribution and truncation bounds.
0075  */
0076 template<class Distribution>
0077 template<class T, class U, class... Args>
0078 CELER_FUNCTION TruncatedDistribution<Distribution>::TruncatedDistribution(
0079     T lower, U upper, Args&&... args)
0080     : sample_(celeritas::forward<Args>(args)...), lower_(lower), upper_(upper)
0081 {
0082     CELER_EXPECT(lower < upper);
0083 }
0084
0085 //---------------------------------------------------------------------------//
0086 /*!
0087  * Sample a random number according to the truncated distribution.
0088  */
0089 template<class Distribution>
0090 template<class Generator>
0091 CELER_FUNCTION auto TruncatedDistribution<Distribution>::operator()(
0092     Generator& rng) -> result_type
0093 {
0094     int num_remaining_samples = max_debug_samples + 1;
0095     result_type result;
0096     do
0097     {
0098         if constexpr (CELERITAS_DEBUG)
0099         {
0100             // Prevent infinite loops (debug assertions only)
0101             if (CELER_UNLIKELY(--num_remaining_samples == 0))
0102             {
0103                 CELER_DEBUG_FAIL(
0104                     "too many samples taken in TruncatedDistribution",
0105                     internal);
0106             }
0107         }
0108         else
0109         {
0110             // CUDA 12.6 causes warning about unused variable
0111             CELER_DISCARD(num_remaining_samples);
0112         }
0113
0114         result = sample_(rng);
0115     } while (result < lower_ || result > upper_);
0116     return result;
0117 }
0118
0119 //---------------------------------------------------------------------------//
0120 }  // namespace celeritas