Draft: Implementation of STS algorithm

inputs/diffusion/conduction.in

@@ -61,6 +61,11 @@ gas = true
 gravity = true
 conduction = true
 drag = true
+sts = false
+
+<sts>
+integrator = rkl1
+sts_max_dt_ratio = 200.0
 
 <gas>
 cfl = 0.3

src/CMakeLists.txt

@@ -74,6 +74,9 @@ set (SRC_LIST
 
   radiation/imc/imc.hpp
 
+  sts/sts.cpp
+  sts/sts.hpp
+
   utils/artemis_utils.cpp
   utils/artemis_utils.hpp
   utils/history.hpp

src/artemis.cpp

@@ -22,6 +22,7 @@
 #include "gravity/gravity.hpp"
 #include "nbody/nbody.hpp"
 #include "rotating_frame/rotating_frame.hpp"
+#include "sts/sts.hpp"
 #include "utils/artemis_utils.hpp"
 #include "utils/history.hpp"
 #include "utils/units.hpp"
@@ -70,6 +71,9 @@ Packages_t ProcessPackages(std::unique_ptr<ParameterInput> &pin) {
   const bool do_viscosity = pin->GetOrAddBoolean("physics", "viscosity", false);
   const bool do_conduction = pin->GetOrAddBoolean("physics", "conduction", false);
   const bool do_radiation = pin->GetOrAddBoolean("physics", "radiation", false);
+  const bool do_sts = pin->GetOrAddBoolean("physics", "sts", false);
+
+  artemis->AddParam("do_sts", do_sts);
   artemis->AddParam("do_gas", do_gas);
   artemis->AddParam("do_dust", do_dust);
   artemis->AddParam("do_gravity", do_gravity);
@@ -81,6 +85,11 @@ Packages_t ProcessPackages(std::unique_ptr<ParameterInput> &pin) {
   artemis->AddParam("do_conduction", do_conduction);
   artemis->AddParam("do_diffusion", do_conduction || do_viscosity);
   artemis->AddParam("do_radiation", do_radiation);
+
+  PARTHENON_REQUIRE(!(do_sts) || (do_sts && do_gas),
+                    "STS requires the gas package, but there is not gas!");
+  PARTHENON_REQUIRE(!(do_sts) || (do_sts && (do_conduction || do_viscosity)),
+                    "STS requires diffusion to be enabled!");
   PARTHENON_REQUIRE(!(do_cooling) || (do_cooling && do_gas),
                     "Cooling requires the gas package, but there is not gas!");
   PARTHENON_REQUIRE(!(do_viscosity) || (do_viscosity && do_gas),
@@ -105,6 +114,8 @@ Packages_t ProcessPackages(std::unique_ptr<ParameterInput> &pin) {
   if (do_rotating_frame) packages.Add(RotatingFrame::Initialize(pin.get()));
   if (do_cooling) packages.Add(Gas::Cooling::Initialize(pin.get()));
   if (do_drag) packages.Add(Drag::Initialize(pin.get()));
+  if (do_sts) packages.Add(STS::Initialize(pin.get()));
+
   if (do_radiation) {
     auto eos_h = packages.Get("gas")->Param<EOS>("eos_h");
     auto opacity_h = packages.Get("gas")->Param<Opacity>("opacity_h");

src/artemis.hpp

@@ -104,6 +104,9 @@ enum class ArtemisBC {
   none
 };
 
+// ...STS integrator types
+enum class STSInt { rkl1, rkl2, null };
+
 // Floating point limits
 template <typename T = Real>
 KOKKOS_FORCEINLINE_FUNCTION constexpr auto Big() {

src/artemis_driver.cpp

@@ -29,6 +29,7 @@
 #include "nbody/nbody.hpp"
 #include "radiation/imc/imc.hpp"
 #include "rotating_frame/rotating_frame.hpp"
+#include "sts/sts.hpp"
 #include "utils/integrators/artemis_integrator.hpp"
 
 using namespace parthenon::driver::prelude;
@@ -65,6 +66,7 @@ ArtemisDriver<GEOM>::ArtemisDriver(ParameterInput *pin, ApplicationInput *app_in
   do_conduction = artemis_pkg->template Param<bool>("do_conduction");
   do_nbody = artemis_pkg->template Param<bool>("do_nbody");
   do_diffusion = do_viscosity || do_conduction;
+  do_sts = artemis_pkg->template Param<bool>("do_sts");
   do_radiation = artemis_pkg->template Param<bool>("do_radiation");
 
   // NBody initialization tasks
@@ -103,10 +105,14 @@ TaskListStatus ArtemisDriver<GEOM>::Step() {
   // Prepare registers
   PreStepTasks();
 
+  // Execute STS first stage
+  if (do_sts) STSFirstStage();
+
   // Execute explicit, unsplit physics
   auto status = StepTasks().Execute();
   if (status != TaskListStatus::complete) return status;
 
+  // STS_second_stage();
   // Execute operator split physics
   if (do_radiation) status = IMC::JaybenneIMC<GEOM>(pmesh, tm.time, tm.dt);
   if (status != TaskListStatus::complete) return status;
@@ -138,6 +144,31 @@ void ArtemisDriver<GEOM>::PreStepTasks() {
   auto &u1 = pmesh->mesh_data.Add("u1", u0);
 }
 
+//----------------------------------------------------------------------------------------
+//! \fn TaskCollection ArtemisDriver::STS_first_stage
+//! \brief Define the tasks for the first stage of the STS integrator
+template <Coordinates GEOM>
+void ArtemisDriver<GEOM>::STSFirstStage() {
+
+  // Assign sts registers and First stage of STS integration
+  auto &gas_pkg = pmesh->packages.Get("gas");
+  auto min_diff_dt = gas_pkg->template Param<Real>("diff_dt");
+  auto &sts_pkg = pmesh->packages.Get("STS");
+  auto info_output = sts_pkg->template Param<bool>("info_output");
+  // compute the number of stages needed for the STS integrator (for rkl1 only)
+  int s_sts = static_cast<int>(0.5 * (-1. + std::sqrt(1. + 8. * tm.dt / min_diff_dt)));
+  if (s_sts % 2 == 0) s_sts += 1;
+
+  if (parthenon::Globals::my_rank == 0 && info_output) {
+    Real ratio = tm.dt / min_diff_dt;
+    std::cout << "STS ratio: " << ratio << ", Taking " << s_sts << " steps." << std::endl;
+    if (ratio > 200.0) {
+      std::cout << "WARNING: ratio is > 200. Proceed at own risk." << std::endl;
+    }
+  }
+  STS::PreStepSTSTasks<GEOM>(pmesh, tm.time, tm.dt, s_sts);
+}
+
 //----------------------------------------------------------------------------------------
 //! \fn TaskCollection ArtemisDriver::PreStepTasks
 //! \brief Defines the main integrator's TaskCollection for the ArtemisDriver
@@ -186,7 +217,7 @@ TaskCollection ArtemisDriver<GEOM>::StepTasks() {
 
       // Compute (gas) diffusive fluxes
       TaskID diff_flx = none;
-      if (do_diffusion && do_gas) {
+      if (do_diffusion && do_gas && !(do_sts)) {
         auto zf = tl.AddTask(none, Gas::ZeroDiffusionFlux, u0.get());
         TaskID vflx = zf, tflx = zf;
         if (do_viscosity) vflx = tl.AddTask(zf, Gas::ViscousFlux<GEOM>, u0.get());
@@ -215,7 +246,8 @@ TaskCollection ArtemisDriver<GEOM>::StepTasks() {
       // NOTE(@pdmullen): I believe set_flx dependency implicitly inside gas_coord_src,
       // but included below explicitly for posterity
       TaskID gas_diff_src = gas_coord_src | diff_flx | set_flx;
-      if (do_diffusion && do_gas) {
+
+      if (do_diffusion && do_gas && !(do_sts)) {
         gas_diff_src = tl.AddTask(gas_coord_src | diff_flx | set_flx,
                                   Gas::DiffusionUpdate<GEOM>, u0.get(), bdt);
       }
@@ -252,18 +284,18 @@ TaskCollection ArtemisDriver<GEOM>::StepTasks() {
           tl.AddTask(cooling_src, ArtemisDerived::SetAuxillaryFields<GEOM>, u0.get());
 
       // Set (remaining) fields to be communicated
-      auto c2p = tl.AddTask(set_aux, PreCommFillDerived<MeshData<Real>>, u0.get());
+      auto pre_comm = tl.AddTask(set_aux, PreCommFillDerived<MeshData<Real>>, u0.get());
 
       // Set boundary conditions (both physical and logical)
-      auto bcs = parthenon::AddBoundaryExchangeTasks(c2p, tl, u0, pmesh->multilevel);
+      auto bcs = parthenon::AddBoundaryExchangeTasks(pre_comm, tl, u0, pmesh->multilevel);
 
-      // Sync fields
-      auto p2c = tl.AddTask(TQ::local_sync, bcs, FillDerived<MeshData<Real>>, u0.get());
+      // Update primitive variables
+      auto c2p = tl.AddTask(TQ::local_sync, bcs, FillDerived<MeshData<Real>>, u0.get());
 
       // Advance nbody integrator
-      TaskID nbadv = p2c;
+      TaskID nbadv = c2p;
       if (do_nbody) {
-        nbadv = tl.AddTask(TQ::once_per_region, p2c, NBody::Advance, pmesh, time, stage,
+        nbadv = tl.AddTask(TQ::once_per_region, c2p, NBody::Advance, pmesh, time, stage,
                            nbody_integrator.get());
       }
     }
@@ -281,6 +313,8 @@ TaskCollection ArtemisDriver<GEOM>::PostStepTasks() {
   TaskCollection tc;
   TaskID none(0);
 
+  // TODO: Implement RKL2 STS integration
+
   const int num_partitions = pmesh->DefaultNumPartitions();
   auto &post_region = tc.AddRegion(num_partitions);
   for (int i = 0; i < num_partitions; i++) {
@@ -305,31 +339,37 @@ typedef ApplicationInput AI;
 template ArtemisDriver<C::cartesian>::ArtemisDriver(PI *p, AI *a, M *m, const bool r);
 template TaskListStatus ArtemisDriver<C::cartesian>::Step();
 template void ArtemisDriver<C::cartesian>::PreStepTasks();
+template void ArtemisDriver<C::cartesian>::STSFirstStage();
 template TaskCollection ArtemisDriver<C::cartesian>::StepTasks();
 template TaskCollection ArtemisDriver<C::cartesian>::PostStepTasks();
 template ArtemisDriver<C::cylindrical>::ArtemisDriver(PI *p, AI *a, M *m, const bool r);
 template TaskListStatus ArtemisDriver<C::cylindrical>::Step();
 template void ArtemisDriver<C::cylindrical>::PreStepTasks();
+template void ArtemisDriver<C::cylindrical>::STSFirstStage();
 template TaskCollection ArtemisDriver<C::cylindrical>::StepTasks();
 template TaskCollection ArtemisDriver<C::cylindrical>::PostStepTasks();
 template ArtemisDriver<C::spherical1D>::ArtemisDriver(PI *p, AI *a, M *m, const bool r);
 template TaskListStatus ArtemisDriver<C::spherical1D>::Step();
 template void ArtemisDriver<C::spherical1D>::PreStepTasks();
+template void ArtemisDriver<C::spherical1D>::STSFirstStage();
 template TaskCollection ArtemisDriver<C::spherical1D>::StepTasks();
 template TaskCollection ArtemisDriver<C::spherical1D>::PostStepTasks();
 template ArtemisDriver<C::spherical2D>::ArtemisDriver(PI *p, AI *a, M *m, const bool r);
 template TaskListStatus ArtemisDriver<C::spherical2D>::Step();
 template void ArtemisDriver<C::spherical2D>::PreStepTasks();
+template void ArtemisDriver<C::spherical2D>::STSFirstStage();
 template TaskCollection ArtemisDriver<C::spherical2D>::StepTasks();
 template TaskCollection ArtemisDriver<C::spherical2D>::PostStepTasks();
 template ArtemisDriver<C::spherical3D>::ArtemisDriver(PI *p, AI *a, M *m, const bool r);
 template TaskListStatus ArtemisDriver<C::spherical3D>::Step();
 template void ArtemisDriver<C::spherical3D>::PreStepTasks();
+template void ArtemisDriver<C::spherical3D>::STSFirstStage();
 template TaskCollection ArtemisDriver<C::spherical3D>::StepTasks();
 template TaskCollection ArtemisDriver<C::spherical3D>::PostStepTasks();
 template ArtemisDriver<C::axisymmetric>::ArtemisDriver(PI *p, AI *a, M *m, const bool r);
 template TaskListStatus ArtemisDriver<C::axisymmetric>::Step();
 template void ArtemisDriver<C::axisymmetric>::PreStepTasks();
+template void ArtemisDriver<C::axisymmetric>::STSFirstStage();
 template TaskCollection ArtemisDriver<C::axisymmetric>::StepTasks();
 template TaskCollection ArtemisDriver<C::axisymmetric>::PostStepTasks();
 

src/artemis_driver.hpp

@@ -46,14 +46,15 @@ class ArtemisDriver : public EvolutionDriver {
                 const bool is_restart_in);
   TaskListStatus Step();
   void PreStepTasks();
+  void STSFirstStage();
   TaskCollection StepTasks();
   TaskCollection PostStepTasks();
 
  protected:
   IntegratorPtr_t integrator, nbody_integrator;
   StateDescriptor *artemis_pkg;
   bool do_gas, do_dust, do_gravity, do_rotating_frame, do_cooling, do_drag, do_viscosity,
-      do_nbody, do_conduction, do_diffusion, do_radiation;
+      do_nbody, do_conduction, do_diffusion, do_sts, do_radiation;
   const bool is_restart;
 };
 

src/gas/gas.cpp

@@ -18,6 +18,7 @@
 #include "artemis.hpp"
 #include "gas.hpp"
 #include "geometry/geometry.hpp"
+#include "sts/sts.hpp"
 #include "utils/artemis_utils.hpp"
 #include "utils/diffusion/diffusion.hpp"
 #include "utils/diffusion/diffusion_coeff.hpp"
@@ -183,6 +184,15 @@ std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin,
   const bool do_conduction = pin->GetOrAddBoolean("physics", "conduction", false);
   params.Add("do_conduction", do_conduction);
 
+  const bool do_sts = pin->GetOrAddBoolean("physics", "sts", false);
+  params.Add("do_sts", do_sts);
+
+  if (do_sts) {
+    params.Add("diff_dt", std::numeric_limits<Real>::max(), Params::Mutability::Mutable);
+    Real sts_max_dt_ratio = pin->GetOrAddReal("sts", "sts_max_dt_ratio", -1.0);
+    params.Add("sts_max_dt_ratio", sts_max_dt_ratio);
+  }
+
   const bool do_diffusion = do_viscosity || do_conduction;
   params.Add("do_diffusion", do_diffusion);
 
@@ -388,6 +398,7 @@ std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin,
 //----------------------------------------------------------------------------------------
 //! \fn  Real Gas::EstimateTimestepMesh
 //! \brief Compute gas hydrodynamics timestep
+//! STS_Flag
 template <Coordinates GEOM>
 Real EstimateTimestepMesh(MeshData<Real> *md) {
   using parthenon::MakePackDescriptor;
@@ -464,7 +475,23 @@ Real EstimateTimestepMesh(MeshData<Real> *md) {
   Real diff_dt = std::min(visc_dt, cond_dt);
 
   const auto cfl_number = params.template Get<Real>("cfl");
-  return cfl_number * std::min(min_dt, diff_dt);
+
+  // STS Time Stepping Control
+  const auto do_sts = params.template Get<bool>("do_sts");
+  if (do_sts) {
+    const auto sts_max_dt_ratio = params.template Get<Real>("sts_max_dt_ratio");
+    auto dt_ratio = min_dt / diff_dt;
+    // limit the timestep within the STS ratio, otherwise use the hyperbolic timestep
+    if (sts_max_dt_ratio > 0.0 && dt_ratio > sts_max_dt_ratio) {
+      min_dt = sts_max_dt_ratio * diff_dt;
+    }
+    // update the parabolic timestep
+    gas_pkg->UpdateParam("diff_dt", cfl_number * diff_dt);
+  } else {
+    min_dt = std::min(min_dt, diff_dt);
+  }
+
+  return cfl_number * min_dt;
 }
 
 //----------------------------------------------------------------------------------------