Skip to content

Commit

Permalink
Geostrophic wind profile (#1520)
Browse files Browse the repository at this point in the history 
* Geostrophic wind profile

* Format fixes

* Tab fixes

* fix double counting of pressure grad, formatting, remove amrex namespace from prob.cpp

* Fix gpu capture.

---------

Co-authored-by: Chandru Dhandapani <[email protected]>
Co-authored-by: AMLattanzi <[email protected]>
Co-authored-by: Aaron M. Lattanzi <[email protected]>
  • Loading branch information
@chandru320 @AMLattanzi
4 people authored Mar 22, 2024
1 parent 883c91b commit 45bdcfe
Show file tree
Hide file tree
Showing 13 changed files with 241 additions and 81 deletions.
1 change: 1 addition & 0 deletions Exec/DevTests/Bomex/input_SAM
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ erf.init_sounding_ideal = true

erf.add_custom_rhotheta_forcing = true
erf.add_custom_moisture_forcing = true
erf.add_custom_geostrophic_profile = true
erf.add_custom_w_subsidence = true
erf.custom_forcing_uses_primitive_vars = false

Expand Down
9 changes: 9 additions & 0 deletions Exec/DevTests/Bomex/prob.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -107,6 +107,15 @@ public:
const amrex::Geometry& geom,
std::unique_ptr<amrex::MultiFab>& z_phys_cc) override;

void update_geostrophic_profile (
const amrex::Real& /*time*/,
amrex::Vector<amrex::Real>& u_geos,
amrex::Gpu::DeviceVector<amrex::Real>& d_u_geos,
amrex::Vector<amrex::Real>& v_geos,
amrex::Gpu::DeviceVector<amrex::Real>& d_v_geos,
const amrex::Geometry& geom,
std::unique_ptr<amrex::MultiFab>& z_phys_cc) override;

protected:
std::string name () override { return "BOMEX"; }

Expand Down
141 changes: 92 additions & 49 deletions Exec/DevTests/Bomex/prob.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ amrex_probinit (const amrex_real* problo, const amrex_real* probhi)
return std::make_unique<Problem>(problo, probhi);
}

