RK4Integrator.hpp

/******************************************************************************
* Copyright (c) ATTX LLC 2024. All Rights Reserved.
*
* This software and associated documentation (the "Software") are the 
* proprietary and confidential information of ATTX, LLC. The Software is 
* furnished under a license agreement between ATTX and the user organization 
* and may be used or copied only in accordance with the terms of the agreement.
* Refer to 'license/attx_license.adoc' for standard license terms.
*
* EXPORT CONTROL NOTICE: THIS SOFTWARE MAY INCLUDE CONTENT CONTROLLED UNDER THE
* INTERNATIONAL TRAFFIC IN ARMS REGULATIONS (ITAR) OR THE EXPORT ADMINISTRATION 
* REGULATIONS (EAR99). No part of the Software may be used, reproduced, or 
* transmitted in any form or by any means, for any purpose, without the express 
* written permission of ATTX, LLC.
******************************************************************************/
/*
RK4 header file
---------------
 
 
Author: Alex Reynolds
*/
#ifndef RK4_HPP
#define RK4_HPP
 
#include <array>
 
#include "Rates.hpp"
#include "Integrator.hpp"
 
namespace clockwerk {
 
    template <typename T, unsigned int N>
    class RK4Integrator : public Integrator<T, N> {
    public:
        RK4Integrator(Rates<T, N>& rate_calculator);
 
        /// Function to take a full integrator step forward from time start to time end
        /// @param   start_time      Start time for integration
        /// @param   end_time        End time for integration
        /// @param   start_state     State at start of integration
        /// @param   out_state       Output state via implicit -- result of integration
        void step(T start_time, T end_time, const std::array<T, N> &start_state, std::array<T, N> &out_state);
 
        /// -------------------------------------------------------------------------
        /// Functions for manual RK4 control -- calling these individually allows
        /// a downstream user to manually control RK4 steps. Each is executed in order
        /// for a full step
        /// -------------------------------------------------------------------------
        /// @brief   Function to configure a full integration step
        /// @param   start_time      Start time for the full integration step
        /// @param   end_time        End time for the full integration step
        /// @param   start_state     State at start of integration
        void configureForStep(T start_time, T end_time, const std::array<T, N> &start_state);
        /// @brief   Function to calculate k1
        /// @param   state_for_k2    The integrated state to be used in the calculation of k2
        void calculateK1(std::array<T, N> &state_for_k2);
        /// @brief   Function to calculate k2
        /// @param   state_in_k2     The state input to calculate k2
        /// @param   state_for_k3    The integrated state to be used in the calculation of k3
        void calculateK2(const std::array<T, N> &state_in_k2, std::array<T, N> &state_for_k3);
        /// @brief   Function to calculate k3
        /// @param   state_in_k3     The state input to calculate k3
        /// @param   state_for_k4    The integrated state to be used in the calculation of k4
        void calculateK3(const std::array<T, N> &state_in_k2, std::array<T, N> &state_for_k4);
        /// @brief   Function to calculate k4
        /// @param   state_in_k4     The state input to calculate k4
        void calculateK4(const std::array<T, N> &state_in_k4);
        /// @brief   Calculation of final integrated step at the end of state
        /// @param   end_state   State at the end of step
        void getValueEndStep(std::array<T, N> &end_state);
    protected:
        /// Step size for integrator
        T _full_step_size;
        T _half_step_size;
        T _step_average;
 
        /// Initial state tracking
        T _start_time;
        T _end_time;
        std::array<T, N> _start_state;
 
        /// RK4 step rate values k1, k2, k3, k4 and var to calc average rate
        std::array<T, N> _k1;
        std::array<T, N> _k2;
        std::array<T, N> _k3;
        std::array<T, N> _k4;
    };
 
