celeritas/grid/PolyEvaluator.hh

0001 //----------------------------------*-C++-*----------------------------------//
0002 // Copyright 2022-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/grid/PolyEvaluator.hh
0007 //---------------------------------------------------------------------------//
0008 #pragma once
0009
0010 #include <cmath>
0011 #include <type_traits>
0012
0013 #include "corecel/Macros.hh"
0014 #include "corecel/Types.hh"
0015 #include "corecel/cont/Array.hh"
0016 #include "corecel/math/Algorithms.hh"
0017
0018 namespace celeritas
0019 {
0020 //---------------------------------------------------------------------------//
0021 /*!
0022  * Functor class to evaluate a polynomial.
0023  *
0024  * This is an efficient and foolproof way of storing and evaluating a
0025  * polynomial expansion:
0026  * \f[
0027   f(x) = a_0 + x * (a_1 + x * (a_2 + ...))
0028   \f]
0029  *
0030  * It replaces opaque expressions such as:
0031  * \code
0032     corr = (zeff * (real_type(1.84035e-4) * zeff - real_type(1.86427e-2))
0033              + real_type(1.41125));
0034  * \endcode
0035  * with
0036  * \code
0037     corr = PolyEvaluator{1.41125, -1.86427e-2, 1.84035e-4}(zeff);
0038    \endcode
0039  * or, to use an explicit type without having to cast each coefficient:
0040  * \code
0041    using PolyQuad = PolyEvaluator<real_type, N>;
0042    corr = PolyQuad{1.41125, -1.86427e-2, 1.84035e-4)(zeff);
0043  * \endcode
0044  */
0045 template<class T, size_type N>
0046 class PolyEvaluator
0047 {
0048   public:
0049     //!@{
0050     //! \name Type aliases
0051     using result_type = T;
0052     using argument_type = T;
0053     using ArrayT = Array<T, N + 1>;
0054     //!@}
0055
0056   public:
0057     //! Construct with the polynomial to evaluate
0058     template<class... Ts>
0059     explicit CELER_CONSTEXPR_FUNCTION PolyEvaluator(Ts... coeffs)
0060         : coeffs_{static_cast<T>(coeffs)...}
0061     {
0062         // Protect against leaving off a coefficient, e.g. PolyQuad(1, 2)
0063         static_assert(sizeof...(coeffs) == N + 1,
0064                       "All coefficients for PolyEvaluator must be explicitly "
0065                       "specified");
0066     }
0067
0068     //! Construct from an array of data
0069     CELER_CONSTEXPR_FUNCTION PolyEvaluator(ArrayT const& coeffs)
0070         : coeffs_{coeffs}
0071     {
0072     }
0073
0074     //! Evaluate the polynomial at the given value
0075     CELER_CONSTEXPR_FUNCTION T operator()(T arg) const
0076     {
0077         return this->calc_impl<0>(arg);
0078     }
0079
0080   private:
0081     ArrayT const coeffs_;
0082
0083     template<unsigned int M, std::enable_if_t<(M < N), bool> = true>
0084     CELER_CONSTEXPR_FUNCTION T calc_impl(T arg) const
0085     {
0086         return fma(arg, calc_impl<M + 1>(arg), coeffs_[M]);
0087     }
0088
0089     template<unsigned int M, std::enable_if_t<(M == N), bool> = true>
0090     CELER_CONSTEXPR_FUNCTION T calc_impl(T) const
0091     {
0092         return coeffs_[N];
0093     }
0094 };
0095
0096 //---------------------------------------------------------------------------//
0097 // DEDUCTION GUIDES
0098 //---------------------------------------------------------------------------//
0099 template<typename T, size_type N>
0100 CELER_FUNCTION PolyEvaluator(Array<T, N> const&) -> PolyEvaluator<T, N - 1>;
0101
0102 template<typename... Ts,
0103          std::enable_if_t<std::is_arithmetic_v<std::common_type_t<Ts...>>, bool>
0104          = true>
0105 CELER_FUNCTION PolyEvaluator(Ts&&...)
0106     -> PolyEvaluator<typename std::common_type_t<Ts...>, sizeof...(Ts) - 1>;
0107
0108 //---------------------------------------------------------------------------//
0109 }  // namespace celeritas