Adaptive integrator

Src/Base/AMReX_IntegratorBase.H

@@ -44,8 +44,8 @@ struct IntegratorOps<T>
     static void Saxpy (T& Y, const amrex::Real a, T& X)
     {
         // Calculate Y += a * X using a particle-level saxpy function supplied by the particle container T
-        using TParIter = amrex::ParIter<T::NStructReal, T::NStructInt, T::NArrayReal, T::NArrayInt>;
-        using ParticleType = amrex::Particle<T::NStructReal, T::NStructInt>;
+        using TParIter = typename T::ParIterType;
+        using ParticleType = typename T::ParticleType;
 
         int lev = 0;
         TParIter pty(Y, lev);
@@ -69,16 +69,33 @@ struct IntegratorOps<T>
             const int npx  = ptx.numParticles();
             AMREX_ALWAYS_ASSERT(npy == npx);
 
-            ParticleType* psy = &(pty.GetArrayOfStructs()[0]);
-            ParticleType* psx = &(ptx.GetArrayOfStructs()[0]);
-
             auto particle_apply_rhs = T::particle_apply_rhs;
 
-            amrex::ParallelFor ( npy, [=] AMREX_GPU_DEVICE (int i) {
-                ParticleType& py = psy[i];
-                const ParticleType& px = psx[i];
-                particle_apply_rhs(py, a, px);
-            });
+            // Hand particle_apply_rhs an array-of-structs particle when the
+            // container has one and the supplied function accepts it.
+            if constexpr ( ! ParticleType::is_soa_particle &&
+                           IsCallable<decltype(particle_apply_rhs),
+                                      ParticleType&, amrex::Real,
+                                      ParticleType&>::value) {
+                ParticleType* psy = &(pty.GetArrayOfStructs()[0]);
+                ParticleType* psx = &(ptx.GetArrayOfStructs()[0]);
+
+                amrex::ParallelFor ( npy, [=] AMREX_GPU_DEVICE (int i) {
+                    ParticleType& py = psy[i];
+                    const ParticleType& px = psx[i];
+                    particle_apply_rhs(py, a, px);
+                });
+
+            // Otherwise pass the particle tile data and an index, which is the only form
+            // available once the real components live in struct-of-arrays storage
+            } else {
+                auto ytd = pty.GetParticleTile().getParticleTileData();
+                auto xtd = ptx.GetParticleTile().getParticleTileData();
+
+                amrex::ParallelFor ( npy, [=] AMREX_GPU_DEVICE (int i) {
+                    particle_apply_rhs(ytd, i, a, xtd);
+                });
+            }
         }
     }
 
@@ -226,6 +243,14 @@ protected:
      */
     std::function<void (T&, amrex::Real)> post_fast_step_action;
 
+    /**
+     * \brief User-supplied error norm. This must be set via set_error_norm
+     * before using adaptive time stepping; no default is provided.
+     * Arguments are (error_state, S_old, S_new, abs_tol, rel_tol) and the
+     * return value is the scaled error used to decide step acceptance.
+     */
+    std::function<amrex::Real (T&, T&, T&, amrex::Real, amrex::Real)> custom_error_norm;
+
     /**
      * \brief Flag to enable/disable adaptive time stepping in single rate
      * methods or at the slow time scale in multirate methods (bool)
@@ -286,6 +311,12 @@ protected:
      */
     amrex::Real fast_abs_tol = 1.0e-9;
 
+    /**
+     * \brief Most recent scaled error computed by the adaptive controller.
+     * Updated each time advance() runs the error norm; otherwise zero.
+     */
+    amrex::Real last_scaled_error = 0.0;
+
 
 public:
     IntegratorBase () {
@@ -335,11 +366,26 @@ public:
         post_fast_step_action = A;
     }
 
+    void set_error_norm (std::function<amrex::Real (T&, T&, T&, amrex::Real, amrex::Real)> F)
+    {
+        custom_error_norm = F;
+    }
+
     amrex::Real get_time_step ()
     {
         return time_step;
     }
 