    template <typename T, unsigned int N>
    RK4Integrator<T, N>::RK4Integrator(Rates<T, N>& rate_calculator)
        : Integrator<T, N>(rate_calculator) {
        /// Initialize steps and whatnot
        _full_step_size = 0.0;
        _half_step_size = 0.0;
        _step_average = 0.0;
        _start_time = 0.0;
        _end_time = 0.0;
 
        /// Initialize k values
        for(unsigned int i = 0; i < N; i++) {
            _start_state[i] = 0.0;
            _k1[i] = 0.0;
            _k2[i] = 0.0;
            _k3[i] = 0.0;
            _k4[i] = 0.0;
        }
    }
 
    template <typename T, unsigned int N>
    void RK4Integrator<T, N>::configureForStep(T start_time, T end_time,
                                            const std::array<T, N> &start_state) {
        /// Set up for step
        _full_step_size = end_time - start_time;
        _half_step_size = 0.5*_full_step_size;
        _step_average = _full_step_size/6.0;
        _start_time = start_time;
        _end_time = end_time;
        for(unsigned int i = 0; i < N; i++) {
            _start_state[i] = start_state[i];
        }
    }
 
    template <typename T, unsigned int N>
    void RK4Integrator<T, N>::calculateK1(std::array<T, N> &state_for_k2) {
        /// Calculate k1 from start state
        Integrator<T, N>::_rate_calculator.calculateRates(_start_time, _start_state, _k1);
 
        /// Calculate state for k2 calculation downstream
        for(unsigned int i = 0; i < N; i++) {
            state_for_k2[i] = _start_state[i] + _half_step_size*_k1[i];
        }
    }
 
    template <typename T, unsigned int N>
    void RK4Integrator<T, N>::calculateK2(const std::array<T, N> &state_in_k2,
                                        std::array<T, N> &state_for_k3) {
        /// Calculate k2 from start state
        Integrator<T, N>::_rate_calculator.calculateRates(_start_time + _half_step_size, state_in_k2, _k2);
 
        /// Calculate state for k3 calculation downstream
        for(unsigned int i = 0; i < N; i++) {
            state_for_k3[i] = _start_state[i] + _half_step_size*_k2[i];
        }
    }
 
    template <typename T, unsigned int N>
    void RK4Integrator<T, N>::calculateK3(const std::array<T, N> &state_in_k3,
                                        std::array<T, N> &state_for_k4) {
        /// Calculate k3 from start state
        Integrator<T, N>::_rate_calculator.calculateRates(_start_time + _half_step_size, state_in_k3, _k3);
 
        /// Calculate state for k4 calculation downstream
        for(unsigned int i = 0; i < N; i++) {
            state_for_k4[i] = _start_state[i] + _full_step_size*_k3[i];
        }
    }
 
    template <typename T, unsigned int N>
    void RK4Integrator<T, N>::calculateK4(const std::array<T, N> &state_in_k4) {
        /// Calculate k4 from start state
        Integrator<T, N>::_rate_calculator.calculateRates(_start_time + _full_step_size, state_in_k4, _k4);
    }
 
    template <typename T, unsigned int N>
    void RK4Integrator<T, N>::getValueEndStep(std::array<T, N> &end_state) {
        for(unsigned int i = 0; i < N; i++) {
            end_state[i] = _start_state[i] + _step_average*(_k1[i] + 2.0*_k2[i] + 2.0*_k3[i] + _k4[i]);
        }
    }
 
    template <typename T, unsigned int N>
    void RK4Integrator<T, N>::step(T start_time, T end_time,
                                const std::array<T, N> &start_state,
                                std::array<T, N> &out_state) {
        /// This is an RK4 integrator, so we need to execute each of our K steps
        /// to execute our full integration. Go step by step, then sum total for
        /// our out state. Call subfunctions for simplicity
        configureForStep(start_time, end_time, start_state);
        calculateK1(out_state);
        calculateK2(out_state, out_state);
        calculateK3(out_state, out_state);
        calculateK4(out_state);
        getValueEndStep(out_state);
    }
 
}
 
#endif