ERF
Energy Research and Forecasting: An Atmospheric Modeling Code
Functions
ERF_Substep_NS.cpp File Reference
#include <ERF_TI_fast_headers.H>
Include dependency graph for ERF_Substep_NS.cpp:

Functions

void erf_substep_NS (int step, int nrk, int level, int finest_level, Vector< MultiFab > &S_slow_rhs, const Vector< MultiFab > &S_prev, Vector< MultiFab > &S_stage_data, const MultiFab &S_stage_prim, const MultiFab &qt, const MultiFab &pi_stage, const MultiFab &fast_coeffs, Vector< MultiFab > &S_data, MultiFab &lagged_delta_rt, MultiFab &avg_xmom, MultiFab &avg_ymom, MultiFab &avg_zmom, const MultiFab &cc_src, const MultiFab &xmom_src, const MultiFab &ymom_src, const MultiFab &zmom_src, const Geometry geom, const Real gravity, amrex::Gpu::DeviceVector< amrex::Real > &stretched_dz_d, const double dtau_d, const Real beta_s, const Real facinv, Vector< std::unique_ptr< MultiFab >> &mapfac, YAFluxRegister *fr_as_crse, YAFluxRegister *fr_as_fine, bool l_use_moisture, bool l_reflux, bool l_real_bc, const amrex::Real *sinesq_stag_d, const Real l_damp_coef)
 

Function Documentation

◆ erf_substep_NS()

void erf_substep_NS ( int  step,
int  nrk,
int  level,
int  finest_level,
Vector< MultiFab > &  S_slow_rhs,
const Vector< MultiFab > &  S_prev,
Vector< MultiFab > &  S_stage_data,
const MultiFab &  S_stage_prim,
const MultiFab &  qt,
const MultiFab &  pi_stage,
const MultiFab &  fast_coeffs,
Vector< MultiFab > &  S_data,
MultiFab &  lagged_delta_rt,
MultiFab &  avg_xmom,
MultiFab &  avg_ymom,
MultiFab &  avg_zmom,
const MultiFab &  cc_src,
const MultiFab &  xmom_src,
const MultiFab &  ymom_src,
const MultiFab &  zmom_src,
const Geometry  geom,
const Real  gravity,
amrex::Gpu::DeviceVector< amrex::Real > &  stretched_dz_d,
const double  dtau_d,
const Real  beta_s,
const Real  facinv,
Vector< std::unique_ptr< MultiFab >> &  mapfac,
YAFluxRegister *  fr_as_crse,
YAFluxRegister *  fr_as_fine,
bool  l_use_moisture,
bool  l_reflux,
bool  l_real_bc,
const amrex::Real *  sinesq_stag_d,
const Real  l_damp_coef 
)

Function for computing the fast RHS with no terrain and variable vertical spacing

