add user source term to pen file "disk.hpp"

src/artemis.hpp

@@ -238,6 +238,8 @@ inline int ProblemDimension(parthenon::ParameterInput *pin) {
 // Custom AMR criteria
 namespace artemis {
 extern std::function<AmrTag(MeshBlockData<Real> *mbd)> ProblemCheckRefinementBlock;
+extern std::function<TaskStatus(MeshData<Real> *md, const Real time, const Real dt)>
+    UserSourceTerm;
 } // namespace artemis
 
 #endif // ARTEMIS_ARTEMIS_HPP_

src/artemis_driver.cpp

@@ -302,9 +302,12 @@ TaskCollection ArtemisDriver<GEOM>::StepTasks() {
             tl.AddTask(rt_src, RotatingFrame::RotatingFrameForce, u0.get(), time, bdt);
       }
 
+      auto user_src =
+          tl.AddTask(rframe_src, ArtemisUtils::ProblemSourceTerm, u0.get(), time, bdt);
+
       // Apply drag source term
       // NOTE(@pdmullen): RK integrated, operator split drag (RHS computed from U)
-      TaskID drag_src = rframe_src;
+      TaskID drag_src = user_src;
       if (do_drag) {
         drag_src = tl.AddTask(rframe_src, Drag::DragSource<GEOM>, u0.get(), time, bdt);
       }

src/drag/drag.hpp

@@ -125,6 +125,7 @@ struct StoppingTimeParams {
   DragModel model;
   ParArray1D<Real> tau;
   Real tau_max, tau_min;
+  Real x1_power;
 
   StoppingTimeParams(std::string block_name, ParameterInput *pin) {
     const std::string choice = pin->GetString(block_name, "type");
@@ -146,6 +147,7 @@ struct StoppingTimeParams {
       scale = pin->GetOrAddReal(block_name, "scale", 1.0);
       tau_max = pin->GetOrAddReal(block_name, "maximum", 1e99);
       tau_min = pin->GetOrAddReal(block_name, "minimum", 0.0);
+      x1_power = pin->GetOrAddReal(block_name, "x1_power", 0.0);
       auto h_tau = tau.GetHostMirror();
       for (int n = 0; n < nd; n++) {
         h_tau(n) = scale;
@@ -463,9 +465,11 @@ TaskStatus SimpleDragSourceImpl(MeshData<Real> *md, const Real time, const Real
         // Extract Stokes specific parameters
         [[maybe_unused]] auto &grain_density_ = grain_density;
         [[maybe_unused]] Real vth = Null<Real>();
+        Real x1fac = 1.0;
         if constexpr (DRAG == DragModel::stokes) {
           const Real gm1 = eos_d.GruneisenParamFromDensityInternalEnergy(dg, sieg);
           vth = std::sqrt(8.0 / M_PI * gm1 * sieg);
+          x1fac = std::pow(xv[0], tp.x1_power);
         }
 
         // First pass to collect \sum rho' and \sum rho' v and compute new vg
@@ -493,8 +497,9 @@ TaskStatus SimpleDragSourceImpl(MeshData<Real> *md, const Real time, const Real
           Real tc = tp.tau(id);
           [[maybe_unused]] auto &sizes_ = sizes;
           if constexpr (DRAG == DragModel::stokes) {
-            tc = std::max(tp.tau_min, std::min(tp.tau_max, tp.scale * grain_density_ /
-                                                               dg * sizes_(id) / vth));
+            tc =
+                std::max(tp.tau_min, std::min(tp.tau_max, tp.scale * grain_density_ / dg *
+                                                              sizes_(id) / vth * x1fac));
           }
           const Real alpha = dt * ((tc <= 0.0) ? Big<Real>() : 1.0 / tc);
           for (int d = 0; d < 3; d++) {
@@ -529,8 +534,9 @@ TaskStatus SimpleDragSourceImpl(MeshData<Real> *md, const Real time, const Real
           Real tc = tp.tau(id);
           [[maybe_unused]] auto &sizes_ = sizes;
           if constexpr (DRAG == DragModel::stokes) {
-            tc = std::max(tp.tau_min, std::min(tp.tau_max, tp.scale * grain_density_ /
-                                                               dg * sizes_(id) / vth));
+            tc =
+                std::max(tp.tau_min, std::min(tp.tau_max, tp.scale * grain_density_ / dg *
+                                                              sizes_(id) / vth * x1fac));
           }
           const Real alpha = dt * ((tc <= 0.0) ? Big<Real>() : 1.0 / tc);
           // Update dust momenta

src/pgen/disk.hpp

@@ -1,3 +1,4 @@
+
 //========================================================================================
 // (C) (or copyright) 2023-2025. Triad National Security, LLC. All rights reserved.
 //
@@ -65,6 +66,7 @@ struct DiskParams {
   Real mdot;
   Real temp_soft2;
   Real kbmu, ar;
+  Real nH_reset;
   bool do_gas, do_dust, do_moment, do_imc;
   bool nbody_temp;
   bool quiet_start;
@@ -201,7 +203,11 @@ KOKKOS_INLINE_FUNCTION State ComputeDiskProfile(
   const Real vp2 = vk2 + (dpdr / res.gdens) * xcyl[0];
   const Real vp = (vp2 < 0.0) ? 0.0 : std::sqrt(vp2);
   const Real nu = ViscosityProfile(pgen, eos_d, rt, xcyl[2]);
-  const Real vr = pgen.quiet_start ? 0.0 : -1.5 * nu / xcyl[0];
+  const Real vr_visc = pgen.alpha > 0.0
+                           ? (-3.0 * pgen.alpha / std::sqrt(xcyl[0]) / res.gdens *
+                              (SQR(xcyl[0]) * dpdr + pres * 2.0 * xcyl[0]))
+                           : (-1.5 * nu / xcyl[0]);
+  const Real vr = pgen.quiet_start ? 0.0 : vr_visc;
 
   // Construct the total cylindrical velocity
   const Real vcyl[3] = {vr, vp - pgen.omf * xcyl[0], 0.0};
@@ -257,6 +263,7 @@ inline void InitDiskParams(MeshBlock *pmb, ParameterInput *pin) {
     disk_params.l0 = pin->GetOrAddReal("problem", "l0", 0.0);
     disk_params.dust_to_gas = pin->GetOrAddReal("problem", "dust_to_gas", 0.01);
     disk_params.temp_soft2 = pin->GetOrAddReal("problem", "temp_soft", 0.0);
+    disk_params.nH_reset = pin->GetOrAddReal("problem", "nH_reset", 0.0);
     const auto mu = gas_pkg->Param<Real>("mu");
     auto &constants = artemis_pkg->Param<ArtemisUtils::Constants>("constants");
     const auto &eos = gas_pkg->Param<ArtemisUtils::EOS>("eos_h");
@@ -890,6 +897,90 @@ void DiskBoundaryExtrap(std::shared_ptr<MeshBlockData<Real>> &mbd, bool coarse)
       });
 }
 
+//----------------------------------------------------------------------------------------
+//! \fn TaskStatus UserSourceTerm()
+//! \brief Custom source term for disk pgen
+template <Coordinates GEOM>
+TaskStatus UserSourceTerm(MeshData<Real> *md, const Real time, const Real dt) {
+  if (GEOM == Coordinates::spherical3D || GEOM == Coordinates::spherical2D) {
+    // reset the gas and dust state above nH_reset: gas to initial values, dust to zeros
+    using parthenon::MakePackDescriptor;
+    auto pm = md->GetParentPointer();
+
+    auto &artemis_pkg = pm->packages.Get("artemis");
+
+    // Disk parameters
+    const auto &pgen = artemis_pkg->template Param<DiskParams>("disk_params");
+    if (pgen.nH_reset <= 0.0) return TaskStatus::complete;
+
+    // Extract gas package and params
+    auto &gas_pkg = pm->packages.Get("gas");
+    const auto &eos_d = gas_pkg->template Param<EOS>("eos_d");
+
+    auto &resolved_pkgs = pm->resolved_packages;
+
+    const bool do_gas = artemis_pkg->template Param<bool>("do_gas");
+    const bool do_dust = artemis_pkg->template Param<bool>("do_dust");
+
+    const auto &cpars =
+        artemis_pkg->template Param<geometry::CoordParams>("coord_params");
+
+    static auto desc =
+        MakePackDescriptor<gas::cons::momentum, gas::cons::total_energy,
+                           dust::cons::momentum, gas::cons::density, gas::prim::density,
+                           dust::prim::density, dust::cons::density>(resolved_pkgs.get());
+    auto vmesh = desc.GetPack(md);
+    const auto ib = md->GetBoundsI(IndexDomain::interior);
+    const auto jb = md->GetBoundsJ(IndexDomain::interior);
+    const auto kb = md->GetBoundsK(IndexDomain::interior);
+
+    parthenon::par_for(
+        DEFAULT_LOOP_PATTERN, "UserSourceTerm", parthenon::DevExecSpace(), 0,
+        md->NumBlocks() - 1, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e,
+        KOKKOS_LAMBDA(const int &b, const int &k, const int &j, const int &i) {
+          // Extract coordinates
+          geometry::Coords<GEOM> coords(cpars, vmesh.GetCoordinates(b), k, j, i);
+          const auto &hx = coords.GetScaleFactors();
+          const auto &xv = coords.GetCellCenter();
+          const auto &xcyl = coords.ConvertToCyl(xv);
+          const Real H = xcyl[0] * pgen.h0 * std::pow(xcyl[0] / pgen.r0, pgen.flare);
+          if (std::abs(xcyl[2]) > pgen.nH_reset * H) {
+            if (do_gas) {
+              const Real nu = ViscosityProfile(pgen, eos_d, xcyl[0], xcyl[2]);
+              const Real vr = pgen.quiet_start ? 0.0 : -1.5 * nu / xcyl[0];
+              for (int n = 0; n < vmesh.GetSize(b, gas::prim::density()); ++n) {
+                const Real etot_res =
+                    vmesh(b, gas::cons::total_energy(n), k, j, i) -
+                    0.5 *
+                        (SQR(vmesh(b, gas::cons::momentum(VI(n, 0)), k, j, i)) +
+                         SQR(vmesh(b, gas::cons::momentum(VI(n, 1)), k, j, i))) /
+                        vmesh(b, gas::cons::density(n), k, j, i);
+                vmesh(b, gas::cons::momentum(VI(n, 0)), k, j, i) = vr * xcyl[0] / xv[0];
+                vmesh(b, gas::cons::momentum(VI(n, 1)), k, j, i) = vr * xcyl[2] / xv[0];
+                vmesh(b, gas::cons::total_energy(n), k, j, i) =
+                    etot_res +
+                    0.5 *
+                        (SQR(vmesh(b, gas::cons::momentum(VI(n, 0)), k, j, i)) +
+                         SQR(vmesh(b, gas::cons::momentum(VI(n, 1)), k, j, i))) /
+                        vmesh(b, gas::cons::density(n), k, j, i);
+              }
+            }
+
+            if (do_dust) {
+              for (int n = 0; n < vmesh.GetSize(b, dust::prim::density()); ++n) {
+                vmesh(b, dust::cons::density(n), k, j, i) = 0.0;
+                vmesh(b, dust::cons::momentum(VI(n, 0)), k, j, i) = 0.0;
+                vmesh(b, dust::cons::momentum(VI(n, 1)), k, j, i) = 0.0;
+                vmesh(b, dust::cons::momentum(VI(n, 2)), k, j, i) = 0.0;
+              }
+            }
+          } // end if (std::abs(xcyl[2]) > 5.0 * H)
+        });
+  }
+
+  return TaskStatus::complete;
+}
+
 //----------------------------------------------------------------------------------------
 //! \fn AmrTag ProblemCheckRefinementBlock()
 //! \brief Refinement criterion for disk pgen

src/pgen/problem_modifier.hpp

@@ -30,7 +30,8 @@
 namespace artemis {
 
 std::function<AmrTag(MeshBlockData<Real> *mbd)> ProblemCheckRefinementBlock = nullptr;
-
+std::function<TaskStatus(MeshData<Real> *md, const Real time, const Real dt)>
+    UserSourceTerm = nullptr;
 } // namespace artemis
 
 // Problem modifiers
@@ -66,6 +67,7 @@ void ProblemModifier(parthenon::ParthenonManager *pman) {
   } else if (artemis_problem == "disk") {
 
     artemis::ProblemCheckRefinementBlock = disk::ProblemCheckRefinementBlock;
+    artemis::UserSourceTerm = disk::UserSourceTerm<G>;
 
     pman->app_input->RegisterBoundaryCondition(BF::inner_x1, "ic",
                                                disk::DiskBoundaryIC<G, ID::inner_x1>);

src/utils/artemis_utils.hpp

@@ -35,6 +35,16 @@ KOKKOS_FORCEINLINE_FUNCTION Real VDot(const V1 &a, const V2 &b) {
   return a[0] * b[0] + a[1] * b[1] + a[2] * b[2];
 }
 
+//----------------------------------------------------------------------------------------
+//! \fn int ArtemisUtils::ProblemSourceTerm
+//! \brief Wrapper function for user-defined source term, checking for nullptr
+static TaskStatus ProblemSourceTerm(MeshData<Real> *md, const Real time, const Real dt) {
+  if (artemis::UserSourceTerm != nullptr) {
+    return artemis::UserSourceTerm(md, time, dt);
+  }
+  return TaskStatus::complete;
+}
+
 //----------------------------------------------------------------------------------------
 //! \fn Real ArtemisUtils::DualEnergySIE(vmesh, const int b, const int n, const int k,
 //!                                      const int j, const int i, const Real de_switch,