main.cpp

#include <AMReX.H>
#include <AMReX_BoxArray.H>
#include <AMReX_DistributionMapping.H>
#include <AMReX_Geometry.H>
#include <AMReX_MultiFab.H>
#include <AMReX_ParallelDescriptor.H>
#include <AMReX_ParmParse.H>
#include <AMReX_PlotFileUtil.H>
#include <AMReX_Print.H>
#include <AMReX_Sundials.H>

#include <arkode/arkode.h>
#include <arkode/arkode_erkstep.h>
#include <sundials/sundials_types.h>

#include <algorithm>
#include <array>
#include <cmath>
#include <cstddef>
#include <sstream>
#include <string>
#include <type_traits>
#include <utility>

namespace
{

static_assert(std::is_same_v<amrex::Real, sunrealtype>,
              "AMReX and SUNDIALS must use the same floating-point precision.");
static_assert(AMREX_SPACEDIM == 2, "This repository is strictly 2D.");

using amrex::Array4;
using amrex::Box;
using amrex::BoxArray;
using amrex::DistributionMapping;
using amrex::Geometry;
using amrex::IntVect;
using amrex::MFIter;
using amrex::MultiFab;
using amrex::ParmParse;
using amrex::Real;
using amrex::RealBox;
using amrex::Vector;

struct ProblemParameters
{
    std::array<int, AMREX_SPACEDIM> n_cell{128, 128};
    std::array<Real, AMREX_SPACEDIM> prob_lo{0.0, 0.0};
    std::array<Real, AMREX_SPACEDIM> prob_hi{1.0, 1.0};
    std::array<Real, AMREX_SPACEDIM> velocity{1.0, 0.35};
    std::array<Real, AMREX_SPACEDIM> pulse_center{0.2, 0.5};
    std::array<Real, AMREX_SPACEDIM> pulse_width{0.08, 0.08};
    int max_grid_size = 32;
    Real diffusion = 2.5e-4;
    Real background = 1.0e-8;
    Real pulse_amplitude = 1.0;
    Real final_time = 0.5;
    Real output_interval = 0.1;
    Real rtol = 1.0e-6;
    Real atol = 1.0e-10;
    Real init_step = 1.0e-4;
    Real max_step = 5.0e-3;
    std::string plot_prefix = "plt";
};

void check_flag(const int flag, const std::string& where)
{
    if (flag < 0) {
        std::ostringstream oss;
        oss << where << " failed with SUNDIALS flag " << flag;
        amrex::Abort(oss.str());
    }
}

void check_ptr(const void* ptr, const std::string& where)
{
    if (ptr == nullptr) {
        amrex::Abort(where + " returned a null pointer");
    }
}

MultiFab& get_multifab(N_Vector y, const std::string& where)
{
    auto* mf = amrex::sundials::N_VGetVectorPointer_MultiFab(y);
    check_ptr(mf, where);
    return *mf;
}

ProblemParameters read_parameters()
{
    ProblemParameters p;

    ParmParse pp("transport");
    Vector<int> n_cell(AMREX_SPACEDIM);
    if (pp.queryarr("n_cell", n_cell)) {
        for (int d = 0; d < AMREX_SPACEDIM; ++d) {
            p.n_cell[d] = n_cell[d];
        }
    }

    pp.query("max_grid_size", p.max_grid_size);

    Vector<Real> real_vals(AMREX_SPACEDIM);
    if (pp.queryarr("prob_lo", real_vals)) {
        for (int d = 0; d < AMREX_SPACEDIM; ++d) {
            p.prob_lo[d] = real_vals[d];
        }
    }

    if (pp.queryarr("prob_hi", real_vals)) {
        for (int d = 0; d < AMREX_SPACEDIM; ++d) {
            p.prob_hi[d] = real_vals[d];
        }
    }

    if (pp.queryarr("velocity", real_vals)) {
        for (int d = 0; d < AMREX_SPACEDIM; ++d) {
            p.velocity[d] = real_vals[d];
        }
    }

    if (pp.queryarr("pulse_center", real_vals)) {
        for (int d = 0; d < AMREX_SPACEDIM; ++d) {
            p.pulse_center[d] = real_vals[d];
        }
    }

    if (pp.queryarr("pulse_width", real_vals)) {
        for (int d = 0; d < AMREX_SPACEDIM; ++d) {
            p.pulse_width[d] = real_vals[d];
        }
    }

    pp.query("diffusion", p.diffusion);
    pp.query("background", p.background);
    pp.query("pulse_amplitude", p.pulse_amplitude);
    pp.query("final_time", p.final_time);
    pp.query("output_interval", p.output_interval);
    pp.query("rtol", p.rtol);
    pp.query("atol", p.atol);
    pp.query("init_step", p.init_step);
    pp.query("max_step", p.max_step);

    ParmParse pp_plot("plot");
    pp_plot.query("prefix", p.plot_prefix);

    return p;
}

class TransportSystem
{
  public:
    explicit TransportSystem(ProblemParameters params)
        : m_params(std::move(params)),
          m_geom(make_geometry(m_params)),
          m_ba(make_box_array(m_geom.Domain(), m_params.max_grid_size)),
          m_dm(m_ba),
          m_state(m_ba, m_dm, 1, 1),
          m_global_cells(static_cast<long>(m_ba.numPts()))
    {}

