Skip to content

Commit

Permalink
cleaning up buoyancy plus adding more ML-related stuff (#1918)
Browse files Browse the repository at this point in the history
  • Loading branch information
@asalmgren
asalmgren authored Oct 30, 2024
1 parent 36e2d55 commit f6dd23b
Show file tree
Hide file tree
Showing 12 changed files with 169 additions and 72 deletions.
4 changes: 4 additions & 0 deletions Source/ERF.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -731,6 +731,10 @@ private:
amrex::Vector<amrex::MultiFab> rW_old;
amrex::Vector<amrex::MultiFab> rW_new;

amrex::Vector<amrex::MultiFab> xmom_crse_rhs;
amrex::Vector<amrex::MultiFab> ymom_crse_rhs;
amrex::Vector<amrex::MultiFab> zmom_crse_rhs;

std::unique_ptr<Microphysics> micro;
amrex::Vector<amrex::Vector<amrex::MultiFab*>> qmoist; // (lev,ncomp) This has up to 8 components: qt, qv, qc, qi, qp, qr, qs, qg

Expand Down
9 changes: 7 additions & 2 deletions Source/ERF.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -197,6 +197,10 @@ ERF::ERF_shared ()
rV_old.resize(nlevs_max);
rW_old.resize(nlevs_max);

xmom_crse_rhs.resize(nlevs_max);
ymom_crse_rhs.resize(nlevs_max);
zmom_crse_rhs.resize(nlevs_max);

for (int lev = 0; lev < nlevs_max; ++lev) {
vars_new[lev].resize(Vars::NumTypes);
vars_old[lev].resize(Vars::NumTypes);
Expand Down Expand Up @@ -460,8 +464,9 @@ ERF::post_timestep (int nstep, Real time, Real dt_lev0)
}
} // mfi

// This call refluxes from the lev/lev+1 interface onto lev
getAdvFluxReg(lev+1)->Reflux(vars_new[lev][Vars::cons], 0, 0, ncomp);
// This call refluxes all "slow" cell-centered variables
// (i.e. not density or (rho theta) or velocities) from the lev/lev+1 interface onto lev
getAdvFluxReg(lev+1)->Reflux(vars_new[lev][Vars::cons], 2, 2, ncomp-2);

