Runge-Kutta support for AMR (AMReX-Codes#2974)

This adds RK2, RK3 and RK4 in a new namespace RungeKutta. Together with
the enhanced FillPatcher class, these functions can be used for RK time
stepping in AMR simulations. A new function AmrLevel::RK is added for
AmrLevel based codes. See CNS::advance in Tests/GPU/CNS/CNS_advance.cpp
for an example of using the new AmrLevel::RK function.

The main motivation for this PR is that ghost cell filling for high
order (> 2) RK methods at coarse/fine boundary is non-trivial when there
is subcycling.

Co-authored-by: Jean M. Sexton <jmsexton@lbl.gov>

Src/Amr/AMReX_AmrLevel.H

@@ -15,6 +15,7 @@
 #include <AMReX_StateDescriptor.H>
 #include <AMReX_StateData.H>
 #include <AMReX_VisMF.H>
+#include <AMReX_RungeKutta.H>
 #include <AMReX_FillPatcher.H>
 #ifdef AMREX_USE_EB
 #include <AMReX_EBSupport.H>
@@ -153,11 +154,10 @@ public:
                           int  ncycle) = 0;
 
     /**
-    * \brief Contains operations to be done after a timestep.  This is a
-    * pure virtual function and hence MUST be implemented by derived
-    * classes.
+    * \brief Contains operations to be done after a timestep.  If this
+    * function is overridden, don't forget to reset FillPatcher.
     */
-    virtual  void post_timestep (int iteration) = 0;
+    virtual void post_timestep (int iteration);
     /**
     * \brief Contains operations to be done only after a full coarse
     * timestep.  The default implementation does nothing.
@@ -397,8 +397,33 @@ public:
                               Real      time,
                               int       index,
                               int       scomp,
-                             int       ncomp,
-                             int       dcomp=0);
+                              int       ncomp,
+                              int       dcomp=0);
+
+    /**
+     * \brief Evolve one step with Runge-Kutta (2, 3, or 4)
+     *
+     * To use RK, the StateData must have all the ghost cells needed.  See
+     * namespace RungeKutta for expected function signatures of the callable
+     * parameters.
+     *
+     * \param order      order of RK
+     * \param state_type index of StateData
+     * \param time       time at the beginning of the step.
+     * \param dt         time step
+     * \param iteration  iteration number on fine level during a coarse time
+     *                   step.  For an AMR simulation with subcycling and a
+     *                   refinement ratio of 2, the number is either 1 or 2,
+     *                   denoting the first and second substep, respectively.
+     * \param ncycle     number of subcyling steps.  It's usually 2 or 4.
+     *                   Without subcycling, this will be 1.
+     * \param f          computing right-hand side for evolving the StateData.
+     *                   One can also register data for flux registers in this.
+     * \param p          optionally post-processing RK stage results
+     */
+    template <typename F, typename P = RungeKutta::PostStageNoOp>
+    void RK (int order, int state_type, Real time, Real dt, int iteration,
+             int ncycle, F&& f, P&& p = RungeKutta::PostStageNoOp());
 
 #ifdef AMREX_USE_EB
     static void SetEBMaxGrowCells (int nbasic, int nvolume, int nfull) noexcept {
@@ -457,6 +482,14 @@ protected:
 
 private:
 
+    template <std::size_t order>
+    void storeRKCoarseData (int state_type, Real time, Real dt,
+                            MultiFab const& S_old,
+                            Array<MultiFab,order> const& rkk);
+
+    void FillRKPatch (int state_index, MultiFab& S, Real time,
+                      int stage, int iteration, int ncycle);
+
     mutable BoxArray      edge_grids[AMREX_SPACEDIM];  // face-centered grids
     mutable BoxArray      nodal_grids;              // all nodal grids
 };
@@ -577,6 +610,74 @@ private:
     std::map< int,Vector< Vector< Vector<FillBoxId> > > > m_fbid; // [grid][level][fillablesubbox][oldnew]
 };
 