    void initialize()
    {
        const auto dx = m_geom.CellSizeArray();
        const auto plo = m_geom.ProbLoArray();

        for (MFIter mfi(m_state); mfi.isValid(); ++mfi) {
            const Box& bx = mfi.validbox();
            auto arr = m_state.array(mfi);
            const auto lo = amrex::lbound(bx);
            const auto hi = amrex::ubound(bx);

            for (int j = lo.y; j <= hi.y; ++j) {
                for (int i = lo.x; i <= hi.x; ++i) {
                    const Real x = plo[0] + (static_cast<Real>(i) + Real(0.5)) * dx[0];
                    const Real y = plo[1] + (static_cast<Real>(j) + Real(0.5)) * dx[1];
                    const Real rx = (x - m_params.pulse_center[0]) / m_params.pulse_width[0];
                    const Real ry = (y - m_params.pulse_center[1]) / m_params.pulse_width[1];
                    arr(i, j, 0) = m_params.background
                                 + m_params.pulse_amplitude
                                       * std::exp(-Real(0.5) * (rx * rx + ry * ry));
                }
            }
        }
    }

    [[nodiscard]] long global_cells() const noexcept { return m_global_cells; }
    [[nodiscard]] const ProblemParameters& params() const noexcept { return m_params; }
    [[nodiscard]] sunindextype vector_length() const noexcept
    {
        return static_cast<sunindextype>(m_state.nComp()) * static_cast<sunindextype>(m_global_cells);
    }
    [[nodiscard]] MultiFab& state() noexcept { return m_state; }

    void write_plotfile(N_Vector y, const Real time, const int index)
    {
        const MultiFab& state = get_multifab(y, "N_VGetVectorPointer_MultiFab");

        const std::string plot_name = amrex::Concatenate(m_params.plot_prefix, index, 5);
        amrex::Print() << "Writing plotfile " << plot_name << " at t = " << time << "\n";
        amrex::WriteSingleLevelPlotfile(
            plot_name, state, Vector<std::string>{"concentration"}, m_geom, time, 0);
    }

    int rhs(Real, N_Vector y, N_Vector ydot)
    {
        MultiFab& state = get_multifab(y, "N_VGetVectorPointer_MultiFab");
        MultiFab& rhs = get_multifab(ydot, "N_VGetVectorPointer_MultiFab");
        state.FillBoundary(m_geom.periodicity());

        const auto dx = m_geom.CellSizeArray();
        const Real inv_dx = Real(1.0) / dx[0];
        const Real inv_dy = Real(1.0) / dx[1];
        const Real inv_dx2 = inv_dx * inv_dx;
        const Real inv_dy2 = inv_dy * inv_dy;
        const Real vel_x = m_params.velocity[0];
        const Real vel_y = m_params.velocity[1];
        const Real diffusivity = m_params.diffusion;

        for (MFIter mfi(state); mfi.isValid(); ++mfi) {
            const Box& bx = mfi.validbox();
            const auto state_arr = state.const_array(mfi);
            auto rhs_arr = rhs.array(mfi);
            const auto lo = amrex::lbound(bx);
            const auto hi = amrex::ubound(bx);

            for (int j = lo.y; j <= hi.y; ++j) {
                for (int i = lo.x; i <= hi.x; ++i) {
                    const Real c = state_arr(i, j, 0);
                    const Real c_im1 = state_arr(i - 1, j, 0);
                    const Real c_ip1 = state_arr(i + 1, j, 0);
                    const Real c_jm1 = state_arr(i, j - 1, 0);
                    const Real c_jp1 = state_arr(i, j + 1, 0);

                    const Real flux_x_hi = (vel_x >= Real(0.0)) ? vel_x * c : vel_x * c_ip1;
                    const Real flux_x_lo = (vel_x >= Real(0.0)) ? vel_x * c_im1 : vel_x * c;
                    const Real flux_y_hi = (vel_y >= Real(0.0)) ? vel_y * c : vel_y * c_jp1;
                    const Real flux_y_lo = (vel_y >= Real(0.0)) ? vel_y * c_jm1 : vel_y * c;

                    const Real advection = -(flux_x_hi - flux_x_lo) * inv_dx
                                         - (flux_y_hi - flux_y_lo) * inv_dy;
                    const Real diffusion = diffusivity
                                         * ((c_ip1 - Real(2.0) * c + c_im1) * inv_dx2
                                            + (c_jp1 - Real(2.0) * c + c_jm1) * inv_dy2);

                    rhs_arr(i, j, 0) = advection + diffusion;
                }
            }
        }

        return 0;
    }