// Here we multiply (rho S) by m^2 after refluxing
for (MFIter mfi(vars_new[lev][Vars::cons], TilingIfNotGPU()); mfi.isValid(); ++mfi) {
Expand Down
9 changes: 9 additions & 0 deletions Source/ERF_make_new_arrays.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -149,6 +149,15 @@ ERF::init_stuff (int lev, const BoxArray& ba, const DistributionMapping& dm,
rW_old[lev].define(convert(ba, IntVect(0,0,1)), dm, 1, ngrow_vels);
rW_new[lev].define(convert(ba, IntVect(0,0,1)), dm, 1, ngrow_vels);

if (lev > 0) {
xmom_crse_rhs[lev].define(convert(ba, IntVect(1,0,0)), dm, 1, IntVect{0});
xmom_crse_rhs[lev].setVal(3.456e22);
ymom_crse_rhs[lev].define(convert(ba, IntVect(0,1,0)), dm, 1, IntVect{0});
ymom_crse_rhs[lev].setVal(3.456e22);
zmom_crse_rhs[lev].define(convert(ba, IntVect(0,0,1)), dm, 1, IntVect{0});
zmom_crse_rhs[lev].setVal(3.456e22);
}

// We do this here just so they won't be undefined in the initial FillPatch
rU_old[lev].setVal(1.2e21);
rV_old[lev].setVal(3.4e22);
Expand Down
51 changes: 25 additions & 26 deletions Source/ERF_make_new_level.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -332,40 +332,39 @@ ERF::RemakeLevel (int lev, Real time, const BoxArray& ba, const DistributionMapp
// *****************************************************************************************************
make_physbcs(lev);

// ********************************************************************************************
// Update the base state at this level by interpolation from coarser level AND copy
// from previous (pre-regrid) base_state array
// ********************************************************************************************
Interpolater* mapper = &cell_cons_interp;

Vector<MultiFab*> fmf = {&base_state[lev ], &base_state[lev ]};
Vector<MultiFab*> cmf = {&base_state[lev-1], &base_state[lev-1]};
Vector<Real> ftime = {time, time};
Vector<Real> ctime = {time, time};

// Call FillPatch which ASSUMES that all ghost cells at lev-1 have already been filled
FillPatchTwoLevels(temp_base_state, temp_base_state.nGrowVect(), IntVect(0,0,0),
time, cmf, ctime, fmf, ftime,
0, 0, temp_base_state.nComp(), geom[lev-1], geom[lev],
refRatio(lev-1), mapper, domain_bcs_type,
BCVars::base_bc);

// Impose bc's outside the domain
(*physbcs_base[lev])(temp_base_state,0,temp_base_state.nComp(),base_state[lev].nGrowVect());

// *************************************************************************************************
// This will fill the temporary MultiFabs with data from vars_new
// NOTE: the momenta here are only used as scratch space, the momenta themselves are not fillpatched
// NOTE: we must create the new base state before calling FillPatch because we will
// interpolate perturbational quantities
// *************************************************************************************************
FillPatch(lev, time, {&temp_lev_new[Vars::cons],&temp_lev_new[Vars::xvel],
&temp_lev_new[Vars::yvel],&temp_lev_new[Vars::zvel]},
{&temp_lev_new[Vars::cons],&rU_new[lev],&rV_new[lev],&rW_new[lev]},
base_state[lev], temp_base_state, false);
base_state[lev],temp_base_state,false);

// Now swap the pointers since we needed both old and new in the FillPatch
std::swap(temp_base_state, base_state[lev]);
// ********************************************************************************************
// Update the base state at this level by interpolation from coarser level AND copy
// from previous (pre-regrid) base_state array
// ********************************************************************************************
if (lev > 0) {
Interpolater* mapper = &cell_cons_interp;

Vector<MultiFab*> fmf = {&base_state[lev ], &base_state[lev ]};
Vector<MultiFab*> cmf = {&base_state[lev-1], &base_state[lev-1]};
Vector<Real> ftime = {time, time};
Vector<Real> ctime = {time, time};

// Call FillPatch which ASSUMES that all ghost cells at lev-1 have already been filled
FillPatchTwoLevels(temp_base_state, temp_base_state.nGrowVect(), IntVect(0,0,0),
time, cmf, ctime, fmf, ftime,
0, 0, temp_base_state.nComp(), geom[lev-1], geom[lev],
refRatio(lev-1), mapper, domain_bcs_type,
BCVars::base_bc);

// Impose bc's outside the domain
(*physbcs_base[lev])(temp_base_state,0,temp_base_state.nComp(),base_state[lev].nGrowVect());

std::swap(temp_base_state, base_state[lev]);
}

// ********************************************************************************************
// Copy from new into old just in case
Expand Down
6 changes: 2 additions & 4 deletions Source/SourceTerms/ERF_Src_headers.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,7 @@ void make_buoyancy (amrex::Vector< amrex::MultiFab>& S_data,
amrex::MultiFab& buoyancy,
const amrex::Geometry geom,
const SolverChoice& solverChoice,
const amrex::MultiFab* r0,
const amrex::MultiFab* p0,
const amrex::MultiFab& base_state,
const int n_qstate,
const int anelastic);

Expand Down Expand Up @@ -56,8 +55,7 @@ void make_mom_sources (int level, int nrk,
amrex::MultiFab& xmom_source,
amrex::MultiFab& ymom_source,
amrex::MultiFab& zmom_source,
amrex::MultiFab* r0,
amrex::MultiFab* p0,
const amrex::MultiFab& base_state,
const amrex::Geometry geom,
const SolverChoice& solverChoice,
std::unique_ptr<amrex::MultiFab>& mapfac_m,
Expand Down
55 changes: 39 additions & 16 deletions Source/SourceTerms/ERF_buoyancy_utils.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,8 @@ buoyancy_dry_anelastic (int& i,
int& j,
int& k,
amrex::Real const& grav_gpu,
const amrex::Array4<const amrex::Real>& p0_arr,
const amrex::Array4<const amrex::Real>& r0_arr,
const amrex::Array4<const amrex::Real>& p0_arr,
const amrex::Array4<const amrex::Real>& cell_data)
{
// Note: this is the same term as the moist anelastic buoyancy when qv = qc = qt = 0
Expand All @@ -29,6 +29,32 @@ buoyancy_dry_anelastic (int& i,
return (-grav_gpu * (theta_d_wface - theta_d_0_wface) / theta_d_0_wface);
}

AMREX_GPU_DEVICE
AMREX_FORCE_INLINE
amrex::Real
buoyancy_dry_anelastic_T (int& i,
int& j,
int& k,
amrex::Real const& grav_gpu,
amrex::Real const& rd_over_cp,
const amrex::Array4<const amrex::Real>& r0_arr,
const amrex::Array4<const amrex::Real>& p0_arr,
const amrex::Array4<const amrex::Real>& cell_data)
{
amrex::Real rt0_hi = getRhoThetagivenP(p0_arr(i,j,k));
amrex::Real t0_hi = getTgivenPandTh(p0_arr(i,j,k), rt0_hi/r0_arr(i,j,k), rd_over_cp);
amrex::Real t_hi = getTgivenPandTh(p0_arr(i,j,k), cell_data(i,j,k,RhoTheta_comp)/r0_arr(i,j,k), rd_over_cp);
amrex::Real qplus = (t_hi-t0_hi)/t0_hi;

amrex::Real rt0_lo = getRhoThetagivenP(p0_arr(i,j,k-1));
amrex::Real t0_lo = getTgivenPandTh(p0_arr(i,j,k-1), rt0_lo/r0_arr(i,j,k-1), rd_over_cp);
amrex::Real t_lo = getTgivenPandTh(p0_arr(i,j,k-1), cell_data(i,j,k-1,RhoTheta_comp)/r0_arr(i,j,k-1), rd_over_cp);
amrex::Real qminus = (t_lo-t0_lo)/t0_lo;

amrex::Real r0_q_avg = amrex::Real(0.5) * (r0_arr(i,j,k) * qplus + r0_arr(i,j,k-1) * qminus);
return (-r0_q_avg * grav_gpu);
}


AMREX_GPU_DEVICE
AMREX_FORCE_INLINE
Expand All @@ -38,28 +64,26 @@ buoyancy_moist_anelastic (int& i,
int& k,
amrex::Real const& grav_gpu,
amrex::Real const& rv_over_rd,
const amrex::Array4<const amrex::Real>& p0_arr,
const amrex::Array4<const amrex::Real>& r0_arr,
const amrex::Array4<const amrex::Real>& r0_arr,
const amrex::Array4<const amrex::Real>& th0_arr,
const amrex::Array4<const amrex::Real>& cell_data)
{
amrex::Real theta_d_lo = cell_data(i,j,k-1,RhoTheta_comp)/r0_arr(i,j,k-1);
amrex::Real qv_lo = cell_data(i,j,k-1,RhoQ1_comp)/r0_arr(i,j,k-1);
amrex::Real qc_lo = cell_data(i,j,k-1,RhoQ2_comp)/r0_arr(i,j,k-1);
amrex::Real qv_lo = cell_data(i,j,k-1,RhoQ1_comp) /r0_arr(i,j,k-1);
amrex::Real qc_lo = cell_data(i,j,k-1,RhoQ2_comp) /r0_arr(i,j,k-1);
amrex::Real qt_lo = qv_lo + qc_lo;
amrex::Real theta_v_lo = theta_d_lo * (1.0 - (1.0 - rv_over_rd)*qt_lo - rv_over_rd*qc_lo);

amrex::Real theta_d_hi = cell_data(i,j,k,RhoTheta_comp)/r0_arr(i,j,k);
amrex::Real qv_hi = cell_data(i,j,k,RhoQ1_comp)/r0_arr(i,j,k);
amrex::Real qc_hi = cell_data(i,j,k,RhoQ2_comp)/r0_arr(i,j,k);
amrex::Real qv_hi = cell_data(i,j,k,RhoQ1_comp) /r0_arr(i,j,k);
amrex::Real qc_hi = cell_data(i,j,k,RhoQ2_comp) /r0_arr(i,j,k);
amrex::Real qt_hi = qv_hi + qc_hi;
amrex::Real theta_v_hi = theta_d_hi * (1.0 - (1.0 - rv_over_rd)*qt_hi - rv_over_rd*qc_hi);

amrex::Real theta_v_wface = amrex::Real(0.5) * (theta_v_lo + theta_v_hi);

amrex::Real theta_d_0_lo = getRhoThetagivenP(p0_arr(i,j,k-1)) / r0_arr(i,j,k-1);
amrex::Real theta_d_0_hi = getRhoThetagivenP(p0_arr(i,j,k )) / r0_arr(i,j,k );
amrex::Real theta_v_0_lo = theta_d_0_lo * (1.0 - (1.0 - rv_over_rd)*qt_lo - rv_over_rd*qc_lo);
amrex::Real theta_v_0_hi = theta_d_0_hi * (1.0 - (1.0 - rv_over_rd)*qt_hi - rv_over_rd*qc_hi);
amrex::Real theta_v_0_lo = th0_arr(i,j,k-1) * (1.0 - (1.0 - rv_over_rd)*qt_lo - rv_over_rd*qc_lo);
amrex::Real theta_v_0_hi = th0_arr(i,j,k ) * (1.0 - (1.0 - rv_over_rd)*qt_hi - rv_over_rd*qc_hi);

amrex::Real theta_v_0_wface = amrex::Real(0.5) * (theta_v_0_lo + theta_v_0_hi);

Expand All @@ -74,17 +98,16 @@ buoyancy_dry_default (int& i,
int& k,
amrex::Real const& grav_gpu,
amrex::Real const& rd_over_cp,
const amrex::Array4<const amrex::Real>& p0_arr,
const amrex::Array4<const amrex::Real>& r0_arr,
const amrex::Array4<const amrex::Real>& p0_arr,
const amrex::Array4<const amrex::Real>& th0_arr,
const amrex::Array4<const amrex::Real>& cell_data)
{
amrex::Real rt0_hi = getRhoThetagivenP(p0_arr(i,j,k));
amrex::Real t0_hi = getTgivenPandTh(p0_arr(i,j,k), rt0_hi/r0_arr(i,j,k), rd_over_cp);
amrex::Real t0_hi = getTgivenPandTh(p0_arr(i,j,k), th0_arr(i,j,k), rd_over_cp);
amrex::Real t_hi = getTgivenRandRTh(cell_data(i,j,k ,Rho_comp), cell_data(i,j,k ,RhoTheta_comp));
amrex::Real qplus = (t_hi-t0_hi)/t0_hi;

amrex::Real rt0_lo = getRhoThetagivenP(p0_arr(i,j,k-1));
amrex::Real t0_lo = getTgivenPandTh(p0_arr(i,j,k-1), rt0_lo/r0_arr(i,j,k-1), rd_over_cp);
amrex::Real t0_lo = getTgivenPandTh(p0_arr(i,j,k-1), th0_arr(i,j,k-1), rd_over_cp);
amrex::Real t_lo = getTgivenRandRTh(cell_data(i,j,k-1,Rho_comp), cell_data(i,j,k-1,RhoTheta_comp));
amrex::Real qminus = (t_lo-t0_lo)/t0_lo;

Expand Down
27 changes: 16 additions & 11 deletions Source/SourceTerms/ERF_make_buoyancy.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,8 +34,7 @@ void make_buoyancy (Vector<MultiFab>& S_data,
MultiFab& buoyancy,
const amrex::Geometry geom,
const SolverChoice& solverChoice,
const MultiFab* r0,
const MultiFab* p0,
const MultiFab& base_state,
const int n_qstate,
const int anelastic)
{
Expand All @@ -50,6 +49,10 @@ void make_buoyancy (Vector<MultiFab>& S_data,
Real rd_over_cp = solverChoice.rdOcp;
Real rv_over_rd = R_v/R_d;

MultiFab r0 (base_state, make_alias, BaseState::r0_comp , 1);
MultiFab p0 (base_state, make_alias, BaseState::p0_comp , 1);
MultiFab th0(base_state, make_alias, BaseState::th0_comp, 1);

if (anelastic == 1) {
#ifdef _OPENMP
#pragma omp parallel if (amrex::Gpu::notInLaunchRegion())
Expand All @@ -66,15 +69,16 @@ void make_buoyancy (Vector<MultiFab>& S_data,
const Array4< Real> & buoyancy_fab = buoyancy.array(mfi);

// Base state density and pressure
const Array4<const Real>& r0_arr = r0->const_array(mfi);
const Array4<const Real>& p0_arr = p0->const_array(mfi);
const Array4<const Real>& r0_arr = r0.const_array(mfi);
const Array4<const Real>& p0_arr = p0.const_array(mfi);
const Array4<const Real>& th0_arr = th0.const_array(mfi);

if (solverChoice.moisture_type == MoistureType::None) {
ParallelFor(tbz, [=] AMREX_GPU_DEVICE (int i, int j, int k)
{
buoyancy_fab(i, j, k) = buoyancy_dry_anelastic(i,j,k,
grav_gpu[2],
p0_arr,r0_arr,cell_data);
r0_arr,p0_arr,cell_data);
});
} else {
// NOTE: For decomposition in the vertical direction, klo may not
Expand All @@ -84,7 +88,7 @@ void make_buoyancy (Vector<MultiFab>& S_data,
{
buoyancy_fab(i, j, k) = buoyancy_moist_anelastic(i,j,k,
grav_gpu[2],rv_over_rd,
p0_arr,r0_arr,cell_data);
r0_arr,th0_arr,cell_data);
});
}
} // mfi
Expand Down Expand Up @@ -113,7 +117,7 @@ void make_buoyancy (Vector<MultiFab>& S_data,
const Array4< Real> & buoyancy_fab = buoyancy.array(mfi);

// Base state density
const Array4<const Real>& r0_arr = r0->const_array(mfi);
const Array4<const Real>& r0_arr = r0.const_array(mfi);

ParallelFor(tbz, [=] AMREX_GPU_DEVICE (int i, int j, int k)
{
Expand Down Expand Up @@ -142,8 +146,9 @@ void make_buoyancy (Vector<MultiFab>& S_data,
if (tbz.bigEnd(2) == khi) tbz.growHi(2,-1);

// Base state density and pressure
const Array4<const Real>& r0_arr = r0->const_array(mfi);
const Array4<const Real>& p0_arr = p0->const_array(mfi);
const Array4<const Real>& r0_arr = r0.const_array(mfi);
const Array4<const Real>& p0_arr = p0.const_array(mfi);
const Array4<const Real>& th0_arr = th0.const_array(mfi);

const Array4<const Real> & cell_data = S_data[IntVars::cons].array(mfi);
const Array4< Real> & buoyancy_fab = buoyancy.array(mfi);
Expand All @@ -152,7 +157,7 @@ void make_buoyancy (Vector<MultiFab>& S_data,
{
buoyancy_fab(i, j, k) = buoyancy_dry_default(i,j,k,
grav_gpu[2],rd_over_cp,
p0_arr,r0_arr,cell_data);
r0_arr,p0_arr,th0_arr,cell_data);
});
} // mfi
} // buoyancy_type
Expand Down Expand Up @@ -187,7 +192,7 @@ void make_buoyancy (Vector<MultiFab>& S_data,
const Array4< Real> & buoyancy_fab = buoyancy.array(mfi);

// Base state density
const Array4<const Real>& r0_arr = r0->const_array(mfi);
const Array4<const Real>& r0_arr = r0.const_array(mfi);

ParallelFor(tbz, [=] AMREX_GPU_DEVICE (int i, int j, int k)
{
Expand Down
4 changes: 2 additions & 2 deletions Source/SourceTerms/ERF_make_mom_sources.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ void make_mom_sources (int level,
MultiFab & xmom_src,
MultiFab & ymom_src,
MultiFab & zmom_src,
MultiFab* r0, MultiFab* p0,
const MultiFab& base_state,
const Geometry geom,
const SolverChoice& solverChoice,
std::unique_ptr<MultiFab>& mapfac_m,
Expand Down Expand Up @@ -190,7 +190,7 @@ void make_mom_sources (int level,
// *****************************************************************************
// 1. Create the BUOYANCY forcing term in the z-direction
// *****************************************************************************
make_buoyancy(S_data, S_prim, zmom_src, geom, solverChoice, r0, p0,
make_buoyancy(S_data, S_prim, zmom_src, geom, solverChoice, base_state,
n_qstate, solverChoice.anelastic[level]);

// *****************************************************************************
Expand Down
Loading

0 comments on commit f6dd23b

Please sign in to comment.