time_integrator.h

Go to the documentation of this file.

 
#ifndef TIME_INTEGRATOR_H
#define TIME_INTEGRATOR_H
 
#include <functional>
#include <string>
#include <unordered_map>
#include <vector>
 
#include <Kokkos_Core.hpp>
 
#include "common.h"
#include "state.h"
 
enum class TimeIntegratorType {
    FE,
    RK4,
    SSPRK3,
};
 
static const std::unordered_map<std::string, TimeIntegratorType> TIME_INTEGRATOR_TYPES = {
    {"FE", TimeIntegratorType::FE},
    {"RK4", TimeIntegratorType::RK4},
    {"SSPRK3", TimeIntegratorType::SSPRK3},
};
 
static const std::unordered_map<TimeIntegratorType, std::string> TIME_INTEGRATOR_NAMES = {
    {TimeIntegratorType::FE, "FE"},
    {TimeIntegratorType::RK4, "RK4"},
    {TimeIntegratorType::SSPRK3, "SSPRK3"},
};
 
using RHSFunction = std::function<void(State solution, State rhs, rtype t)>;
 
void axpby(const rtype a, const State & x, const rtype b, const State & y, const rtype * a_device = nullptr);
 
enum DtScale { DT = 0, DT_HALF = 1, DT_SIXTH = 2, DT_THIRD = 3, N_DT_SCALES = 4 };
 
KOKKOS_INLINE_FUNCTION
void dt_scales(const rtype dt, rtype * scales) {
    scales[DT] = dt;
    scales[DT_HALF] = 0.5_r * dt;
    scales[DT_SIXTH] = dt / 6.0_r;
    scales[DT_THIRD] = dt / 3.0_r;
}
 
void lerp(const rtype a, const State & x, const State & y);
 
class TimeIntegrator {
    public:
        virtual ~TimeIntegrator() = default;
 
        TimeIntegratorType get_type() const { return type; }
        uint8_t get_n_solution_vectors() const { return n_solution_vectors; }
        uint8_t get_n_rhs_vectors() const { return n_rhs_vectors; }
 
        virtual void take_step(const rtype t, const rtype dt,
                               std::vector<State> & solution_vec,
                               std::vector<State> & rhs_vec,
                               const RHSFunction & calc_rhs, const rtype * dt_device = nullptr) = 0;
 
    protected:
        TimeIntegratorType type;
        uint8_t n_solution_vectors;
        uint8_t n_rhs_vectors;
};
 
class FE : public TimeIntegrator {
    public:
        FE();
        void take_step(const rtype t, const rtype dt,
                       std::vector<State> & solution_vec,
                       std::vector<State> & rhs_vec,
                       const RHSFunction & calc_rhs, const rtype * dt_device = nullptr) override;
};
 
class RK4 : public TimeIntegrator {
    public:
        RK4();
        void take_step(const rtype t, const rtype dt,
                       std::vector<State> & solution_vec,
                       std::vector<State> & rhs_vec,
                       const RHSFunction & calc_rhs, const rtype * dt_device = nullptr) override;
};
 
class SSPRK3 : public TimeIntegrator {
    public:
        SSPRK3();
        void take_step(const rtype t, const rtype dt,
                       std::vector<State> & solution_vec,
                       std::vector<State> & rhs_vec,
                       const RHSFunction & calc_rhs, const rtype * dt_device = nullptr) override;
};
 
#endif // TIME_INTEGRATOR_H