celeritas/field/RungeKuttaStepper.hh

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