rk2.hpp

// SPDX-License-Identifier: MIT
#pragma once
#include <array>
#include <type_traits>
 
#include "ddc_alias_inline_functions.hpp"
#include "ddc_aliases.hpp"
#include "ddc_helper.hpp"
#include "itimestepper.hpp"
#include "vector_field_common.hpp"
 
template <
        class FieldMem,
        class DerivFieldMem = FieldMem,
        class ExecSpace = Kokkos::DefaultExecutionSpace>
class RK2 : public ITimeStepper<FieldMem, DerivFieldMem, ExecSpace>
{
    using base_type = ITimeStepper<FieldMem, DerivFieldMem, ExecSpace>;
 
public:
    using typename base_type::IdxRange;
 
    using typename base_type::ValConstField;
    using typename base_type::ValField;
 
    using typename base_type::DerivConstField;
    using typename base_type::DerivField;
 
private:
    IdxRange const m_idx_range;
 
public:
    using base_type::update;
 
public:
    explicit RK2(IdxRange idx_range) : m_idx_range(idx_range) {}
 
    void update(
            ExecSpace const& exec_space,
            ValField y,
            double dt,
            std::function<void(DerivField, ValConstField)> dy_calculator,
            std::function<void(ValField, DerivConstField, double)> y_update) const final
    {
        DerivFieldMem k1_alloc(m_idx_range);
        DerivFieldMem k2_alloc(m_idx_range);
        FieldMem y_prime_alloc(m_idx_range);
 
        DerivField k1(k1_alloc);
        DerivField k2(k2_alloc);
        ValField y_prime(y_prime_alloc);
 
        // Save initial conditions
        base_type::copy(y_prime, get_const_field(y));
 
        // --------- Calculate k1 ------------
        // Calculate k1 = f(y)
        dy_calculator(k1, get_const_field(y));
 
        // --------- Calculate k2 ------------
        // Calculate y_new := y_n + h/2*k_1
        y_update(y_prime, get_const_field(k1), 0.5 * dt);
 
        // Calculate k2 = f(y_new)
        dy_calculator(k2, get_const_field(y_prime));
 
        // ----------- Update y --------------
        // Calculate y_{n+1} := y_n + h*k_2
        y_update(y, get_const_field(k2), dt);
    }
};