Parameters
[in]stepwhich fast time step within each Runge-Kutta step
[in]nrkwhich Runge-Kutta step
[in]levellevel of resolution
[in]finest_levelfinest level of resolution
[in]S_slow_rhsslow RHS computed in erf_slow_rhs_pre
[in]S_previf step == 0, this is S_old, else the previous fast solution
[in]S_stage_datasolution at previous RK stage
[in]S_stage_primprimitive variables at previous RK stage
[in]pi_stageExner function at previous RK stage
[in]fast_coeffscoefficients for the tridiagonal solve used in the fast integrator
[out]S_datacurrent solution
[in,out]lagged_delta_rt
[in,out]avg_xmomtime-averaged x-momentum to be used for updating slow variables
[in,out]avg_ymomtime-averaged y-momentum to be used for updating slow variables
[in,out]avg_zmomtime-averaged z-momentum to be used for updating slow variables
[in]cc_srcsource terms for conserved variables
[in]xmom_srcsource terms for x-momentum
[in]ymom_srcsource terms for y-momentum
[in]zmom_srcsource terms for z-momentum
[in]geomcontainer for geometric information
[in]gravitymagnitude of gravity
[in]stretched_dz_d
[in]dtaufast time step
[in]beta_sCoefficient which determines how implicit vs explicit the solve is
[in]facinvinverse factor for time-averaging the momenta
[in]mapfacvector of map factors
[in,out]fr_as_crseYAFluxRegister at level l at level l / l+1 interface
[in,out]fr_as_fineYAFluxRegister at level l at level l-1 / l interface
[in]l_use_moisture
[in]l_refluxshould we add fluxes to the FluxRegisters?
73 {
74  //
75  // NOTE: for step > 0, S_data and S_prev point to the same MultiFab data!!
76  //
77 
78  BL_PROFILE_REGION("erf_substep_NS()");
79 
80  Real dtau = static_cast<Real>(dtau_d);
81 
82  const Box& domain = geom.Domain();
83  auto const domlo = lbound(domain);
84  auto const domhi = ubound(domain);
85 
86  int ilo = domlo.x;
87  int ihi = domhi.x + 1;
88  int jlo = domlo.y;
89  int jhi = domhi.y + 1;
90 
91  Real beta_1 = myhalf * (one - beta_s); // multiplies explicit terms
92  Real beta_2 = myhalf * (one + beta_s); // multiplies implicit terms
93 
94  // How much do we project forward the (rho theta) that is used in the horizontal momentum equations
95  Real beta_d = Real(0.1);
96 
97  bool l_rayleigh_impl_for_w = (sinesq_stag_d != nullptr);
98 
99  const Real* dx = geom.CellSize();
100  const GpuArray<Real, AMREX_SPACEDIM> dxInv = geom.InvCellSizeArray();
101 
102  Real dxi = dxInv[0];
103  Real dyi = dxInv[1];
104  Real dzi = dxInv[2];
105 
106  auto dz_ptr = stretched_dz_d.data();
107 
108  const auto& ba = S_stage_data[IntVars::cons].boxArray();
109  const auto& dm = S_stage_data[IntVars::cons].DistributionMap();
110 
111  MultiFab Delta_rho_theta( ba , dm, 1, 1);
112  MultiFab Delta_rho_w (convert(ba,IntVect(0,0,1)), dm, 1, IntVect(1,1,0));
113 
114  MultiFab coeff_A_mf(fast_coeffs, make_alias, 0, 1);
115  MultiFab inv_coeff_B_mf(fast_coeffs, make_alias, 1, 1);
116  MultiFab coeff_C_mf(fast_coeffs, make_alias, 2, 1);
117  MultiFab coeff_P_mf(fast_coeffs, make_alias, 3, 1);
118  MultiFab coeff_Q_mf(fast_coeffs, make_alias, 4, 1);
119 
120  // *************************************************************************
121  // Set gravity as a vector
122  const Array<Real,AMREX_SPACEDIM> grav{zero, zero, -gravity};
123  const GpuArray<Real,AMREX_SPACEDIM> grav_gpu{grav[0], grav[1], grav[2]};
124 
125  // This will hold theta extrapolated forward in time
126  MultiFab extrap(S_data[IntVars::cons].boxArray(),S_data[IntVars::cons].DistributionMap(),1,1);
127 
128  // This will hold the update for (rho) and (rho theta)
129  MultiFab temp_rhs(S_stage_data[IntVars::zmom].boxArray(),S_stage_data[IntVars::zmom].DistributionMap(),2,0);
130 
131  // This will hold the new x- and y-momenta temporarily (so that we don't overwrite values we need when tiling)
132  MultiFab temp_cur_xmom(S_stage_data[IntVars::xmom].boxArray(),S_stage_data[IntVars::xmom].DistributionMap(),1,0);
133  MultiFab temp_cur_ymom(S_stage_data[IntVars::ymom].boxArray(),S_stage_data[IntVars::ymom].DistributionMap(),1,0);
134 
135  // We assume that in the first step (nrk == 0) we are only doing one substep.
136  AMREX_ALWAYS_ASSERT(nrk > 0 || step == 0);
137 
138  // *************************************************************************
139  // First set up some arrays we'll need
140  // *************************************************************************
141 
142 #ifdef _OPENMP
143 #pragma omp parallel if (Gpu::notInLaunchRegion())
144 #endif
145  for ( MFIter mfi(S_stage_data[IntVars::cons],TilingIfNotGPU()); mfi.isValid(); ++mfi)
146  {
147  const Array4<const Real>& prev_cons = S_prev[IntVars::cons].const_array(mfi);
148  const Array4<const Real>& prev_zmom = S_prev[IntVars::zmom].const_array(mfi);
149 
150  const Array4<const Real>& stage_cons = S_stage_data[IntVars::cons].const_array(mfi);
151  const Array4<const Real>& stage_zmom = S_stage_data[IntVars::zmom].const_array(mfi);
152 
153  const Array4<Real>& prev_drho_w = Delta_rho_w.array(mfi);
154  const Array4<Real>& prev_drho_theta = Delta_rho_theta.array(mfi);
155  const Array4<Real>& lagged_arr = lagged_delta_rt.array(mfi);
156  const Array4<Real>& theta_extrap = extrap.array(mfi);
157  const Array4<const Real>& prim = S_stage_prim.const_array(mfi);
158 
159  Box gbx = mfi.growntilebox(1);
160  ParallelFor(gbx, [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
161  {
162  prev_drho_theta(i,j,k) = prev_cons(i,j,k,RhoTheta_comp) - stage_cons(i,j,k,RhoTheta_comp);
163 
164  if (step == 0) {
165  theta_extrap(i,j,k) = prev_drho_theta(i,j,k);
166  } else {
167  theta_extrap(i,j,k) = prev_drho_theta(i,j,k) + beta_d *
168  ( prev_drho_theta(i,j,k) - lagged_arr(i,j,k) );
169  }
170 
171  // NOTE: qv is not changing over the fast steps so we use the stage data
172  Real qv = (l_use_moisture) ? prim(i,j,k,PrimQ1_comp) : zero;
173  theta_extrap(i,j,k) *= (one + RvoRd*qv);
174 
175  // We define lagged_delta_rt for our next step as the current delta_rt
176  // (after using it above to extrapolate theta for this step)
177  lagged_arr(i,j,k) = prev_drho_theta(i,j,k);
178  });
179 
180  // NOTE: We must do this here because for step > 0, prev_zmom and cur_zmom both point to the same data,
181  // so by the time we would use prev_zmom to define zflux, it would have already been over-written.
182  Box gtbz = mfi.nodaltilebox(2);
183  gtbz.grow(IntVect(1,1,0));
184  ParallelFor(gtbz, [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept {
185  prev_drho_w(i,j,k) = prev_zmom(i,j,k) - stage_zmom(i,j,k);
186  });
187  } // mfi
188 
189  // *************************************************************************
190  // Define updates in the current RK stage
191  // *************************************************************************
192 
193 #ifdef _OPENMP
194 #pragma omp parallel if (Gpu::notInLaunchRegion())
195 #endif
196  for ( MFIter mfi(S_stage_data[IntVars::cons],TilingIfNotGPU()); mfi.isValid(); ++mfi)
197  {
198  Box tbx = mfi.nodaltilebox(0);
199  Box tby = mfi.nodaltilebox(1);
200 
201  const Array4<Real const>& xmom_src_arr = xmom_src.const_array(mfi);
202  const Array4<Real const>& ymom_src_arr = ymom_src.const_array(mfi);
203 
204  const Array4<const Real> & stage_xmom = S_stage_data[IntVars::xmom].const_array(mfi);
205  const Array4<const Real> & stage_ymom = S_stage_data[IntVars::ymom].const_array(mfi);
206  const Array4<const Real> & qt_arr = qt.const_array(mfi);
207 
208  const Array4<const Real>& slow_rhs_rho_u = S_slow_rhs[IntVars::xmom].const_array(mfi);
209  const Array4<const Real>& slow_rhs_rho_v = S_slow_rhs[IntVars::ymom].const_array(mfi);
210 
211  const Array4<Real>& temp_cur_xmom_arr = temp_cur_xmom.array(mfi);
212  const Array4<Real>& temp_cur_ymom_arr = temp_cur_ymom.array(mfi);
213 
214  const Array4<const Real>& prev_xmom = S_prev[IntVars::xmom].const_array(mfi);
215  const Array4<const Real>& prev_ymom = S_prev[IntVars::ymom].const_array(mfi);
216 
217  // These store the advection momenta which we will use to update the slow variables
218  const Array4< Real>& avg_xmom_arr = avg_xmom.array(mfi);
219  const Array4< Real>& avg_ymom_arr = avg_ymom.array(mfi);
220 
221  const Array4<const Real>& pi_stage_ca = pi_stage.const_array(mfi);
222 
223  const Array4<Real>& theta_extrap = extrap.array(mfi);
224 
225  // Map factors
226  const Array4<const Real>& mf_ux = mapfac[MapFacType::u_x]->const_array(mfi);
227  const Array4<const Real>& mf_uy = mapfac[MapFacType::u_y]->const_array(mfi);
228  const Array4<const Real>& mf_vx = mapfac[MapFacType::v_x]->const_array(mfi);
229  const Array4<const Real>& mf_vy = mapfac[MapFacType::v_y]->const_array(mfi);
230 
231  // *********************************************************************
232  // Define updates in the RHS of {x, y, z}-momentum equations
233  //
234  // NOTE: avg_{x,y}mom is a *flux* that is later handed to the scalar
235  // advection in erf_slow_rhs_post, which applies (mfsq/detJ) to its
236  // divergence. So what we accumulate here must be in the same
237  // convention as the base value set in AdvectionSrcForRho, namely
238  // ax*rho_u/mf_uy (and ay*rho_v/mf_vx). Two weightings are needed:
239  //
240  // (1) the map factor, matching the density fluxes formed below;
241  //
242  // (2) h_zeta = ax = detJ, which for this (laterally homogeneous)
243  // mesh is just dz_ptr[k]/dz. Note this factor cancels out of
244  // the rho update we do here -- temp_rhs uses a bare dxi and no
245  // 1/detJ -- but it does NOT cancel for avg_{x,y}mom, because
246  // the consumer downstream re-applies the 1/detJ that we never
247  // applied. It is identically 1 for MeshType::ConstantDz.
248  // *********************************************************************
249  if (nrk == 0 and step == 0) { // prev == stage
250  ParallelFor(tbx, tby,
251  [=] AMREX_GPU_DEVICE (int i, int j, int k)
252  {
253  Real h_zeta = dz_ptr[k] * dzi;
254  Real new_drho_u = dtau * slow_rhs_rho_u(i,j,k) + dtau * xmom_src_arr(i,j,k);;
255  avg_xmom_arr(i,j,k) += facinv * new_drho_u * h_zeta / mf_uy(i,j,0);
256  temp_cur_xmom_arr(i,j,k) = stage_xmom(i,j,k) + new_drho_u;
257  },
258  [=] AMREX_GPU_DEVICE (int i, int j, int k)
259  {
260  Real h_zeta = dz_ptr[k] * dzi;
261  Real new_drho_v = dtau * slow_rhs_rho_v(i,j,k) + dtau * ymom_src_arr(i,j,k);
262  avg_ymom_arr(i,j,k) += facinv * new_drho_v * h_zeta / mf_vx(i,j,0);
263  temp_cur_ymom_arr(i,j,k) = stage_ymom(i,j,k) + new_drho_v;
264  });
265  } else {
266  ParallelFor(tbx, tby,
267  [=] AMREX_GPU_DEVICE (int i, int j, int k)
268  {
269  // Add (negative) gradient of (rho theta) multiplied by lagged "pi"
270  Real gpx = (l_real_bc && (level==0) && (i==ilo || i==ihi)) ? Real(0.) :
271  (theta_extrap(i,j,k) - theta_extrap(i-1,j,k))*dxi;
272  gpx *= mf_ux(i,j,0);
273 
274  Real q = (l_use_moisture) ? myhalf * (qt_arr(i,j,k) + qt_arr(i-1,j,k)) : zero;
275 
276  Real pi_c = myhalf * (pi_stage_ca(i-1,j,k,0) + pi_stage_ca(i,j,k,0));
277  Real fast_rhs_rho_u = -Gamma * R_d * pi_c * gpx / (one + q);
278 
279  Real new_drho_u = prev_xmom(i,j,k) - stage_xmom(i,j,k)
280  + dtau * fast_rhs_rho_u + dtau * slow_rhs_rho_u(i,j,k)
281  + dtau * xmom_src_arr(i,j,k);
282 
283  Real h_zeta = dz_ptr[k] * dzi;
284  avg_xmom_arr(i,j,k) += facinv * new_drho_u * h_zeta / mf_uy(i,j,0);
285 
286  temp_cur_xmom_arr(i,j,k) = stage_xmom(i,j,k) + new_drho_u;
287  },
288  [=] AMREX_GPU_DEVICE (int i, int j, int k)
289  {
290  // Add (negative) gradient of (rho theta) multiplied by lagged "pi"
291  Real gpy = (l_real_bc && (level==0) && (j==jlo || j==jhi)) ? Real(0.) :
292  (theta_extrap(i,j,k) - theta_extrap(i,j-1,k))*dyi;
293  gpy *= mf_vy(i,j,0);
294 
295  Real q = (l_use_moisture) ? myhalf * (qt_arr(i,j,k) + qt_arr(i,j-1,k)) : zero;
296 
297  Real pi_c = myhalf * (pi_stage_ca(i,j-1,k,0) + pi_stage_ca(i,j,k,0));
298  Real fast_rhs_rho_v = -Gamma * R_d * pi_c * gpy / (one + q);
299 
300  Real new_drho_v = prev_ymom(i,j,k) - stage_ymom(i,j,k)
301  + dtau * fast_rhs_rho_v + dtau * slow_rhs_rho_v(i,j,k)
302  + dtau * ymom_src_arr(i,j,k);
303 
304  Real h_zeta = dz_ptr[k] * dzi;
305  avg_ymom_arr(i,j,k) += facinv * new_drho_v * h_zeta / mf_vx(i,j,0);
306 
307  temp_cur_ymom_arr(i,j,k) = stage_ymom(i,j,k) + new_drho_v;
308  });
309  } // nrk > 0 and/or step > 0
310  } //mfi
311 
312 #ifdef _OPENMP
313 #pragma omp parallel if (Gpu::notInLaunchRegion())
314 #endif
315  {
316  std::array<FArrayBox,AMREX_SPACEDIM> flux;
317  for ( MFIter mfi(S_stage_data[IntVars::cons],TileNoZ()); mfi.isValid(); ++mfi)
318  {
319  Box bx = mfi.tilebox();
320  Box tbz = surroundingNodes(bx,2);
321 
322  Box vbx = mfi.validbox();
323  const auto& vbx_hi = ubound(vbx);
324 
325  const Array4<Real const>& zmom_src_arr = zmom_src.const_array(mfi);
326 
327  const Array4<const Real>& stage_xmom = S_stage_data[IntVars::xmom].const_array(mfi);
328  const Array4<const Real>& stage_ymom = S_stage_data[IntVars::ymom].const_array(mfi);
329  const Array4<const Real>& stage_zmom = S_stage_data[IntVars::zmom].const_array(mfi);
330  const Array4<const Real> & prim = S_stage_prim.const_array(mfi);
331 
332  const Array4<const Real>& prev_drho_theta = Delta_rho_theta.array(mfi);
333 
334  const Array4<const Real>& prev_cons = S_prev[IntVars::cons].const_array(mfi);
335  const Array4<const Real>& stage_cons = S_stage_data[IntVars::cons].const_array(mfi);
336 
337  const Array4<const Real>& slow_rhs_cons = S_slow_rhs[IntVars::cons].const_array(mfi);
338  const Array4<const Real>& slow_rhs_rho_w = S_slow_rhs[IntVars::zmom].const_array(mfi);
339 
340  const Array4<const Real>& prev_zmom = S_prev[IntVars::zmom].const_array(mfi);
341  const Array4< Real>& cur_zmom = S_data[IntVars::zmom].array(mfi);
342 
343  const Array4<Real>& temp_cur_xmom_arr = temp_cur_xmom.array(mfi);
344  const Array4<Real>& temp_cur_ymom_arr = temp_cur_ymom.array(mfi);
345 
346  // These store the advection momenta which we will use to update the slow variables
347  const Array4< Real>& avg_zmom_arr = avg_zmom.array(mfi);
348 
349  // Map factors
350  const Array4<const Real>& mf_mx = mapfac[MapFacType::m_x]->const_array(mfi);
351  const Array4<const Real>& mf_my = mapfac[MapFacType::m_y]->const_array(mfi);
352  const Array4<const Real>& mf_uy = mapfac[MapFacType::u_y]->const_array(mfi);
353  const Array4<const Real>& mf_vx = mapfac[MapFacType::v_x]->const_array(mfi);
354 
355  FArrayBox RHS_fab;
356  RHS_fab.resize(tbz,1, The_Async_Arena());
357 
358  FArrayBox soln_fab;
359  soln_fab.resize(tbz,1, The_Async_Arena());
360 
361  auto const& RHS_a = RHS_fab.array();
362  auto const& soln_a = soln_fab.array();
363 
364  auto const& temp_rhs_arr = temp_rhs.array(mfi);
365 
366  auto const& coeffA_a = coeff_A_mf.array(mfi);
367  auto const& inv_coeffB_a = inv_coeff_B_mf.array(mfi);
368  auto const& coeffC_a = coeff_C_mf.array(mfi);
369  auto const& coeffP_a = coeff_P_mf.array(mfi);
370  auto const& coeffQ_a = coeff_Q_mf.array(mfi);
371 
372  // *************************************************************************
373  // Define flux arrays for use in advection
374  // *************************************************************************
375  for (int dir = 0; dir < AMREX_SPACEDIM; ++dir) {
376  flux[dir].resize(surroundingNodes(bx,dir),2,The_Async_Arena());
377  flux[dir].setVal<RunOn::Device>(0);
378  }
379  const GpuArray<const Array4<Real>, AMREX_SPACEDIM>
380  flx_arr{{AMREX_D_DECL(flux[0].array(), flux[1].array(), flux[2].array())}};
381 
382  // *********************************************************************
383  ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept {
384  Real xflux_lo = (temp_cur_xmom_arr(i ,j,k) - stage_xmom(i ,j,k)) / mf_uy(i ,j,0);
385  Real xflux_hi = (temp_cur_xmom_arr(i+1,j,k) - stage_xmom(i+1,j,k)) / mf_uy(i+1,j,0);
386  Real yflux_lo = (temp_cur_ymom_arr(i,j ,k) - stage_ymom(i,j ,k)) / mf_vx(i,j ,0);
387  Real yflux_hi = (temp_cur_ymom_arr(i,j+1,k) - stage_ymom(i,j+1,k)) / mf_vx(i,j+1,0);
388 
389  Real mfsq = mf_mx(i,j,0) * mf_my(i,j,0);
390 
391  temp_rhs_arr(i,j,k,Rho_comp ) = ( xflux_hi - xflux_lo ) * dxi * mfsq
392  + ( yflux_hi - yflux_lo ) * dyi * mfsq;
393  temp_rhs_arr(i,j,k,RhoTheta_comp) = (( xflux_hi * (prim(i,j,k,0) + prim(i+1,j,k,0)) -
394  xflux_lo * (prim(i,j,k,0) + prim(i-1,j,k,0)) ) * dxi * mfsq +
395  ( yflux_hi * (prim(i,j,k,0) + prim(i,j+1,k,0)) -
396  yflux_lo * (prim(i,j,k,0) + prim(i,j-1,k,0)) ) * dyi * mfsq) * myhalf;
397 
398  if (l_reflux) {
399  (flx_arr[0])(i,j,k,0) = xflux_lo;
400  (flx_arr[0])(i,j,k,1) = (flx_arr[0])(i ,j,k,0) * myhalf * (prim(i,j,k,0) + prim(i-1,j,k,0));
401 
402  (flx_arr[1])(i,j,k,0) = yflux_lo;
403  (flx_arr[1])(i,j,k,1) = (flx_arr[1])(i,j ,k,0) * myhalf * (prim(i,j,k,0) + prim(i,j-1,k,0));
404 
405  if (i == vbx_hi.x) {
406  (flx_arr[0])(i+1,j,k,0) = xflux_hi;
407  (flx_arr[0])(i+1,j,k,1) = (flx_arr[0])(i+1,j,k,0) * myhalf * (prim(i,j,k,0) + prim(i+1,j,k,0));
408  }
409  if (j == vbx_hi.y) {
410  (flx_arr[1])(i,j+1,k,0) = yflux_hi;
411  (flx_arr[1])(i,j+1,k,1) = (flx_arr[1])(i,j+1,k,0) * myhalf * (prim(i,j,k,0) + prim(i,j+1,k,0));
412  }
413  }
414  });
415 
416  Box bx_shrunk_in_k = bx;
417  int klo = tbz.smallEnd(2);
418  int khi = tbz.bigEnd(2);
419  bx_shrunk_in_k.setSmall(2,klo+1);
420  bx_shrunk_in_k.setBig(2,khi-1);
421 
422  // Note that the notes use "g" to mean the magnitude of gravity, so it is positive
423  // We set grav_gpu[2] to be the vector component which is negative
424  // We define halfg to match the notes (which is why we take the absolute value)
425  Real halfg = std::abs(myhalf * grav_gpu[2]);
426 
427  // *********************************************************************
428  // fast_loop_on_shrunk
429  // *********************************************************************
430  //Note we don't act on the bottom or top boundaries of the domain
431  ParallelFor(bx_shrunk_in_k, [=] AMREX_GPU_DEVICE (int i, int j, int k)
432  {
433  Real coeff_P = coeffP_a(i,j,k);
434  Real coeff_Q = coeffQ_a(i,j,k);
435 
436  Real theta_t_lo = myhalf * ( prim(i,j,k-2,PrimTheta_comp) + prim(i,j,k-1,PrimTheta_comp) );
437  Real theta_t_mid = myhalf * ( prim(i,j,k-1,PrimTheta_comp) + prim(i,j,k ,PrimTheta_comp) );
438  Real theta_t_hi = myhalf * ( prim(i,j,k ,PrimTheta_comp) + prim(i,j,k+1,PrimTheta_comp) );
439 
440  Real Omega_kp1 = prev_zmom(i,j,k+1) - stage_zmom(i,j,k+1);
441  Real Omega_k = prev_zmom(i,j,k ) - stage_zmom(i,j,k );
442  Real Omega_km1 = prev_zmom(i,j,k-1) - stage_zmom(i,j,k-1);
443 
444  // line 2 last two terms (order dtau)
445  Real old_drho_k = prev_cons(i,j,k ,Rho_comp) - stage_cons(i,j,k ,Rho_comp);
446  Real old_drho_km1 = prev_cons(i,j,k-1,Rho_comp) - stage_cons(i,j,k-1,Rho_comp);
447  Real R0_tmp = coeff_P * prev_drho_theta(i,j,k) + coeff_Q * prev_drho_theta(i,j,k-1)
448  - halfg * ( old_drho_k + old_drho_km1 );
449 
450  // lines 3-5 residuals (order dtau^2) one <-> beta_2
451  Real R1_tmp = halfg * (-slow_rhs_cons(i,j,k ,Rho_comp) - slow_rhs_cons(i,j,k-1,Rho_comp)
452  + temp_rhs_arr(i,j,k ,Rho_comp) + temp_rhs_arr(i,j,k-1,Rho_comp) )
453  + ( coeff_P * (slow_rhs_cons(i,j,k ,RhoTheta_comp) - temp_rhs_arr(i,j,k ,RhoTheta_comp)) +
454  coeff_Q * (slow_rhs_cons(i,j,k-1,RhoTheta_comp) - temp_rhs_arr(i,j,k-1,RhoTheta_comp)) );
455 
456  // lines 6&7 consolidated (reuse Omega & metrics) (order dtau^2)
457  R1_tmp += beta_1 * ( ( (Omega_kp1 - Omega_k) / dz_ptr[k] + (Omega_k - Omega_km1) / dz_ptr[k-1] ) * halfg
458  +(-(Omega_kp1*theta_t_hi - Omega_k *theta_t_mid) * coeff_P / dz_ptr[k]
459  -(Omega_k *theta_t_mid - Omega_km1*theta_t_lo ) * coeff_Q / dz_ptr[k-1]) );
460 
461  // line 1
462  RHS_a(i,j,k) = Omega_k + dtau * (slow_rhs_rho_w(i,j,k) + R0_tmp + dtau * beta_2 * R1_tmp + zmom_src_arr(i,j,k));
463 
464  }); // bx_shrunk_in_k
465 
466  Box b2d = tbz; // Copy constructor
467  b2d.setRange(2,0);
468 
469  auto const lo = lbound(bx);
470  auto const hi = ubound(bx);
471 
472  ParallelFor(b2d, [=] AMREX_GPU_DEVICE (int i, int j, int)
473  {
474  // w at bottom boundary of grid is 0 if at domain boundary, otherwise w = w_old + dtau * slow_rhs
475  RHS_a (i,j,lo.z) = prev_zmom(i,j,lo.z) - stage_zmom(i,j,lo.z)
476  + dtau * slow_rhs_rho_w(i,j,lo.z)
477  + dtau * zmom_src_arr(i,j,lo.z);
478 
479  // w at top boundary of grid is 0 if at domain boundary, otherwise w = w_old + dtau * slow_rhs
480  RHS_a (i,j,hi.z+1) = prev_zmom(i,j,hi.z+1) - stage_zmom(i,j,hi.z+1)
481  + dtau * slow_rhs_rho_w(i,j,hi.z+1)
482  + dtau * zmom_src_arr(i,j,hi.z+1);
483  }); // b2d
484 
485 #ifdef AMREX_USE_GPU
486  ParallelFor(b2d, [=] AMREX_GPU_DEVICE (int i, int j, int)
487  {
488  // w = specified Dirichlet value at k = lo.z
489  soln_a(i,j,lo.z) = RHS_a(i,j,lo.z) * inv_coeffB_a(i,j,lo.z);
490  cur_zmom(i,j,lo.z) = stage_zmom(i,j,lo.z) + soln_a(i,j,lo.z);
491 
492  for (int k = lo.z+1; k <= hi.z+1; k++) {
493  soln_a(i,j,k) = (RHS_a(i,j,k)-coeffA_a(i,j,k)*soln_a(i,j,k-1)) * inv_coeffB_a(i,j,k);
494  }
495 
496  cur_zmom(i,j,hi.z+1) = stage_zmom(i,j,hi.z+1) + soln_a(i,j,hi.z+1);
497 
498  for (int k = hi.z; k >= lo.z; k--) {
499  soln_a(i,j,k) -= ( coeffC_a(i,j,k) * inv_coeffB_a(i,j,k) ) *soln_a(i,j,k+1);
500  cur_zmom(i,j,k) = stage_zmom(i,j,k) + soln_a(i,j,k);
501  }
502  }); // b2d
503 #else
504  for (int j = lo.y; j <= hi.y; ++j) {
505  AMREX_PRAGMA_SIMD
506  for (int i = lo.x; i <= hi.x; ++i) {
507  soln_a(i,j,lo.z) = RHS_a(i,j,lo.z) * inv_coeffB_a(i,j,lo.z);
508  }
509  }
510  for (int k = lo.z+1; k <= hi.z+1; ++k) {
511  for (int j = lo.y; j <= hi.y; ++j) {
512  AMREX_PRAGMA_SIMD
513  for (int i = lo.x; i <= hi.x; ++i) {
514  soln_a(i,j,k) = (RHS_a(i,j,k)-coeffA_a(i,j,k)*soln_a(i,j,k-1)) * inv_coeffB_a(i,j,k);
515  }
516  }
517  }
518  for (int j = lo.y; j <= hi.y; ++j) {
519  AMREX_PRAGMA_SIMD
520  for (int i = lo.x; i <= hi.x; ++i) {
521  cur_zmom(i,j,hi.z+1) = stage_zmom(i,j,hi.z+1) + soln_a(i,j,hi.z+1);
522  }
523  }
524  for (int k = hi.z; k >= lo.z; --k) {
525  for (int j = lo.y; j <= hi.y; ++j) {
526  AMREX_PRAGMA_SIMD
527  for (int i = lo.x; i <= hi.x; ++i) {
528  soln_a(i,j,k) -= ( coeffC_a(i,j,k) * inv_coeffB_a(i,j,k) ) * soln_a(i,j,k+1);
529  cur_zmom(i,j,k) = stage_zmom(i,j,k) + soln_a(i,j,k);
530  }
531  }
532  }
533 #endif
534  if (l_rayleigh_impl_for_w) {
535  ParallelFor(bx_shrunk_in_k, [=] AMREX_GPU_DEVICE (int i, int j, int k)
536  {
537  Real damping_coeff = l_damp_coef * dtau * sinesq_stag_d[k];
538  cur_zmom(i,j,k) /= (one + damping_coeff);
539  });
540  }
541 
542  // **************************************************************************
543  // Define updates in the RHS of rho and (rho theta)
544  // **************************************************************************
545  const Array4<Real>& prev_drho_w = Delta_rho_w.array(mfi);
546  ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
547  {
548  Real zflux_lo = beta_2 * soln_a(i,j,k ) + beta_1 * prev_drho_w(i,j,k );
549  Real zflux_hi = beta_2 * soln_a(i,j,k+1) + beta_1 * prev_drho_w(i,j,k+1);
550 
551  avg_zmom_arr(i,j,k) += facinv*zflux_lo / (mf_mx(i,j,0) * mf_my(i,j,0));
552  if (l_reflux) {
553  (flx_arr[2])(i,j,k,0) = zflux_lo / (mf_mx(i,j,0) * mf_my(i,j,0));
554  (flx_arr[2])(i,j,k,1) = (flx_arr[2])(i,j,k,0) * myhalf * (prim(i,j,k) + prim(i,j,k-1));
555  }
556 
557  if (k == vbx_hi.z) {
558  avg_zmom_arr(i,j,k+1) += facinv * zflux_hi / (mf_mx(i,j,0) * mf_my(i,j,0));
559  if (l_reflux) {
560  (flx_arr[2])(i,j,k+1,0) = zflux_hi / (mf_mx(i,j,0) * mf_my(i,j,0));
561  (flx_arr[2])(i,j,k+1,1) = (flx_arr[2])(i,j,k+1,0) * myhalf * (prim(i,j,k) + prim(i,j,k+1));
562  }
563  }
564 
565  Real dz_inv = one / dz_ptr[k];
566  temp_rhs_arr(i,j,k,Rho_comp ) += dz_inv * ( zflux_hi - zflux_lo );
567  temp_rhs_arr(i,j,k,RhoTheta_comp) += myhalf * dz_inv * ( zflux_hi * (prim(i,j,k) + prim(i,j,k+1))
568  - zflux_lo * (prim(i,j,k) + prim(i,j,k-1)) );
569  });
570 
571  // We only add to the flux registers in the final RK step
572  if (l_reflux) {
573  int strt_comp_reflux = 0;
574  // For now we don't reflux (rho theta) because it seems to create issues at c/f boundaries
575  int num_comp_reflux = 1;
576  if (level < finest_level) {
577  fr_as_crse->CrseAdd(mfi,
578  {{AMREX_D_DECL(&(flux[0]), &(flux[1]), &(flux[2]))}},
579  dx, dtau, strt_comp_reflux, strt_comp_reflux, num_comp_reflux, RunOn::Device);
580  }
581  if (level > 0) {
582  fr_as_fine->FineAdd(mfi,
583  {{AMREX_D_DECL(&(flux[0]), &(flux[1]), &(flux[2]))}},
584  dx, dtau, strt_comp_reflux, strt_comp_reflux, num_comp_reflux, RunOn::Device);
585  }
586 
587  // This is necessary here so we don't go on to the next FArrayBox without
588  // having finished copying the fluxes into the FluxRegisters (since the fluxes
589  // are stored in temporary FArrayBox's)
590  Gpu::streamSynchronize();
591 
592  } // two-way coupling
593  } // mfi
594  } // OMP
595 
596 #ifdef _OPENMP
597 #pragma omp parallel if (Gpu::notInLaunchRegion())
598 #endif
599  for ( MFIter mfi(S_stage_data[IntVars::cons],TilingIfNotGPU()); mfi.isValid(); ++mfi)
600  {
601  const Box& bx = mfi.tilebox();
602 
603  const Array4< Real>& cur_cons = S_data[IntVars::cons].array(mfi);
604  const Array4<const Real>& prev_cons = S_prev[IntVars::cons].const_array(mfi);
605  auto const& temp_rhs_arr = temp_rhs.const_array(mfi);
606  auto const& slow_rhs_cons = S_slow_rhs[IntVars::cons].const_array(mfi);
607  const Array4<Real const>& cc_src_arr = cc_src.const_array(mfi);
608 
609  if (step == 0) {
610  ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
611  {
612  cur_cons(i,j,k,Rho_comp) = prev_cons(i,j,k,Rho_comp) +
613  dtau * (slow_rhs_cons(i,j,k,Rho_comp) - temp_rhs_arr(i,j,k,Rho_comp));
614  cur_cons(i,j,k,RhoTheta_comp) = prev_cons(i,j,k,RhoTheta_comp) +
615  dtau * (slow_rhs_cons(i,j,k,RhoTheta_comp) - temp_rhs_arr(i,j,k,RhoTheta_comp));
616 
617  // add in source terms for cell-centered conserved variables
618  cur_cons(i,j,k,Rho_comp) += dtau * cc_src_arr(i,j,k,Rho_comp);
619  cur_cons(i,j,k,RhoTheta_comp) += dtau * cc_src_arr(i,j,k,RhoTheta_comp);
620  });
621  } else {
622  ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
623  {
624  //
625  // We didn't need to set cur_cons = prev_cons above because they point to the same data for step > 0
626  //
627  cur_cons(i,j,k,Rho_comp) += dtau * (slow_rhs_cons(i,j,k,Rho_comp) - temp_rhs_arr(i,j,k,Rho_comp));
628  cur_cons(i,j,k,RhoTheta_comp) += dtau * (slow_rhs_cons(i,j,k,RhoTheta_comp) - temp_rhs_arr(i,j,k,RhoTheta_comp));
629 
630  // add in source terms for cell-centered conserved variables
631  cur_cons(i,j,k,Rho_comp) += dtau * cc_src_arr(i,j,k,Rho_comp);
632  cur_cons(i,j,k,RhoTheta_comp) += dtau * cc_src_arr(i,j,k,RhoTheta_comp);
633  });
634  } // step = 0
635 
636  const Array4<Real>& cur_xmom = S_data[IntVars::xmom].array(mfi);
637  const Array4<Real>& cur_ymom = S_data[IntVars::ymom].array(mfi);
638 
639  const Array4<Real const>& temp_cur_xmom_arr = temp_cur_xmom.const_array(mfi);
640  const Array4<Real const>& temp_cur_ymom_arr = temp_cur_ymom.const_array(mfi);
641 
642  Box tbx = surroundingNodes(bx,0);
643  Box tby = surroundingNodes(bx,1);
644 
645  ParallelFor(tbx, tby,
646  [=] AMREX_GPU_DEVICE (int i, int j, int k)
647  {
648  cur_xmom(i,j,k) = temp_cur_xmom_arr(i,j,k);
649  },
650  [=] AMREX_GPU_DEVICE (int i, int j, int k)
651  {
652  cur_ymom(i,j,k) = temp_cur_ymom_arr(i,j,k);
653  });
654 
655  } // mfi
656 }
one
constexpr amrex::Real one
Definition: ERF_Constants.H:9
zero
constexpr amrex::Real zero
Definition: ERF_Constants.H:8
myhalf
constexpr amrex::Real myhalf
Definition: ERF_Constants.H:13
RvoRd
constexpr amrex::Real RvoRd
Definition: ERF_Constants.H:56
R_d
constexpr amrex::Real R_d
Definition: ERF_Constants.H:47
Gamma
constexpr amrex::Real Gamma
Definition: ERF_Constants.H:62
v_x
@ v_x
Definition: ERF_DataStruct.H:27
u_y
@ u_y
Definition: ERF_DataStruct.H:28
v_y
@ v_y
Definition: ERF_DataStruct.H:28
m_y
@ m_y
Definition: ERF_DataStruct.H:28
u_x
@ u_x
Definition: ERF_DataStruct.H:27
m_x
@ m_x
Definition: ERF_DataStruct.H:27
PrimQ1_comp
#define PrimQ1_comp
Definition: ERF_IndexDefines.H:58
Rho_comp
#define Rho_comp
Definition: ERF_IndexDefines.H:36
RhoTheta_comp
#define RhoTheta_comp
Definition: ERF_IndexDefines.H:37
PrimTheta_comp
#define PrimTheta_comp
Definition: ERF_IndexDefines.H:55
dxInv
amrex::GpuArray< Real, AMREX_SPACEDIM > dxInv
Definition: ERF_InitCustomPertVels_ParticleTests.H:17
dx
const Real dx
Definition: ERF_InitCustomPert_ABL.H:23
khi
const int khi
Definition: ERF_InitCustomPert_Bubble.H:21
AMREX_ALWAYS_ASSERT
AMREX_ALWAYS_ASSERT(bx.length()[2]==khi+1)
ParallelFor
ParallelFor(fab_box, [=] AMREX_GPU_DEVICE(int i, int j, int k) { qrcuten_arr(i, j, k)=Real(0);qscuten_arr(i, j, k)=Real(0);qicuten_arr(i, j, k)=Real(0);})
Real
amrex::Real Real
Definition: ERF_ShocInterface.H:19
TileNoZ
AMREX_FORCE_INLINE amrex::IntVect TileNoZ()
Definition: ERF_TileNoZ.H:11
GpVars::gpy
@ gpy
Definition: ERF_IndexDefines.H:187
GpVars::gpx
@ gpx
Definition: ERF_IndexDefines.H:186
IntVars::ymom
@ ymom
Definition: ERF_IndexDefines.H:196
IntVars::cons
@ cons
Definition: ERF_IndexDefines.H:194
IntVars::zmom
@ zmom
Definition: ERF_IndexDefines.H:197
IntVars::xmom
@ xmom
Definition: ERF_IndexDefines.H:195
MicVar_Kess::qt
@ qt
Definition: ERF_Kessler.H:29
MicVar_Kess::qv
@ qv
Definition: ERF_Kessler.H:30
WSM6Ind::q
@ q
Definition: ERF_WSM6.H:184
Here is the call graph for this function: