random/distribution/DistributionVisitor.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/DistributionVisitor.hh
0006 //---------------------------------------------------------------------------//
0007 #pragma once
0008
0009 #include "corecel/Types.hh"
0010 #include "corecel/random/data/DistributionData.hh"
0011
0012 #include "DistributionTypeTraits.hh"
0013
0014 namespace celeritas
0015 {
0016 //---------------------------------------------------------------------------//
0017 /*!
0018  * Expand a macro to a switch statement over all possible distribution types.
0019  */
0020 class DistributionVisitor
0021 {
0022   public:
0023     // Construct with a reference to distribution data
0024     explicit inline CELER_FUNCTION DistributionVisitor(
0025         NativeCRef<DistributionParamsData> const&);
0026
0027     // Apply a functor to a 1D dietribution
0028     template<class F>
0029     CELER_CONSTEXPR_FUNCTION decltype(auto) operator()(F&& func,
0030                                                        OnedDistributionId id);
0031
0032     // Apply a functor to a 3D dietribution
0033     template<class F>
0034     CELER_CONSTEXPR_FUNCTION decltype(auto) operator()(
0035         F&& func, ThreedDistributionId id);
0036
0037   private:
0038     //// TYPES ////
0039
0040     using ODT = OnedDistributionType;
0041     using TDT = ThreedDistributionType;
0042
0043     //// DATA ////
0044
0045     NativeCRef<DistributionParamsData> const& params_;
0046 };
0047
0048 //---------------------------------------------------------------------------//
0049 // INLINE DEFINITIONS
0050 //---------------------------------------------------------------------------//
0051 /*!
0052  * Construct with a reference to distribution data.
0053  */
0054 CELER_FUNCTION DistributionVisitor::DistributionVisitor(
0055     NativeCRef<DistributionParamsData> const& params)
0056     : params_(params)
0057 {
0058     CELER_EXPECT(params_);
0059 }
0060
0061 //---------------------------------------------------------------------------//
0062 /*!
0063  * Expand a macro to a switch statement over all 1D distribution types.
0064  */
0065 template<class F>
0066 CELER_CONSTEXPR_FUNCTION decltype(auto) DistributionVisitor::operator()(
0067     F&& func, OnedDistributionId id)
0068 {
0069     CELER_EXPECT(id < params_.oned_types.size());
0070
0071     ODT type = params_.oned_types[id];
0072     size_type idx = params_.oned_indices[id];
0073
0074 #define CELER_DISTRIB_CASE(ENUM, FIELD) \
0075     case ODT::ENUM: { \
0076         using TypeT = typename OnedDistributionTypeTraits<ODT::ENUM>::type; \
0077         return celeritas::forward<F>(func)( \
0078             TypeT{params_.FIELD[ItemId<TypeT::RecordT>(idx)]}); \
0079     }
0080     switch (type)
0081     {
0082         CELER_DISTRIB_CASE(delta, delta_real);
0083         CELER_DISTRIB_CASE(normal, normal);
0084         CELER_DISTRIB_CASE(truncated_normal, truncated_normal);
0085         default:
0086             CELER_ASSERT_UNREACHABLE();
0087     }
0088 #undef CELER_DISTRIB_CASE
0089 }
0090
0091 //---------------------------------------------------------------------------//
0092 /*!
0093  * Expand a macro to a switch statement over all 3D distribution types.
0094  */
0095 template<class F>
0096 CELER_CONSTEXPR_FUNCTION decltype(auto) DistributionVisitor::operator()(
0097     F&& func, ThreedDistributionId id)
0098 {
0099     CELER_EXPECT(id < params_.threed_types.size());
0100
0101     TDT type = params_.threed_types[id];
0102     size_type idx = params_.threed_indices[id];
0103
0104 #define CELER_DISTRIB_CASE(ENUM, FIELD) \
0105     case TDT::ENUM: { \
0106         using TypeT = typename ThreedDistributionTypeTraits<TDT::ENUM>::type; \
0107         return celeritas::forward<F>(func)( \
0108             TypeT{params_.FIELD[ItemId<TypeT::RecordT>(idx)]}); \
0109     }
0110     switch (type)
0111     {
0112         CELER_DISTRIB_CASE(delta, delta_real3);
0113         CELER_DISTRIB_CASE(isotropic, isotropic);
0114         CELER_DISTRIB_CASE(uniform_box, uniform_box);
0115         default:
0116             CELER_ASSERT_UNREACHABLE();
0117     }
0118 #undef CELER_DISTRIB_CASE
0119 }
0120
0121 //---------------------------------------------------------------------------//
0122 // FREE FUNCTIONS
0123 //---------------------------------------------------------------------------//
0124 /*!
0125  * Helper function for sampling from a distribution.
0126  */
0127 template<class Visitor, class Dist, class Engine>
0128 CELER_FUNCTION decltype(auto) sample_with(
0129     Visitor&& visit, Dist&& dist, Engine& rng)
0130 {
0131     return celeritas::forward<Visitor>(visit)(
0132         [&rng](auto&& d) -> decltype(auto) { return d(rng); },
0133         celeritas::forward<Dist>(dist));
0134 }
0135
0136 //---------------------------------------------------------------------------//
0137 }  // namespace celeritas