+    amrex::Real get_previous_time_step ()
+    {
+        return previous_time_step;
+    }
+
+    amrex::Real get_scaled_error ()
+    {
+        return last_scaled_error;
+    }
+
     void set_time_step (amrex::Real dt)
     {
         time_step = dt;

Src/Base/AMReX_RKIntegrator.H

@@ -5,6 +5,7 @@
 #include <AMReX_ParmParse.H>
 #include <AMReX_IntegratorBase.H>
 #include <functional>
+#include <limits>
 
 namespace amrex {
 
@@ -14,6 +15,7 @@ enum struct ButcherTableauTypes {
     Trapezoid,
     SSPRK3,
     RK4,
+    DormandPrince,
     NumTypes
 };
 
@@ -41,9 +43,22 @@ private:
     // RK embedded method b vector
     amrex::Vector<amrex::Real> extended_weights;
 
+    // Order of the embedded error estimate (e.g. 4 for Dormand-Prince)
+    int error_order;
+
+    // Safety factor for adaptive step-size control
+    amrex::Real safety_factor;
+
+    // Min/max allowed step-size growth/shrink factors per step
+    amrex::Real min_factor;
+    amrex::Real max_factor;
+
     // RK stage right-hand sides
     amrex::Vector<std::unique_ptr<T> > F_nodes;
 
+    // Scratch space for low-order solution in embedded error estimate
+    amrex::Vector<std::unique_ptr<T> > S_error_vec;
+
     // Current (internal) state and time
     amrex::Vector<std::unique_ptr<T> > S_current;
     amrex::Real time_current;
@@ -84,6 +99,28 @@ private:
                         {0.0, 0.0, 1.0, 0.0}};
                 weights = {1./6., 1./3., 1./3., 1./6.};
                 break;
+            case ButcherTableauTypes::DormandPrince:
+                // Dormand-Prince RK4(5) embedded pair
+                nodes = {0.0,
+                         1.0/5.0,
+                         3.0/10.0,
+                         4.0/5.0,
+                         8.0/9.0,
+                         1.0,
+                         1.0};
+                tableau = {{0.0},
+                           {1.0/5.0, 0.0},
+                           {3.0/40.0, 9.0/40.0, 0.0},
+                           {44.0/45.0, -56.0/15.0, 32.0/9.0, 0.0},
+                           {19372.0/6561.0, -25360.0/2187.0, 64448.0/6561.0, -212.0/729.0, 0.0},
+                           {9017.0/3168.0, -355.0/33.0, 46732.0/5247.0, 49.0/176.0, -5103.0/18656.0, 0.0},
+                           {35.0/384.0, 0.0, 500.0/1113.0, 125.0/192.0, -2187.0/6784.0, 11.0/84.0, 0.0}};
+                // 5th-order weights
+                weights = {35.0/384.0, 0.0, 500.0/1113.0, 125.0/192.0, -2187.0/6784.0, 11.0/84.0, 0.0};
+                // 4th-order weights for embedded error estimate
+                extended_weights = {5179.0/57600.0, 0.0, 7571.0/16695.0, 393.0/640.0, -92097.0/339200.0, 187.0/2100.0, 1.0/40.0};
+                error_order = 4;
+                break;
             default:
                 amrex::Error("Invalid RK Integrator tableau type");
                 break;
@@ -104,6 +141,16 @@ private:
         // By default, define no extended weights
         extended_weights = {};
 
+        // Safety factor for adaptive step-size control (default 0.9)
+        safety_factor = 0.9;
+        pp.query("safety_factor", safety_factor);
+
+        // Min/max allowed step-size growth/shrink factors per step
+        min_factor = 0.1;
+        max_factor = 5.0;
+        pp.query("min_factor", min_factor);
+        pp.query("max_factor", max_factor);
+
         if (tableau_type == ButcherTableauTypes::User)
         {
             // Read weights/nodes/butcher tableau
@@ -169,6 +216,11 @@ private:
         IntegratorOps<T>::CreateLike(S_current, S_data, true);
         IntegratorOps<T>::Copy(*S_current[0], S_data);
 
+        // Allocate error scratch space for embedded methods
+        if (!extended_weights.empty()) {
+            IntegratorOps<T>::CreateLike(S_error_vec, S_data);
+        }
+
         // Set the initial time
         time_current = time;
     }
@@ -231,13 +283,49 @@ public:
             IntegratorOps<T>::Saxpy(S_new, dt * weights[i], *F_nodes[i]);
         }
 