+template <typename F, typename P>
+void AmrLevel::RK (int order, int state_type, Real time, Real dt, int iteration,
+                   int ncycle, F&& f, P&& p)
+{
+    BL_PROFILE("AmrLevel::RK()");
+
+    AMREX_ASSERT(AmrLevel::desc_lst[state_type].nExtra() > 0); // Need ghost cells in StateData
+
+    MultiFab& S_old = get_old_data(state_type);
+    MultiFab& S_new = get_new_data(state_type);
+    const Real t_old = state[state_type].prevTime();
+    const Real t_new = state[state_type].curTime();
+    AMREX_ALWAYS_ASSERT(amrex::almostEqual(time,t_old) && amrex::almostEqual(time+dt,t_new));
+
+    if (order == 2) {
+        RungeKutta::RK2(S_old, S_new, time, dt, std::forward<F>(f),
+                        [&] (int /*stage*/, MultiFab& mf, Real t) {
+                            FillPatcherFill(mf, 0, mf.nComp(), mf.nGrow(), t,
+                                            state_type, 0); },
+                        std::forward<P>(p));
+    } else if (order == 3) {
+        RungeKutta::RK3(S_old, S_new, time, dt, std::forward<F>(f),
+                        [&] (int stage, MultiFab& mf, Real t) {
+                            FillRKPatch(state_type, mf, t, stage, iteration, ncycle);
+                        },
+                        [&] (Array<MultiFab,3> const& rkk) {
+                            if (level < parent->finestLevel()) {
+                                storeRKCoarseData(state_type, time, dt, S_old, rkk);
+                            }
+                        },
+                        std::forward<P>(p));
+    } else if (order == 4) {
+        RungeKutta::RK4(S_old, S_new, time, dt, std::forward<F>(f),
+                        [&] (int stage, MultiFab& mf, Real t) {
+                            FillRKPatch(state_type, mf, t, stage, iteration, ncycle);
+                        },
+                        [&] (Array<MultiFab,4> const& rkk) {
+                            if (level < parent->finestLevel()) {
+                                storeRKCoarseData(state_type, time, dt, S_old, rkk);
+                            }
+                        },
+                        std::forward<P>(p));
+    } else {
+        amrex::Abort("AmrLevel::RK: order = "+std::to_string(order)+" is not supported");
+    }
+}
+
+template <std::size_t order>
+void AmrLevel::storeRKCoarseData (int state_type, Real time, Real dt,
+                                  MultiFab const& S_old,
+                                  Array<MultiFab,order> const& rkk)
+{
+    if (level == parent->finestLevel()) { return; }
+
+    const StateDescriptor& desc = AmrLevel::desc_lst[state_type];
+
+    auto& fillpatcher = parent->getLevel(level+1).m_fillpatcher[state_type];
+    fillpatcher = std::make_unique<FillPatcher<MultiFab>>
+        (parent->boxArray(level+1), parent->DistributionMap(level+1),
+         parent->Geom(level+1),
+         parent->boxArray(level), parent->DistributionMap(level),
+         parent->Geom(level),
+         IntVect(desc.nExtra()), desc.nComp(), desc.interp(0));
+
+    fillpatcher->storeRKCoarseData(time, dt, S_old, rkk);
+}
+
+
 }
 
 #endif /*_AmrLevel_H_*/

Src/Amr/AMReX_AmrLevel.cpp

@@ -31,6 +31,14 @@ EBSupport AmrLevel::m_eb_support_level = EBSupport::volume;
 DescriptorList AmrLevel::desc_lst;
 DeriveList     AmrLevel::derive_lst;
 
+void
+AmrLevel::post_timestep (int /*iteration*/)
+{
+    if (level < parent->finestLevel()) {
+        parent->getLevel(level+1).resetFillPatcher();
+    }
+}
+
 void
 AmrLevel::postCoarseTimeStep (Real time)
 {
@@ -2223,4 +2231,23 @@ AmrLevel::CreateLevelDirectory (const std::string &dir)
     levelDirectoryCreated = true;
 }
 
+void
+AmrLevel::FillRKPatch (int state_index, MultiFab& S, Real time,
+                       int stage, int iteration, int ncycle)
+{
+    StateDataPhysBCFunct physbcf(state[state_index], 0, geom);
+
+    if (level == 0) {
+        S.FillBoundary(geom.periodicity());
+        physbcf(S, 0, S.nComp(), S.nGrowVect(), time, 0);
+    } else {
+        auto& crse_level = parent->getLevel(level-1);
+        StateDataPhysBCFunct physbcf_crse(crse_level.state[state_index], 0,
+                                          crse_level.geom);
+        auto& fillpatcher = m_fillpatcher[state_index];
+        fillpatcher->fillRK(stage, iteration, ncycle, S, time, physbcf_crse,
+                            physbcf, AmrLevel::desc_lst[state_index].getBCs());
+    }
+}
+
 }