Problem::Problem (const amrex::Real* problo, const amrex::Real* probhi)
Problem::Problem (const Real* problo, const Real* probhi)
{
// Parse params
ParmParse pp("prob");
Expand Down Expand Up @@ -60,28 +60,28 @@ Problem::Problem (const amrex::Real* problo, const amrex::Real* probhi)

void
Problem::init_custom_pert (
const amrex::Box& bx,
const amrex::Box& xbx,
const amrex::Box& ybx,
const amrex::Box& zbx,
amrex::Array4<amrex::Real const> const& /*state*/,
amrex::Array4<amrex::Real > const& state_pert,
amrex::Array4<amrex::Real > const& x_vel_pert,
amrex::Array4<amrex::Real > const& y_vel_pert,
amrex::Array4<amrex::Real > const& z_vel_pert,
amrex::Array4<amrex::Real > const& /*r_hse*/,
amrex::Array4<amrex::Real > const& /*p_hse*/,
amrex::Array4<amrex::Real const> const& /*z_nd*/,
amrex::Array4<amrex::Real const> const& /*z_cc*/,
amrex::GeometryData const& geomdata,
amrex::Array4<amrex::Real const> const& /*mf_m*/,
amrex::Array4<amrex::Real const> const& /*mf_u*/,
amrex::Array4<amrex::Real const> const& /*mf_v*/,
const Box& bx,
const Box& xbx,
const Box& ybx,
const Box& zbx,
Array4<Real const> const& /*state*/,
Array4<Real > const& state_pert,
Array4<Real > const& x_vel_pert,
Array4<Real > const& y_vel_pert,
Array4<Real > const& z_vel_pert,
Array4<Real > const& /*r_hse*/,
Array4<Real > const& /*p_hse*/,
Array4<Real const> const& /*z_nd*/,
Array4<Real const> const& /*z_cc*/,
GeometryData const& geomdata,
Array4<Real const> const& /*mf_m*/,
Array4<Real const> const& /*mf_u*/,
Array4<Real const> const& /*mf_v*/,
const SolverChoice& sc)
{
const bool use_moisture = (sc.moisture_type != MoistureType::None);

ParallelForRNG(bx, [=, parms=parms] AMREX_GPU_DEVICE(int i, int j, int k, const amrex::RandomEngine& engine) noexcept
ParallelForRNG(bx, [=, parms=parms] AMREX_GPU_DEVICE(int i, int j, int k, const RandomEngine& engine) noexcept
{
// Geometry
const Real* prob_lo = geomdata.ProbLo();
Expand Down Expand Up @@ -117,7 +117,7 @@ Problem::init_custom_pert (
});

// Set the x-velocity
ParallelForRNG(xbx, [=, parms=parms] AMREX_GPU_DEVICE(int i, int j, int k, const amrex::RandomEngine& engine) noexcept
ParallelForRNG(xbx, [=, parms=parms] AMREX_GPU_DEVICE(int i, int j, int k, const RandomEngine& engine) noexcept
{
const Real* prob_lo = geomdata.ProbLo();
const Real* dx = geomdata.CellSize();
Expand All @@ -142,7 +142,7 @@ Problem::init_custom_pert (
});

// Set the y-velocity
ParallelForRNG(ybx, [=, parms=parms] AMREX_GPU_DEVICE(int i, int j, int k, const amrex::RandomEngine& engine) noexcept
ParallelForRNG(ybx, [=, parms=parms] AMREX_GPU_DEVICE(int i, int j, int k, const RandomEngine& engine) noexcept
{
const Real* prob_lo = geomdata.ProbLo();
const Real* dx = geomdata.CellSize();
Expand All @@ -167,7 +167,7 @@ Problem::init_custom_pert (
});

// Set the z-velocity
ParallelForRNG(zbx, [=, parms=parms] AMREX_GPU_DEVICE(int i, int j, int k, const amrex::RandomEngine& engine) noexcept
ParallelForRNG(zbx, [=, parms=parms] AMREX_GPU_DEVICE(int i, int j, int k, const RandomEngine& engine) noexcept
{
const int dom_lo_z = geomdata.Domain().smallEnd()[2];
const int dom_hi_z = geomdata.Domain().bigEnd()[2];
Expand All @@ -190,24 +190,24 @@ Problem::init_custom_pert (
// USER-DEFINED FUNCTION
//=============================================================================
void
Problem::update_rhotheta_sources (const amrex::Real& /*time*/,
amrex::Vector<amrex::Real>& src,
amrex::Gpu::DeviceVector<amrex::Real>& d_src,
const amrex::Geometry& geom,
std::unique_ptr<amrex::MultiFab>& z_phys_cc)
Problem::update_rhotheta_sources (const Real& /*time*/,
Vector<Real>& src,
Gpu::DeviceVector<Real>& d_src,
const Geometry& geom,
std::unique_ptr<MultiFab>& z_phys_cc)
{
if (src.empty()) return;

const int khi = geom.Domain().bigEnd()[2];
const amrex::Real* prob_lo = geom.ProbLo();
const auto dx = geom.CellSize();
const int khi = geom.Domain().bigEnd()[2];
const Real* prob_lo = geom.ProbLo();
const auto dx = geom.CellSize();

// Note: If z_phys_cc, then use_terrain=1 was set. If the z coordinate
// varies in time and or space, then the the height needs to be
// calculated at each time step. Here, we assume that only grid
// stretching exists.
if (z_phys_cc && zlevels.empty()) {
amrex::Print() << "Initializing z levels on stretched grid" << std::endl;
Print() << "Initializing z levels on stretched grid" << std::endl;
zlevels.resize(khi+1);
reduce_to_max_per_level(zlevels, z_phys_cc);
}
Expand All @@ -232,24 +232,24 @@ Problem::update_rhotheta_sources (const amrex::Real& /*time*/,
// USER-DEFINED FUNCTION
//=============================================================================
void
Problem::update_rhoqt_sources (const amrex::Real& /*time*/,
amrex::Vector<amrex::Real>& qsrc,
amrex::Gpu::DeviceVector<amrex::Real>& d_qsrc,
const amrex::Geometry& geom,
std::unique_ptr<amrex::MultiFab>& z_phys_cc)
Problem::update_rhoqt_sources (const Real& /*time*/,
Vector<Real>& qsrc,
Gpu::DeviceVector<Real>& d_qsrc,
const Geometry& geom,
std::unique_ptr<MultiFab>& z_phys_cc)
{
if (qsrc.empty()) return;

const int khi = geom.Domain().bigEnd()[2];
const amrex::Real* prob_lo = geom.ProbLo();
const auto dx = geom.CellSize();
const int khi = geom.Domain().bigEnd()[2];
const Real* prob_lo = geom.ProbLo();
const auto dx = geom.CellSize();

// Note: If z_phys_cc, then use_terrain=1 was set. If the z coordinate
// varies in time and or space, then the the height needs to be
// calculated at each time step. Here, we assume that only grid
// stretching exists.
if (z_phys_cc && zlevels.empty()) {
amrex::Print() << "Initializing z levels on stretched grid" << std::endl;
Print() << "Initializing z levels on stretched grid" << std::endl;
zlevels.resize(khi+1);
reduce_to_max_per_level(zlevels, z_phys_cc);
}
Expand All @@ -274,24 +274,24 @@ Problem::update_rhoqt_sources (const amrex::Real& /*time*/,
// USER-DEFINED FUNCTION
//=============================================================================
void
Problem::update_w_subsidence (const amrex::Real& /*time*/,
amrex::Vector<amrex::Real>& wbar,
amrex::Gpu::DeviceVector<amrex::Real>& d_wbar,
const amrex::Geometry& geom,
std::unique_ptr<amrex::MultiFab>& z_phys_cc)
Problem::update_w_subsidence (const Real& /*time*/,
Vector<Real>& wbar,
Gpu::DeviceVector<Real>& d_wbar,
const Geometry& geom,
std::unique_ptr<MultiFab>& z_phys_cc)
{
if (wbar.empty()) return;

const int khi = geom.Domain().bigEnd()[2];
const amrex::Real* prob_lo = geom.ProbLo();
const auto dx = geom.CellSize();
const int khi = geom.Domain().bigEnd()[2];
const Real* prob_lo = geom.ProbLo();
const auto dx = geom.CellSize();

// Note: If z_phys_cc, then use_terrain=1 was set. If the z coordinate
// varies in time and or space, then the the height needs to be
// calculated at each time step. Here, we assume that only grid
// stretching exists.
if (z_phys_cc && zlevels.empty()) {
amrex::Print() << "Initializing z levels on stretched grid" << std::endl;
Print() << "Initializing z levels on stretched grid" << std::endl;
zlevels.resize(khi+1);
reduce_to_max_per_level(zlevels, z_phys_cc);
}
Expand All @@ -315,3 +315,46 @@ Problem::update_w_subsidence (const amrex::Real& /*time*/,
// Copy from host version to device version
amrex::Gpu::copy(amrex::Gpu::hostToDevice, wbar.begin(), wbar.end(), d_wbar.begin());
}

//=============================================================================
// USER-DEFINED FUNCTION
//=============================================================================
void
Problem::update_geostrophic_profile (const Real& /*time*/,
Vector<Real>& u_geos,
Gpu::DeviceVector<Real>& d_u_geos,
Vector<Real>& v_geos,
Gpu::DeviceVector<Real>& d_v_geos,
const Geometry& geom,
std::unique_ptr<MultiFab>& z_phys_cc)
{
if (u_geos.empty()) return;

const int khi = geom.Domain().bigEnd()[2];
const Real* prob_lo = geom.ProbLo();
const auto dx = geom.CellSize();

// Note: If z_phys_cc, then use_terrain=1 was set. If the z coordinate
// varies in time and or space, then the the height needs to be
// calculated at each time step. Here, we assume that only grid
// stretching exists.
if (z_phys_cc && zlevels.empty()) {
Print() << "Initializing z levels on stretched grid" << std::endl;
zlevels.resize(khi+1);
reduce_to_max_per_level(zlevels, z_phys_cc);
}

const Real coriolis = 0.376E-4;

// Only apply temperature source below nominal inversion height
for (int k = 0; k <= khi; k++) {
const Real z_cc = (z_phys_cc) ? zlevels[k] : prob_lo[2] + (k+0.5)* dx[2];
const Real u_geo_wind = -10.0 + z_cc * 0.0018;
u_geos[k] = 0; // -coriolis_factor * v_geo_wind
v_geos[k] = coriolis * u_geo_wind;
}

// Copy from host version to device version
amrex::Gpu::copy(amrex::Gpu::hostToDevice, u_geos.begin(), u_geos.end(), d_u_geos.begin());
amrex::Gpu::copy(amrex::Gpu::hostToDevice, v_geos.begin(), v_geos.end(), d_v_geos.begin());
}
2 changes: 2 additions & 0 deletions Source/DataStructs/DataStruct.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -216,6 +216,7 @@ struct SolverChoice {
pp.query("add_custom_rhotheta_forcing", custom_rhotheta_forcing);
pp.query("add_custom_moisture_forcing", custom_moisture_forcing);
pp.query("add_custom_w_subsidence", custom_w_subsidence);
pp.query("add_custom_geostrophic_profile", custom_geostrophic_profile);
pp.query("custom_forcing_uses_primitive_vars", custom_forcing_prim_vars);

pp.query("Ave_Plane", ave_plane);
Expand Down Expand Up @@ -403,6 +404,7 @@ struct SolverChoice {
bool custom_rhotheta_forcing = false;
bool custom_moisture_forcing = false;
bool custom_w_subsidence = false;
bool custom_geostrophic_profile = false;
bool custom_forcing_prim_vars = false;

// Numerical diffusion
Expand Down
6 changes: 6 additions & 0 deletions Source/ERF.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -818,6 +818,12 @@ private:
amrex::Vector< amrex::Vector<amrex::Real> > h_w_subsid;
amrex::Vector<amrex::Gpu::DeviceVector<amrex::Real> > d_w_subsid;

amrex::Vector< amrex::Vector<amrex::Real> > h_u_geos;
amrex::Vector<amrex::Gpu::DeviceVector<amrex::Real> > d_u_geos;

amrex::Vector< amrex::Vector<amrex::Real> > h_v_geos;
amrex::Vector<amrex::Gpu::DeviceVector<amrex::Real> > d_v_geos;

// This is a vector over levels of vectors across quantities of Vectors
amrex::Vector<amrex::Vector< amrex::Vector<amrex::Real> > > h_rayleigh_ptrs;

Expand Down
19 changes: 19 additions & 0 deletions Source/ERF.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -677,6 +677,25 @@ ERF::InitData ()
}
}

if (solverChoice.custom_geostrophic_profile)
{
h_u_geos.resize(max_level+1, Vector<Real>(0));
d_u_geos.resize(max_level+1, Gpu::DeviceVector<Real>(0));
h_v_geos.resize(max_level+1, Vector<Real>(0));
d_v_geos.resize(max_level+1, Gpu::DeviceVector<Real>(0));
for (int lev = 0; lev <= finest_level; lev++) {
const int domlen = geom[lev].Domain().length(2);
h_u_geos[lev].resize(domlen, 0.0_rt);
d_u_geos[lev].resize(domlen, 0.0_rt);
h_v_geos[lev].resize(domlen, 0.0_rt);
d_v_geos[lev].resize(domlen, 0.0_rt);
prob->update_geostrophic_profile(t_new[0],
h_u_geos[lev], d_u_geos[lev],
h_v_geos[lev], d_v_geos[lev],
geom[lev], z_phys_cc[lev]);
}
}

if (solverChoice.custom_moisture_forcing)
{
h_rhoqt_src.resize(max_level+1, Vector<Real>(0));
Expand Down
5 changes: 0 additions & 5 deletions Source/Microphysics/SAM/IceFall.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,15 +36,10 @@ void SAM::IceFall () {
auto rho_array = rho->array(mfi);
auto fz_array = fz.array(mfi);

const auto dJ_array = (m_detJ_cc) ? m_detJ_cc->const_array(mfi) : Array4<const Real>{};

const auto& box3d = mfi.tilebox();

ParallelFor(box3d, [=] AMREX_GPU_DEVICE(int i, int j, int k) noexcept
{
// Jacobian determinant
Real dJinv = (dJ_array) ? 1.0/dJ_array(i,j,k) : 1.0;

Real rho_avg, qci_avg;
if (k==k_lo) {
rho_avg = rho_array(i,j,k);
Expand Down
9 changes: 9 additions & 0 deletions Source/TimeIntegration/ERF_advance_dycore.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,8 @@ void ERF::advance_dycore(int level,
Real* dptr_rhotheta_src = solverChoice.custom_rhotheta_forcing ? d_rhotheta_src[level].data() : nullptr;
Real* dptr_rhoqt_src = solverChoice.custom_moisture_forcing ? d_rhoqt_src[level].data() : nullptr;
Real* dptr_wbar_sub = solverChoice.custom_w_subsidence ? d_w_subsid[level].data() : nullptr;
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;

Vector<Real*> d_rayleigh_ptrs_at_lev;
d_rayleigh_ptrs_at_lev.resize(Rayleigh::nvars);
Expand Down Expand Up @@ -210,6 +212,13 @@ void ERF::advance_dycore(int level,
fine_geom, z_phys_cc[level]);
}

if (solverChoice.custom_geostrophic_profile) {
prob->update_geostrophic_profile(old_time,
h_u_geos[level], d_u_geos[level],
h_v_geos[level], d_v_geos[level],
fine_geom, z_phys_cc[level]);
}

// ***********************************************************************************************
// Convert old velocity available on faces to old momentum on faces to be used in time integration
// ***********************************************************************************************
Expand Down
14 changes: 14 additions & 0 deletions Source/TimeIntegration/ERF_slow_rhs_inc.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,8 @@ using namespace amrex;
* @param[in] mapfac_v map factor at y-faces
* @param[in] dptr_rhotheta_src custom temperature source term
* @param[in] dptr_rhoqt_src custom moisture source term
* @param[in] dptr_u_geos custom geostrophic wind profile
* @param[in] dptr_v_geos custom geostrophic wind profile
* @param[in] d_rayleigh_ptrs_at_lev Vector of {strength of Rayleigh damping, reference value for xvel/yvel/zvel/theta} used to define Rayleigh damping
*/

Expand Down Expand Up @@ -97,6 +99,8 @@ void erf_slow_rhs_inc (int /*level*/, int nrk,
std::unique_ptr<MultiFab>& mapfac_u,
std::unique_ptr<MultiFab>& mapfac_v,
const Real* dptr_rhotheta_src,
const Real* dptr_u_geos,
const Real* dptr_v_geos,
const Real* dptr_wbar_sub,
const Vector<Real*> d_rayleigh_dptrs)
{
Expand Down Expand Up @@ -752,6 +756,11 @@ void erf_slow_rhs_inc (int /*level*/, int nrk,
rho_u_rhs(i, j, k) += - solverChoice.abl_pressure_grad[0]
+ rho_on_u_face * solverChoice.abl_geo_forcing[0];

if (solverChoice.custom_geostrophic_profile) {
rho_u_rhs(i, j, k) += - solverChoice.abl_pressure_grad[0]
+ rho_on_u_face * dptr_u_geos[k];
}

// Add Coriolis forcing (that assumes east is +x, north is +y)
if (solverChoice.use_coriolis)
{
Expand Down Expand Up @@ -790,6 +799,11 @@ void erf_slow_rhs_inc (int /*level*/, int nrk,
rho_v_rhs(i, j, k) += - solverChoice.abl_pressure_grad[1]
+ rho_v_face * solverChoice.abl_geo_forcing[1];

if (solverChoice.custom_geostrophic_profile) {
rho_v_rhs(i, j, k) += - solverChoice.abl_pressure_grad[1]
+ rho_v_face * dptr_v_geos[k];
}

// Add Coriolis forcing (that assumes east is +x, north is +y)
if (solverChoice.use_coriolis)
{
Expand Down
Loading

0 comments on commit 45bdcfe

Please sign in to comment.