-        BaseT::post_step_action(S_new, time + dt);
+        BaseT::previous_time_step = dt;
 
         // If we are working with an extended Butcher tableau, we can estimate the error here,
         // and then calculate an adaptive time step.
+        amrex::Real dt_next = dt;
+        if (!extended_weights.empty() && BaseT::use_adaptive_time_step)
+        {
+            // Compute low-order solution: S_error = S_old + dt * sum(extended_weights[i] * F[i])
+            IntegratorOps<T>::Copy(*S_error_vec[0], S_old);
+            for (int i = 0; i < number_nodes; ++i)
+                IntegratorOps<T>::Saxpy(*S_error_vec[0], dt * extended_weights[i], *F_nodes[i]);
 
-        // Save last completed step size for time_interpolate
-        BaseT::previous_time_step = dt;
+            // S_error = S_low - S_high (pointwise error estimate)
+            IntegratorOps<T>::Saxpy(*S_error_vec[0], -1.0, S_new);
+
+            AMREX_ALWAYS_ASSERT_WITH_MESSAGE(static_cast<bool>(BaseT::custom_error_norm),
+                "Adaptive time stepping requires an error norm to be set via set_error_norm.");
+            const amrex::Real scaled_err = BaseT::custom_error_norm(*S_error_vec[0], S_old, S_new,
+                                                                    BaseT::abs_tol, BaseT::rel_tol);
+            BaseT::last_scaled_error = scaled_err;
+
+            // Step-size controller: exponent is 1/(error_order+1) per standard embedded RK theory
+            amrex::Real factor = safety_factor * std::pow(1.0 / scaled_err, 1.0 / (error_order + 1));
+            factor = amrex::min(amrex::max(factor, min_factor), max_factor);
+            dt_next = dt * factor;
+
+            if (scaled_err > 1.0)
+            {
+                // Step rejected: retry with a smaller dt without calling post_step_action
+                return advance(S_old, S_new, time, dt_next);
+            }
+
+            // Step accepted. Set time_step to max before calling post_step_action so that
+            // any CFL-based set_time_step call inside it can be detected and used to cap
+            // the adaptive suggestion.
+            BaseT::time_step = std::numeric_limits<amrex::Real>::max();
+        }
+
+        BaseT::post_step_action(S_new, time + dt);
+
+        // Cap dt_next with any CFL limit set by post_step_action, then keep time_step in sync.
+        dt_next = amrex::min(dt_next, BaseT::time_step);
+        BaseT::time_step = dt_next;
 
         // Return time step
         return dt;

Src/Base/AMReX_TimeIntegrator.H

@@ -167,11 +167,26 @@ public:
         integrator_ptr->set_post_fast_step_action(A);
     }
 
+    void set_error_norm (std::function<amrex::Real (T&, T&, T&, amrex::Real, amrex::Real)> F)
+    {
+        integrator_ptr->set_error_norm(F);
+    }
+
     amrex::Real get_time_step ()
     {
         return integrator_ptr->get_time_step();
     }
 
+    amrex::Real get_previous_time_step ()
+    {
+        return integrator_ptr->get_previous_time_step();
+    }
+
+    amrex::Real get_scaled_error ()
+    {
+        return integrator_ptr->get_scaled_error();
+    }
+
     void set_time_step (amrex::Real dt)
     {
         integrator_ptr->set_time_step(dt);
@@ -237,8 +252,11 @@ public:
             // Call the time integrator advance
             integrator_ptr->advance(S_old, S_new, m_time, m_timestep);
 
-            // Update our time variable
-            m_time += m_timestep;
+            // Update our time variable using the actual step size taken
+            m_time += integrator_ptr->get_previous_time_step();
+
+            // Pick up any timestep changes from post_step_action or adaptive control
+            m_timestep = integrator_ptr->get_time_step();
         }
     }