add P_theta flux register analogous to P_radial

Source/driver/Castro.H

@@ -1301,6 +1301,9 @@ protected:
 #if (AMREX_SPACEDIM <= 2)
     amrex::MultiFab         P_radial;
 #endif
+#if (AMREX_SPACEDIM == 2)
+    amrex::MultiFab         P_theta;
+#endif
 #ifdef RADIATION
     amrex::Vector<std::unique_ptr<amrex::MultiFab> > rad_fluxes;
 #endif

Source/driver/Castro.cpp

@@ -886,6 +886,12 @@ Castro::initMFs()
     }
 #endif
 
+#if (AMREX_SPACEDIM == 2)
+    if (Geom().IsSPHERICAL()) {
+      P_theta.define(getEdgeBoxArray(1), dmap, 1, 0);
+    }
+#endif
+
 #ifdef RADIATION
     if (Radiation::rad_hydro_combined) {
         rad_fluxes.resize(AMREX_SPACEDIM);
@@ -2590,6 +2596,12 @@ Castro::FluxRegCrseInit() {
     }
 #endif
 
+#if (AMREX_SPACEDIM == 2)
+    if (Geom().IsSPHERICAL()) {
+      fine_level.pres_reg.CrseInit(P_theta, 1, 0, 0, 1, pres_crse_scale);
+    }
+#endif
+
 #ifdef RADIATION
     if (Radiation::rad_hydro_combined) {
       for (int i = 0; i < AMREX_SPACEDIM; ++i) {
@@ -2620,6 +2632,12 @@ Castro::FluxRegFineAdd() {
     }
 #endif
 
+#if (AMREX_SPACEDIM == 2)
+    if (Geom().IsSPHERICAL()) {
+      getLevel(level).pres_reg.FineAdd(P_theta, 1, 0, 0, 1, pres_fine_scale);
+    }
+#endif
+
 #ifdef RADIATION
     if (Radiation::rad_hydro_combined) {
       for (int i = 0; i < AMREX_SPACEDIM; ++i) {
@@ -2889,6 +2907,55 @@ Castro::reflux (int crse_level, int fine_level, bool in_post_timestep)
         }
 #endif
 
+#if (AMREX_SPACEDIM == 2)
+        if (Geom().IsSPHERICAL()) {
+
+            reg = &getLevel(lev).pres_reg;
+
+            MultiFab rdtheta(crse_lev.grids, crse_lev.dmap, 1, 0);
+
+            const auto* problo = Geom().ProbLo();
+            auto dr = crse_lev.geom.CellSize(0);
+            auto dtheta = crse_lev.geom.CellSize(1);
+
+            auto const& ma = rdtheta.arrays();
+
+            amrex::ParallelFor(rdtheta,
+            [=] AMREX_GPU_DEVICE (int b, int i, int j, int k) noexcept
+            {
+                Real r = problo[0] + static_cast<Real>(i + 0.5_rt) * dr;
+                ma[b](i,j,k) = r * dtheta;
+            });
+
+            reg->ClearInternalBorders(crse_lev.geom);
+
+            reg->Reflux(crse_state, rdtheta, 1.0, 0, UMY, 1, crse_lev.geom);
+
+            if (update_sources_after_reflux || !in_post_timestep) {
+
+                MultiFab tmp_fluxes(crse_lev.P_theta.boxArray(),
+                                    crse_lev.P_theta.DistributionMap(),
+                                    crse_lev.P_theta.nComp(), crse_lev.P_theta.nGrow());
+
+                tmp_fluxes.setVal(0.0);
+
+                for (OrientationIter fi; fi.isValid(); ++fi)
+                {
+                    const FabSet& fs = (*reg)[fi()];
+                    if (fi().coordDir() == 1) {
+                        fs.copyTo(tmp_fluxes, 0, 0, 0, tmp_fluxes.nComp());
+                    }
+                }
+
+                MultiFab::Add(crse_lev.P_theta, tmp_fluxes, 0, 0, crse_lev.P_theta.nComp(), 0);
+
+            }
+
+            reg->setVal(0.0);
+
+        }
+#endif
+
 #ifdef RADIATION
 
         // This follows the same logic as the pure hydro fluxes; see above for details.

Source/driver/Castro_advance.cpp

@@ -551,6 +551,12 @@ Castro::initialize_advance(Real time, Real dt, int amr_iteration)
     }
 #endif
 
+#if (AMREX_SPACEDIM == 2)
+    if (Geom().IsSPHERICAL()) {
+        P_theta.setVal(0.0);
+    }
+#endif
+
 #ifdef RADIATION
     if (Radiation::rad_hydro_combined) {
         for (int dir = 0; dir < AMREX_SPACEDIM; ++dir) {

Source/driver/Castro_advance_ctu.cpp

@@ -204,6 +204,12 @@ Castro::retry_advance_ctu(Real dt, const advance_status& status)
             }
 #endif
 
+#if (AMREX_SPACEDIM == 2)
+            if (Geom().IsSPHERICAL()) {
+                P_theta.setVal(0.0);
+            }
+#endif
+
 #ifdef RADIATION
             if (Radiation::rad_hydro_combined) {
                 for (int dir = 0; dir < AMREX_SPACEDIM; ++dir) {

Source/hydro/Castro_ctu_hydro.cpp

@@ -179,6 +179,9 @@ Castro::construct_ctu_hydro_source(Real time, Real dt)  // NOLINT(readability-co
 #if AMREX_SPACEDIM <= 2
     FArrayBox pradial(The_Async_Arena());
 #endif
+#if AMREX_SPACEDIM == 2
+    FArrayBox ptheta(The_Async_Arena());
+#endif
 #if AMREX_SPACEDIM == 3
     FArrayBox qmyx(The_Async_Arena()), qpyx(The_Async_Arena());
     FArrayBox qmzx(The_Async_Arena()), qpzx(The_Async_Arena());
@@ -444,6 +447,13 @@ Castro::construct_ctu_hydro_source(Real time, Real dt)  // NOLINT(readability-co
       fab_size += pradial.nBytes();
 #endif
 
+#if AMREX_SPACEDIM == 2
+      if (Geom().IsSPHERICAL()) {
+          ptheta.resize(ybx, 1);
+      }
+      fab_size += ptheta.nBytes();
+#endif
+
 #ifdef SIMPLIFIED_SDC
 #ifdef REACTIONS
       Array4<Real> const sdc_src_arr = SDC_react_source.array(mfi);
@@ -1270,23 +1280,33 @@ Castro::construct_ctu_hydro_source(Real time, Real dt)  // NOLINT(readability-co
         scale_rad_flux(nbx, rad_flux_arr, area_arr, dt);
 #endif
 
-        if (idir == 0) {
 #if AMREX_SPACEDIM <= 2
+        // get the scaled radial pressure -- we need to treat this specially
+
+        if (idir == 0 && !mom_flux_has_p(0, 0, coord)) {
             Array4<Real> pradial_fab = pradial.array();
+
+            amrex::ParallelFor(nbx,
+            [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
+            {
+                pradial_fab(i,j,k) = qex_arr(i,j,k,GDPRES) * dt;
+            });
+        }
 #endif
 
-            // get the scaled radial pressure -- we need to treat this specially
-#if AMREX_SPACEDIM <= 2
-            if (!mom_flux_has_p(0, 0, coord)) {
-                amrex::ParallelFor(nbx,
-                [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
-                {
-                    pradial_fab(i,j,k) = qex_arr(i,j,k,GDPRES) * dt;
-                });
-            }
+#if AMREX_SPACEDIM == 2
+        // get the scaled pressure in the theta direction
 
-#endif
+        if (idir == 1 && !mom_flux_has_p(1, 1, coord)) {
+            Array4<Real> ptheta_fab = ptheta.array();
+
+            amrex::ParallelFor(nbx,
+            [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
+            {
+                ptheta_fab(i,j,k) = qey_arr(i,j,k,GDPRES) * dt;
+            });
         }
+#endif
 
         // Store the fluxes from this advance. For simplified SDC integration we
         // only need to do this on the last iteration.
@@ -1331,7 +1351,19 @@ Castro::construct_ctu_hydro_source(Real time, Real dt)  // NOLINT(readability-co
                     P_radial_fab(i,j,k,0) += pradial_fab(i,j,k,0);
                 });
             }
+#endif
 
+#if AMREX_SPACEDIM == 2
+            if (idir == 1 && !mom_flux_has_p(1, 1, coord)) {
+                Array4<Real> ptheta_fab = ptheta.array();
+                Array4<Real> P_theta_fab = P_theta.array(mfi);
+
+                amrex::ParallelFor(mfi.nodaltilebox(0),
+                [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
+                {
+                    P_theta_fab(i,j,k,0) += ptheta_fab(i,j,k,0);
+                });
+            }
 #endif
 
         } // add_fluxes

Source/hydro/Castro_mol_hydro.cpp

@@ -84,6 +84,9 @@ Castro::construct_mol_hydro_source(Real time, Real dt, MultiFab& A_update)
     }
 #if AMREX_SPACEDIM <= 2
     FArrayBox pradial(The_Async_Arena());
+#endif
+#if AMREX_SPACEDIM == 2
+    FArrayBox ptheta(The_Async_Arena());
 #endif
     FArrayBox avis(The_Async_Arena());
 
@@ -657,6 +660,15 @@ Castro::construct_mol_hydro_source(Real time, Real dt, MultiFab& A_update)
 
         Array4<Real> pradial_fab = pradial.array();
 #endif
+
+#if AMREX_SPACEDIM == 2
+        if (Geom().IsSPHERICAL()) {
+          ptheta.resize(ybx, 1);
+        }
+
+        Array4<Real> ptheta_fab = ptheta.array();
+#endif
+
 #if AMREX_SPACEDIM == 1
         Array4<Real> const qex_arr = qe[0].array();
 #endif
@@ -674,23 +686,34 @@ Castro::construct_mol_hydro_source(Real time, Real dt, MultiFab& A_update)
 #endif
                      flux_arr, area_arr, dt);
 
-
-          if (idir == 0) {
+#if AMREX_SPACEDIM <= 2
+          if (idir == 0 && !mom_flux_has_p(0, 0, coord)) {
             // get the scaled radial pressure -- we need to treat this specially
-            Array4<Real> const qex_fab = qe[idir].array();
-            const int prescomp = GDPRES;
 
+            Array4<Real> const qex_fab = qe[idir].array();
 
-#if AMREX_SPACEDIM <= 2
-            if (!mom_flux_has_p(0, 0, coord)) {
-              amrex::ParallelFor(nbx,
-              [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
-              {
-                pradial_fab(i,j,k) = qex_fab(i,j,k,prescomp) * dt;
-              });
-            }
+            amrex::ParallelFor(nbx,
+            [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
+            {
+              pradial_fab(i,j,k) = qex_fab(i,j,k,GDPRES) * dt;
+            });
+          }
 #endif
+
+#if AMREX_SPACEDIM == 2
+          if (idir == 1 && !mom_flux_has_p(1, 1, coord)) {
+            // get the scaled pressure in the theta direction
+
+            Array4<Real> const qey_fab = qe[idir].array();
+
+            amrex::ParallelFor(nbx,
+            [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
+            {
+              ptheta_fab(i,j,k) = qey_fab(i,j,k,GDPRES) * dt;
+            });
           }
+#endif
+
         }
 
 
@@ -729,6 +752,20 @@ Castro::construct_mol_hydro_source(Real time, Real dt, MultiFab& A_update)
 
           }
 #endif
+
+#if AMREX_SPACEDIM == 2
+          if (Geom().IsSPHERICAL()) {
+
+            Array4<Real> P_theta_fab = P_theta.array(mfi);
+            const Real scale = stage_weight;
+
+            AMREX_HOST_DEVICE_FOR_4D(mfi.nodaltilebox(0), 1, i, j, k, n,
+            {
+                P_theta_fab(i,j,k,0) += scale * ptheta_fab(i,j,k,0);
+            });
+
+          }
+#endif
         }
 
       } // MFIter loop