celeritas/field/RungeKuttaIntegrator.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 celeritas/field/RungeKuttaIntegrator.hh
0006 //---------------------------------------------------------------------------//
0007 #pragma once
0008
0009 #include "corecel/Types.hh"
0010 #include "corecel/math/Algorithms.hh"
0011
0012 #include "Types.hh"
0013
0014 #include "detail/FieldUtils.hh"
0015
0016 namespace celeritas
0017 {
0018 //---------------------------------------------------------------------------//
0019 /*!
0020  * Integrate the field ODEs using the 4th order classical Runge-Kutta method.
0021  *
0022  * This method estimates the updated state from an initial state and evaluates
0023  * the truncation error, with fourth-order accuracy based on description in
0024  * Numerical Recipes in C, The Art of Scientific Computing, Sec. 16.2,
0025  * Adaptive Stepsize Control for Runge-Kutta.
0026  *
0027  * For a magnetic field equation \em f along a charged particle trajectory
0028  * with state \em y, which includes position and momentum but may also include
0029  * time and spin. For N-variables (\em i = 1, ... N), the right hand side of
0030  * the equation
0031  * \f[
0032  *  \difd{y_{i}}{s} = f_i (s, y_{i})
0033  * \f]
0034  * and the fouth order Runge-Kutta solution for a given step size, \em h is
0035  * \f[
0036  *  y_{n+1} - y_{n} = h f(x_n, y_n) = \frac{h}{6} (k_1 + 2 k_2 + 2 k_3 + k_4)
0037  * \f]
0038  * which is the average slope at four different points,
0039  * The truncation error is the difference of the final states of one full step
0040  * (\em y1) and two half steps (\em y2)
0041  * \f[
0042  *  \Delta = y_2 - y_1, y(x+h) = y_2 + \frac{\Delta}{15} + \mathrm{O}(h^{6})
0043  * \f]
0044  */
0045 template<class EquationT>
0046 class RungeKuttaIntegrator
0047 {
0048   public:
0049     //!@{
0050     //! \name Type aliases
0051     using result_type = FieldIntegration;
0052     //!@}
0053
0054   public:
0055     //! Construct with the equation of motion
0056     explicit CELER_FUNCTION RungeKuttaIntegrator(EquationT&& eq)
0057         : calc_rhs_(::celeritas::forward<EquationT>(eq))
0058     {
0059     }
0060
0061     // Advance the ODE state according to the field equations
0062     CELER_FUNCTION result_type operator()(real_type step,
0063                                           OdeState const& beg_state) const;
0064
0065   private:
0066     // Return the final state by the 4th order Runge-Kutta method
0067     CELER_FUNCTION auto do_step(real_type step,
0068                                 OdeState const& beg_state,
0069                                 OdeState const& end_slope) const -> OdeState;
0070
0071   private:
0072     // Equation of the motion
0073     EquationT calc_rhs_;
0074 };
0075
0076 //---------------------------------------------------------------------------//
0077 // DEDUCTION GUIDES
0078 //---------------------------------------------------------------------------//
0079 template<class EquationT>
0080 CELER_FUNCTION
0081 RungeKuttaIntegrator(EquationT&&) -> RungeKuttaIntegrator<EquationT>;
0082
0083 //---------------------------------------------------------------------------//
0084 // INLINE DEFINITIONS
0085 //---------------------------------------------------------------------------//
0086 /*!
0087  * Numerically integrate and return the updated state with estimated error.
0088  */
0089 template<class E>
0090 CELER_FUNCTION auto RungeKuttaIntegrator<E>::operator()(
0091     real_type step, OdeState const& beg_state) const -> result_type
0092 {
0093     using celeritas::axpy;
0094     real_type half_step = step / real_type(2);
0095     constexpr real_type fourth_order_correction = 1 / real_type(15);
0096
0097     result_type result;
0098     OdeState beg_slope = calc_rhs_(beg_state);
0099
0100     // Do two half steps
0101     result.mid_state = this->do_step(half_step, beg_state, beg_slope);
0102     result.end_state = this->do_step(
0103         half_step, result.mid_state, calc_rhs_(result.mid_state));
0104
0105     // Do a full step
0106     OdeState yt = this->do_step(step, beg_state, beg_slope);
0107
0108     // Integrator error: difference between the full step and two half steps
0109     result.err_state = result.end_state;
0110     axpy(real_type(-1), yt, &result.err_state);
0111
0112     // Output correction with the 4th order coefficient (1/15)
0113     axpy(fourth_order_correction, result.err_state, &result.end_state);
0114
0115     return result;
0116 }
0117
0118 //---------------------------------------------------------------------------//
0119 /*!
0120  * The classical RungeKuttaIntegrator integrator (the 4th order).
0121  */
0122 template<class E>
0123 CELER_FUNCTION auto
0124 RungeKuttaIntegrator<E>::do_step(real_type step,
0125                                  OdeState const& beg_state,
0126                                  OdeState const& beg_slope) const -> OdeState
0127 {
0128     using celeritas::axpy;
0129     real_type half_step = step / real_type(2);
0130     constexpr real_type sixth = 1 / real_type(6);
0131
0132     // 1st step k1 = (step/2)*beg_slope
0133     OdeState mid_est = beg_state;
0134     axpy(half_step, beg_slope, &mid_est);
0135     OdeState mid_est_slope = calc_rhs_(mid_est);
0136
0137     // 2nd step k2 = (step/2)*mid_est_slope
0138     OdeState mid_state = beg_state;
0139     axpy(half_step, mid_est_slope, &mid_state);
0140     OdeState mid_slope = calc_rhs_(mid_state);
0141
0142     // 3rd step k3 = step*mid_slope
0143     OdeState end_est = beg_state;
0144     axpy(step, mid_slope, &end_est);
0145     OdeState end_slope = calc_rhs_(end_est);
0146
0147     // Average slope at all 4 points
0148     axpy(real_type(1), beg_slope, &end_slope);
0149     axpy(real_type(2), mid_slope, &end_slope);
0150     axpy(real_type(2), mid_est_slope, &end_slope);
0151
0152     // 4th Step k4 = h*dydxt and the final RK4 output: k1/6+k4/6+(k2+k3)/3
0153     OdeState end_state = beg_state;
0154     axpy(sixth * step, end_slope, &end_state);
0155
0156     return end_state;
0157 }
0158
0159 //---------------------------------------------------------------------------//
0160 }  // namespace celeritas