  private:
    static Geometry make_geometry(const ProblemParameters& params)
    {
        IntVect dom_lo(AMREX_D_DECL(0, 0, 0));
        IntVect dom_hi(AMREX_D_DECL(params.n_cell[0] - 1, params.n_cell[1] - 1, 0));
        Box domain(dom_lo, dom_hi);
        RealBox real_box(
            {AMREX_D_DECL(params.prob_lo[0], params.prob_lo[1], 0.0)},
            {AMREX_D_DECL(params.prob_hi[0], params.prob_hi[1], 1.0)});
        Vector<int> is_periodic(AMREX_SPACEDIM, 1);
        return Geometry(domain, &real_box, 0, is_periodic.data());
    }

    static BoxArray make_box_array(const Box& domain, const int max_grid_size)
    {
        BoxArray ba(domain);
        ba.maxSize(max_grid_size);
        return ba;
    }

    ProblemParameters m_params;
    Geometry m_geom;
    BoxArray m_ba;
    DistributionMapping m_dm;
    MultiFab m_state;
    long m_global_cells = 0;
};

int rhs_callback(sunrealtype t, N_Vector y, N_Vector ydot, void* user_data)
{
    auto* transport = static_cast<TransportSystem*>(user_data);
    return transport->rhs(t, y, ydot);
}

} // namespace

int main(int argc, char* argv[])
{
    amrex::Initialize(argc, argv);

    int status = 0;

    try {
        ProblemParameters params = read_parameters();
        TransportSystem transport(std::move(params));
        transport.initialize();

        auto* sunctx = amrex::sundials::The_Sundials_Context();
        check_ptr(sunctx, "amrex::sundials::The_Sundials_Context");

        N_Vector y =
            amrex::sundials::N_VMake_MultiFab(transport.vector_length(), &transport.state(), sunctx);
        check_ptr(y, "amrex::sundials::N_VMake_MultiFab");

        void* arkode_mem = ERKStepCreate(rhs_callback, Real(0.0), y, *sunctx);
        check_ptr(arkode_mem, "ERKStepCreate");

        check_flag(ARKodeSetUserData(arkode_mem, &transport), "ARKodeSetUserData");
        check_flag(ARKodeSStolerances(arkode_mem, transport.params().rtol, transport.params().atol),
                   "ARKodeSStolerances");
        check_flag(ARKodeSetInitStep(arkode_mem, transport.params().init_step), "ARKodeSetInitStep");
        check_flag(ARKodeSetMaxStep(arkode_mem, transport.params().max_step), "ARKodeSetMaxStep");
        check_flag(ARKodeSetMaxNumSteps(arkode_mem, 1000000), "ARKodeSetMaxNumSteps");

        amrex::Print() << "Chemical transport solve on " << transport.global_cells()
                       << " cells with " << amrex::ParallelDescriptor::NProcs()
                       << " MPI rank(s)\n";

        Real t = Real(0.0);
        int plot_index = 0;
        transport.write_plotfile(y, t, plot_index++);

        const Real final_time = transport.params().final_time;
        const Real output_interval = transport.params().output_interval;
        const Real eps = Real(1.0e-12) * std::max<Real>(Real(1.0), final_time);

        if (output_interval > Real(0.0)) {
            Real next_output = std::min(output_interval, final_time);
            while (t < final_time - eps) {
                check_flag(ARKodeEvolve(arkode_mem, next_output, y, &t, ARK_NORMAL),
                           "ARKodeEvolve");
                transport.write_plotfile(y, t, plot_index++);
                next_output = std::min(next_output + output_interval, final_time);
            }
        } else {
            check_flag(ARKodeEvolve(arkode_mem, final_time, y, &t, ARK_NORMAL),
                       "ARKodeEvolve");
            transport.write_plotfile(y, t, plot_index++);
        }

        long int num_steps = 0;
        check_flag(ARKodeGetNumSteps(arkode_mem, &num_steps), "ARKodeGetNumSteps");
        amrex::Print() << "Completed solve at t = " << t << " in " << num_steps
                       << " internal step(s)\n";

        ARKodeFree(&arkode_mem);
        N_VDestroy(y);
    } catch (const std::exception& ex) {
        amrex::Print() << "Unhandled exception: " << ex.what() << "\n";
        status = 1;
    }

    amrex::Finalize();
    return status;
}