Skip to content

Commit

Permalink
clean up to remove code duplication (#1617)
Browse files Browse the repository at this point in the history
  • Loading branch information
@asalmgren
asalmgren authored May 11, 2024
1 parent 9349d66 commit 96958c1
Show file tree
Hide file tree
Showing 7 changed files with 107 additions and 544 deletions.
7 changes: 2 additions & 5 deletions Source/Advection/Advection.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,11 +27,8 @@ void AdvectionSrcForRho (const amrex::Box& bx,
const amrex::Array4<const amrex::Real>& mf_m,
const amrex::Array4<const amrex::Real>& mf_u,
const amrex::Array4<const amrex::Real>& mf_v,
const amrex::GpuArray<const amrex::Array4<amrex::Real>, AMREX_SPACEDIM>& flx_arr
#ifdef ERF_USE_POISSON_SOLVE
,const bool const_rho = false
#endif
);
const amrex::GpuArray<const amrex::Array4<amrex::Real>, AMREX_SPACEDIM>& flx_arr,
const bool const_rho);

/** Compute advection tendency for all scalars other than density and potential temperature */
void AdvectionSrcForScalars (const amrex::Box& bx,
Expand Down
8 changes: 2 additions & 6 deletions Source/Advection/AdvectionSrcForState.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,10 +43,8 @@ AdvectionSrcForRho (const Box& bx,
const Array4<const Real>& mf_m,
const Array4<const Real>& mf_u,
const Array4<const Real>& mf_v,
const GpuArray<const Array4<Real>, AMREX_SPACEDIM>& flx_arr
#ifdef ERF_USE_POISSON_SOLVE
,const bool const_rho
#endif
const GpuArray<const Array4<Real>, AMREX_SPACEDIM>& flx_arr,
const bool const_rho
)
{
BL_PROFILE_VAR("AdvectionSrcForRho", AdvectionSrcForRho);
Expand All @@ -73,14 +71,12 @@ AdvectionSrcForRho (const Box& bx,
avg_zmom(i,j,k ) = (flx_arr[2])(i,j,k,0);
});

#ifdef ERF_USE_POISSON_SOLVE
if (const_rho) {
ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
{
advectionSrc(i,j,k,0) = 0.0;
});
} else
#endif
{
ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
{
Expand Down
500 changes: 0 additions & 500 deletions Source/TimeIntegration/ERF_slow_rhs_inc.cpp
View file

This file was deleted.

89 changes: 72 additions & 17 deletions Source/TimeIntegration/ERF_slow_rhs_pre.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ using namespace amrex;
* @param[in] nrk which RK stage
* @param[in] dt slow time step
* @param[out] S_rhs RHS computed here
* @param[in] S_old old-time solution -- used only for incompressible
* @param[in] S_data current solution
* @param[in] S_prim primitive variables (i.e. conserved variables divided by density)
* @param[in] S_scratch scratch space
Expand Down Expand Up @@ -66,6 +67,7 @@ void erf_slow_rhs_pre (int level, int finest_level,
int nrk,
Real dt,
Vector<MultiFab>& S_rhs,
Vector<MultiFab>& S_old,
Vector<MultiFab>& S_data,
const MultiFab& S_prim,
Vector<MultiFab>& S_scratch,
Expand Down Expand Up @@ -148,6 +150,18 @@ void erf_slow_rhs_pre (int level, int finest_level,
const bool use_most = (most != nullptr);
const bool exp_most = (solverChoice.use_explicit_most);

#ifdef ERF_USE_POISSON_SOlVE
const bool l_incompressible = solverChoice.incompressible;
const bool l_const_rho = solverChoice.constant_density;

// We cannot use incompressible with terrain or with moisture
AMREX_ALWAYS_ASSERT(!l_use_terrain || !incompressible);
AMREX_ALWAYS_ASSERT(!l_use_moisture || !incompressible);
#else
const bool l_incompressible = false;
const bool l_const_rho = false;
#endif

const Box& domain = geom.Domain();
const int domhi_z = domain.bigEnd(2);
const int domlo_z = domain.smallEnd(2);
Expand Down Expand Up @@ -222,10 +236,23 @@ void erf_slow_rhs_pre (int level, int finest_level,
const Array4<const Real> & cell_prim = S_prim.array(mfi);
const Array4<Real> & cell_rhs = S_rhs[IntVars::cons].array(mfi);

const Array4<const Real> & cell_old = S_old[IntVars::cons].array(mfi);

const Array4<Real const>& xmom_src_arr = xmom_src.const_array(mfi);
const Array4<Real const>& ymom_src_arr = ymom_src.const_array(mfi);
const Array4<Real const>& zmom_src_arr = zmom_src.const_array(mfi);

const Array4<Real>& rho_u_old = S_old[IntVars::xmom].array(mfi);
const Array4<Real>& rho_v_old = S_old[IntVars::ymom].array(mfi);
const Array4<Real>& rho_w_old = S_old[IntVars::zmom].array(mfi);

if (l_incompressible) {
// When incompressible we must reset these to 0 each RK step
S_scratch[IntVars::xmom][mfi].template setVal<RunOn::Device>(0.0,tbx);
S_scratch[IntVars::ymom][mfi].template setVal<RunOn::Device>(0.0,tby);
S_scratch[IntVars::zmom][mfi].template setVal<RunOn::Device>(0.0,tbz);
}

Array4<Real> avg_xmom = S_scratch[IntVars::xmom].array(mfi);
Array4<Real> avg_ymom = S_scratch[IntVars::ymom].array(mfi);
Array4<Real> avg_zmom = S_scratch[IntVars::zmom].array(mfi);
Expand Down Expand Up @@ -277,21 +304,26 @@ void erf_slow_rhs_pre (int level, int finest_level,
// *****************************************************************************
// Perturbational pressure field
// *****************************************************************************
Box gbx = mfi.tilebox(); gbx.grow(IntVect(1,1,1));
if (gbx.smallEnd(2) < 0) gbx.setSmall(2,0);
FArrayBox pprime; pprime.resize(gbx,1,The_Async_Arena());
const Array4<Real > & pp_arr = pprime.array();
{
BL_PROFILE("slow_rhs_pre_pprime");
ParallelFor(gbx, [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
{
//if (cell_data(i,j,k,RhoTheta_comp) < 0.) printf("BAD THETA AT %d %d %d %e %e \n",
// i,j,k,cell_data(i,j,k,RhoTheta_comp),cell_data(i,j,k+1,RhoTheta_comp));
AMREX_ASSERT(cell_data(i,j,k,RhoTheta_comp) > 0.);
Real qv_for_p = (use_moisture) ? cell_data(i,j,k,RhoQ1_comp)/cell_data(i,j,k,Rho_comp) : 0.0;
pp_arr(i,j,k) = getPgivenRTh(cell_data(i,j,k,RhoTheta_comp),qv_for_p) - p0_arr(i,j,k);
});
} // end profile
FArrayBox pprime;
if (!l_incompressible) {
Box gbx = mfi.tilebox(); gbx.grow(IntVect(1,1,1));
if (gbx.smallEnd(2) < 0) gbx.setSmall(2,0);
pprime.resize(gbx,1,The_Async_Arena());
const Array4<Real>& pptemp_arr = pprime.array();
ParallelFor(gbx, [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
{
//if (cell_data(i,j,k,RhoTheta_comp) < 0.) printf("BAD THETA AT %d %d %d %e %e \n",
// i,j,k,cell_data(i,j,k,RhoTheta_comp),cell_data(i,j,k+1,RhoTheta_comp));
AMREX_ASSERT(cell_data(i,j,k,RhoTheta_comp) > 0.);
Real qv_for_p = (use_moisture) ? cell_data(i,j,k,RhoQ1_comp)/cell_data(i,j,k,Rho_comp) : 0.0;
pptemp_arr(i,j,k) = getPgivenRTh(cell_data(i,j,k,RhoTheta_comp),qv_for_p) - p0_arr(i,j,k);
});
}
#ifdef ERF_USE_POISSON_SOLVE
const Array4<const Real>& pp_arr = (l_incompressible) ? pp_inc.const_array(mfi) : pprime.const_array();
#else
const Array4<const Real>& pp_arr = pprime.const_array();
#endif

// *****************************************************************************
// Contravariant flux field
Expand Down Expand Up @@ -393,7 +425,7 @@ void erf_slow_rhs_pre (int level, int finest_level,
avg_xmom, avg_ymom, avg_zmom, // these are being defined from the fluxes
ax_arr, ay_arr, az_arr, detJ_arr,
dxInv, mf_m, mf_u, mf_v,
flx_arr);
flx_arr, l_const_rho);

int icomp = RhoTheta_comp; int ncomp = 1;
AdvectionSrcForScalars(bx, icomp, ncomp,
Expand Down Expand Up @@ -453,6 +485,19 @@ void erf_slow_rhs_pre (int level, int finest_level,
});
}

// If incompressible and in second RK stage, take average of old-time and new-time source
if ( l_incompressible && (nrk == 1) )
{
ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
{
cell_rhs(i,j,k, Rho_comp) *= 0.5;
cell_rhs(i,j,k,RhoTheta_comp) *= 0.5;

cell_rhs(i,j,k, Rho_comp) += 0.5 / dt * (cell_data(i,j,k, Rho_comp) - cell_old(i,j,k, Rho_comp));
cell_rhs(i,j,k,RhoTheta_comp) += 0.5 / dt * (cell_data(i,j,k,RhoTheta_comp) - cell_old(i,j,k,RhoTheta_comp));
});
}

// *****************************************************************************
// Define updates in the RHS of {x, y, z}-momentum equations
// *****************************************************************************
Expand Down Expand Up @@ -546,6 +591,11 @@ void erf_slow_rhs_pre (int level, int finest_level,
Real h_zeta = Compute_h_zeta_AtIface(i, j, k, dxInv, z_nd);
rho_u_rhs(i, j, k) *= h_zeta;
}

if ( l_incompressible && (nrk == 1) ) {
rho_u_rhs(i,j,k) *= 0.5;
rho_u_rhs(i,j,k) += 0.5 / dt * (rho_u(i,j,k) - rho_u_old(i,j,k));
}
});

ParallelFor(tby, [=] AMREX_GPU_DEVICE (int i, int j, int k)
Expand Down Expand Up @@ -589,6 +639,11 @@ void erf_slow_rhs_pre (int level, int finest_level,
Real h_zeta = Compute_h_zeta_AtJface(i, j, k, dxInv, z_nd);
rho_v_rhs(i, j, k) *= h_zeta;
}

if ( l_incompressible && (nrk == 1) ) {
rho_v_rhs(i,j,k) *= 0.5;
rho_v_rhs(i,j,k) += 0.5 / dt * (rho_v(i,j,k) - rho_v_old(i,j,k));
}
});

// *****************************************************************************
Expand Down Expand Up @@ -666,7 +721,7 @@ void erf_slow_rhs_pre (int level, int finest_level,
// NOTE: for now we are only refluxing density not (rho theta) since the latter seems to introduce
// a problem at top and bottom boundaries
if (l_reflux && nrk == 2) {
int strt_comp_reflux = 0;
int strt_comp_reflux = (l_const_rho) ? 1 : 0;
int num_comp_reflux = 1;
if (level < finest_level) {
fr_as_crse->CrseAdd(mfi,
Expand Down
4 changes: 0 additions & 4 deletions Source/TimeIntegration/Make.package
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,10 +12,6 @@ CEXE_sources += ERF_fast_rhs_N.cpp
CEXE_sources += ERF_fast_rhs_T.cpp
CEXE_sources += ERF_fast_rhs_MT.cpp

ifeq ($(USE_POISSON_SOLVE),TRUE)
CEXE_sources += ERF_slow_rhs_inc.cpp
endif

CEXE_headers += TI_fast_rhs_fun.H
CEXE_headers += TI_slow_rhs_fun.H
CEXE_headers += TI_no_substep_fun.H
Expand Down
5 changes: 3 additions & 2 deletions Source/TimeIntegration/TI_slow_headers.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -57,9 +57,10 @@ void erf_make_tau_terms (int level, int nrk,
void erf_slow_rhs_pre (int level, int finest_level, int nrk,
amrex::Real dt,
amrex::Vector<amrex::MultiFab>& S_rhs,
amrex::Vector<amrex::MultiFab>& S_old,
amrex::Vector<amrex::MultiFab>& S_data,
const amrex::MultiFab& S_prim,
amrex::Vector<amrex::MultiFab >& S_scratch,
const amrex::MultiFab & S_prim,
amrex::Vector<amrex::MultiFab>& S_scratch,
const amrex::MultiFab& xvel,
const amrex::MultiFab& yvel,
const amrex::MultiFab& zvel,
Expand Down
38 changes: 28 additions & 10 deletions Source/TimeIntegration/TI_slow_rhs_fun.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,7 @@
* Wrapper for calling the routine that creates the slow RHS
*/
auto slow_rhs_fun_pre = [&](Vector<MultiFab>& S_rhs,
#ifdef ERF_USE_NETCDF
Vector<MultiFab>& S_old,
#else
Vector<MultiFab>& /*S_old*/,
#endif
Vector<MultiFab>& S_data,
Vector<MultiFab>& S_scratch,
const Real old_step_time,
Expand Down Expand Up @@ -119,7 +115,7 @@
dptr_u_geos, dptr_v_geos, dptr_wbar_sub,
d_rayleigh_ptrs_at_lev, d_sponge_ptrs_at_lev, n_qstate);

erf_slow_rhs_pre(level, finest_level, nrk, slow_dt, S_rhs, S_data, S_prim, S_scratch,
erf_slow_rhs_pre(level, finest_level, nrk, slow_dt, S_rhs, S_old, S_data, S_prim, S_scratch,
xvel_new, yvel_new, zvel_new,
z_t_rk[level], Omega, cc_src, xmom_src, ymom_src, zmom_src,
Tau11_lev[level].get(), Tau22_lev[level].get(), Tau33_lev[level].get(), Tau12_lev[level].get(),
Expand Down Expand Up @@ -222,7 +218,7 @@
dptr_u_geos, dptr_v_geos, dptr_wbar_sub,
d_rayleigh_ptrs_at_lev, d_sponge_ptrs_at_lev, n_qstate);

erf_slow_rhs_pre(level, finest_level, nrk, slow_dt, S_rhs, S_data, S_prim, S_scratch,
erf_slow_rhs_pre(level, finest_level, nrk, slow_dt, S_rhs, S_old, S_data, S_prim, S_scratch,
xvel_new, yvel_new, zvel_new,
z_t_rk[level], Omega, cc_src, xmom_src, ymom_src, zmom_src,
Tau11_lev[level].get(), Tau22_lev[level].get(), Tau33_lev[level].get(), Tau12_lev[level].get(),
Expand Down Expand Up @@ -392,6 +388,21 @@

Real slow_dt = new_stage_time - old_step_time;

// *************************************************************************
// Set up flux registers if using two_way coupling
// *************************************************************************
YAFluxRegister* fr_as_crse = nullptr;
YAFluxRegister* fr_as_fine = nullptr;
if (solverChoice.coupling_type == CouplingType::TwoWay) {
if (level < finest_level) {
fr_as_crse = getAdvFluxReg(level+1);
fr_as_crse->reset();
}
if (level > 0) {
fr_as_fine = getAdvFluxReg(level);
}
}

Real* dptr_u_geos = solverChoice.custom_geostrophic_profile ? d_u_geos[level].data(): nullptr;
Real* dptr_v_geos = solverChoice.custom_geostrophic_profile ? d_v_geos[level].data(): nullptr;

Expand All @@ -413,16 +424,23 @@
dptr_u_geos, dptr_v_geos, dptr_wbar_sub,
d_rayleigh_ptrs_at_lev, d_sponge_ptrs_at_lev, n_qstate);

erf_slow_rhs_inc(level, nrk, slow_dt,
erf_slow_rhs_pre(level, finest_level, nrk, slow_dt,
S_rhs, S_old, S_data, S_prim, S_scratch,
xvel_new, yvel_new, zvel_new,
Omega, cc_src, xmom_src, ymom_src, zmom_src,
z_t_rk[level], Omega, cc_src, xmom_src, ymom_src, zmom_src,
Tau11_lev[level].get(), Tau22_lev[level].get(), Tau33_lev[level].get(), Tau12_lev[level].get(),
Tau13_lev[level].get(), Tau21_lev[level].get(), Tau23_lev[level].get(), Tau31_lev[level].get(),
Tau32_lev[level].get(), SmnSmn, eddyDiffs, Hfx3, Diss,
fine_geom, solverChoice, m_most, domain_bcs_type_d, domain_bcs_type,
z_phys_nd[level], ax[level], ay[level], az[level], detJ_cc[level], p0, pp_inc[level],
mapfac_m[level], mapfac_u[level], mapfac_v[level]);
z_phys_nd[level], ax[level], ay[level], az[level], detJ_cc[level], p0,
#ifdef ERF_USE_POISSON_SOLVE
pp_inc[level],
#endif
mapfac_m[level], mapfac_u[level], mapfac_v[level],
#ifdef ERF_USE_EB
EBFactory(level),
#endif
fr_as_crse, fr_as_fine);

add_thin_body_sources(xmom_src, ymom_src, zmom_src,
xflux_imask[level], yflux_imask[level], zflux_imask[level],
Expand Down

0 comments on commit 96958c1

Please sign in to comment.