Skip to content

Commit

Permalink
Fix GPU issue for MetGrid with new Newton iters.
Browse files Browse the repository at this point in the history
  • Loading branch information
@AMLattanzi
AMLattanzi committed Sep 26, 2024
1 parent 88b928b commit e2cf47a
Show file tree
Hide file tree
Showing 2 changed files with 132 additions and 90 deletions.
84 changes: 0 additions & 84 deletions Source/Initialization/ERF_Metgrid_utils.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -101,90 +101,6 @@ init_base_state_from_metgrid (const bool use_moisture,
amrex::Vector<amrex::Vector<amrex::FArrayBox>>& fabs_for_bcs,
const amrex::Array4<const int>& mask_c_arr);

AMREX_FORCE_INLINE
AMREX_GPU_DEVICE
void
calc_rho_p (const int& kmax,
const int& flag_psfc,
const amrex::Real& psfc,
amrex::Real* Thetad_vec,
amrex::Real* Thetam_vec,
amrex::Real* Q_vec,
amrex::Real* z_vec,
amrex::Real* Rhod_vec,
amrex::Real* Rhom_vec,
amrex::Real* Pd_vec,
amrex::Real* Pm_vec)
{
const int maxiter = 10;
const amrex::Real tol = 1.0e-12;

// Calculate or use moist pressure at the surface.
amrex::Real Psurf;
if (flag_psfc == 1) {
Psurf = psfc;
} else {
amrex::Real t_0 = 290.0; // WRF's model_config_rec%base_temp
amrex::Real a = 50.0; // WRF's model_config_rec%base_lapse
Psurf = p_0*exp(-t_0/a+std::pow((std::pow(t_0/a, 2)-2.0*CONST_GRAV*z_vec[0]/(a*R_d)), 0.5));
}

// Iterations for the first CC point that is 1/2 dz off the surface
{
amrex::Real half_dz = z_vec[0];
amrex::Real qvf = 1.0+(R_v/R_d)*Q_vec[0];
Thetam_vec[0] = Thetad_vec[0]*qvf;
for (int it=0; it<maxiter; it++) {
Pm_vec[0] = Psurf-half_dz*(Rhom_vec[0])*(1.0+Q_vec[0])*CONST_GRAV;
if (Pm_vec[0] < 0.0) Pm_vec[0] = 0.0;
Rhom_vec[0] = (p_0/(R_d*Thetam_vec[0]))*std::pow(Pm_vec[0]/p_0, iGamma);
} // it
}

// Integrate from the first CC point to the top boundary.
for (int k=1; k<=kmax; k++) {

// Assignment only for downward integration in following loop
amrex::Real qvf = 1.0+(R_v/R_d)*Q_vec[k];
Thetam_vec[k] = Thetad_vec[k]*qvf;

// Vertical grid spacing
amrex::Real dz = z_vec[k]-z_vec[k-1];

// Establish known constant
amrex::Real rho_tot_lo = Rhom_vec[k-1] * (1. + Q_vec[k-1]);
amrex::Real C = -Pm_vec[k-1] + 0.5*rho_tot_lo*CONST_GRAV*dz;

// Initial guess and residual
Pm_vec[k] = Pm_vec[k-1];
Rhom_vec[k] = getRhogivenThetaPress(Thetad_vec[k],
Pm_vec[k],
R_d/Cp_d,
Q_vec[k]);
amrex::Real rho_tot_hi = Rhom_vec[k] * (1. + Q_vec[k]);
amrex::Real F = Pm_vec[k] + 0.5*rho_tot_hi*CONST_GRAV*dz + C;

// Do iterations
if (std::abs(F)>tol) HSEutils::Newton_Raphson_hse(tol, R_d/Cp_d, dz,
CONST_GRAV, C, Thetad_vec[k],
Q_vec[k], Q_vec[k],
Pm_vec[k], Rhom_vec[k], F);
} // k

// integrate from the top back down to get dry pressure and density.
Pd_vec[kmax] = Pm_vec[kmax];
Rhod_vec[kmax] = (p_0/(R_d*Thetam_vec[kmax]))*std::pow(Pd_vec[kmax]/p_0, iGamma);
for (int k=kmax-1; k>=0; k--) {
amrex::Real dz = z_vec[k+1]-z_vec[k];
Rhod_vec[k] = Rhod_vec[k+1]; // an initial guess.
for (int it=0; it<maxiter; it++) {
Pd_vec[k] = Pd_vec[k+1]+0.5*dz*(Rhod_vec[k]+Rhod_vec[k+1])*CONST_GRAV;
if (Pd_vec[k] < 0.0) Pd_vec[k] = 0.0;
Rhod_vec[k] = (p_0/(R_d*Thetam_vec[k]))*std::pow(Pd_vec[k]/p_0, iGamma);
} // it
} // k
}

AMREX_FORCE_INLINE
AMREX_GPU_DEVICE
amrex::Real
Expand Down
138 changes: 132 additions & 6 deletions Source/Initialization/ERF_init_from_metgrid.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -862,6 +862,8 @@ init_base_state_from_metgrid (const bool use_moisture,
auto new_data = state_fab.array();
auto const new_z = z_phys_cc_fab.const_array();

const int maxiter = 10;
const amrex::Real tol = 1.0e-12;
ParallelFor(valid_bx2d, [=] AMREX_GPU_DEVICE (int i, int j, int) noexcept
{
for (int k=0; k<=kmax; k++) {
Expand All @@ -871,9 +873,70 @@ init_base_state_from_metgrid (const bool use_moisture,
}
z_vec[kmax+1] = new_z(i,j,kmax+1);

calc_rho_p(kmax, flag_psfc_vec[0], orig_psfc(i,j,0),
Thetad_vec, Thetam_vec, Q_vec, z_vec,
Rhod_vec, Rhom_vec, Pd_vec, Pm_vec);
// Get Psurf
amrex::Real Psurf;
if (flag_psfc_vec[0] == 1) {
Psurf = orig_psfc(i,j,0);
} else {
amrex::Real t_0 = 290.0; // WRF's model_config_rec%base_temp
amrex::Real a = 50.0; // WRF's model_config_rec%base_lapse
Psurf = p_0*exp(-t_0/a+std::pow((std::pow(t_0/a, 2)-2.0*CONST_GRAV*z_vec[0]/(a*R_d)), 0.5));
}

// Iterations for the first CC point that is 1/2 dz off the surface
{
amrex::Real half_dz = z_vec[0];
amrex::Real qvf = 1.0+(R_v/R_d)*Q_vec[0];
Thetam_vec[0] = Thetad_vec[0]*qvf;
for (int iter=0; iter<maxiter; iter++) {
Pm_vec[0] = Psurf-half_dz*(Rhom_vec[0])*(1.0+Q_vec[0])*CONST_GRAV;
if (Pm_vec[0] < 0.0) Pm_vec[0] = 0.0;
Rhom_vec[0] = (p_0/(R_d*Thetam_vec[0]))*std::pow(Pm_vec[0]/p_0, iGamma);
} // it
}

// Integrate from the first CC point to the top boundary.
for (int k=1; k<=kmax; k++) {

// Assignment only for downward integration in following loop
amrex::Real qvf = 1.0+(R_v/R_d)*Q_vec[k];
Thetam_vec[k] = Thetad_vec[k]*qvf;

// Vertical grid spacing
amrex::Real dz = z_vec[k]-z_vec[k-1];

// Establish known constant
amrex::Real rho_tot_lo = Rhom_vec[k-1] * (1. + Q_vec[k-1]);
amrex::Real C = -Pm_vec[k-1] + 0.5*rho_tot_lo*CONST_GRAV*dz;

// Initial guess and residual
Pm_vec[k] = Pm_vec[k-1];
Rhom_vec[k] = getRhogivenThetaPress(Thetad_vec[k],
Pm_vec[k],
R_d/Cp_d,
Q_vec[k]);
amrex::Real rho_tot_hi = Rhom_vec[k] * (1. + Q_vec[k]);
amrex::Real F = Pm_vec[k] + 0.5*rho_tot_hi*CONST_GRAV*dz + C;

// Do iterations
if (std::abs(F)>tol) HSEutils::Newton_Raphson_hse(tol, R_d/Cp_d, dz,
CONST_GRAV, C, Thetad_vec[k],
Q_vec[k], Q_vec[k],
Pm_vec[k], Rhom_vec[k], F);
} // k

// integrate from the top back down to get dry pressure and density.
Pd_vec[kmax] = Pm_vec[kmax];
Rhod_vec[kmax] = (p_0/(R_d*Thetam_vec[kmax]))*std::pow(Pd_vec[kmax]/p_0, iGamma);
for (int k=kmax-1; k>=0; k--) {
amrex::Real dz = z_vec[k+1]-z_vec[k];
Rhod_vec[k] = Rhod_vec[k+1]; // an initial guess.
for (int iter=0; iter<maxiter; iter++) {
Pd_vec[k] = Pd_vec[k+1]+0.5*dz*(Rhod_vec[k]+Rhod_vec[k+1])*CONST_GRAV;
if (Pd_vec[k] < 0.0) Pd_vec[k] = 0.0;
Rhod_vec[k] = (p_0/(R_d*Thetam_vec[k]))*std::pow(Pd_vec[k]/p_0, iGamma);
} // it
} // k

for (int k=0; k<=kmax; k++) {
p_hse_arr(i,j,k) = Pm_vec[k];
Expand Down Expand Up @@ -966,6 +1029,8 @@ init_base_state_from_metgrid (const bool use_moisture,
auto Q_arr = (use_moisture ) ? fabs_for_bcs[it][MetGridBdyVars::QV].array() : Array4<Real>{};
auto p_hse_arr = p_hse_bcs_fab.array();

const int maxiter = 10;
const amrex::Real tol = 1.0e-12;
ParallelFor(valid_bx2d, [=] AMREX_GPU_DEVICE (int i, int j, int) noexcept
{
for (int k=0; k<=kmax; k++) {
Expand All @@ -975,9 +1040,70 @@ init_base_state_from_metgrid (const bool use_moisture,
}
z_vec[kmax+1] = new_z(i,j,kmax+1);

calc_rho_p(kmax, flag_psfc_vec[it], orig_psfc(i,j,0),
Thetad_vec, Thetam_vec, Q_vec, z_vec,
Rhod_vec, Rhom_vec, Pd_vec, Pm_vec);
// Get Psurf
amrex::Real Psurf;
if (flag_psfc_vec[it] == 1) {
Psurf = orig_psfc(i,j,0);
} else {
amrex::Real t_0 = 290.0; // WRF's model_config_rec%base_temp
amrex::Real a = 50.0; // WRF's model_config_rec%base_lapse
Psurf = p_0*exp(-t_0/a+std::pow((std::pow(t_0/a, 2)-2.0*CONST_GRAV*z_vec[0]/(a*R_d)), 0.5));
}

// Iterations for the first CC point that is 1/2 dz off the surface
{
amrex::Real half_dz = z_vec[0];
amrex::Real qvf = 1.0+(R_v/R_d)*Q_vec[0];
Thetam_vec[0] = Thetad_vec[0]*qvf;
for (int iter=0; iter<maxiter; iter++) {
Pm_vec[0] = Psurf-half_dz*(Rhom_vec[0])*(1.0+Q_vec[0])*CONST_GRAV;
if (Pm_vec[0] < 0.0) Pm_vec[0] = 0.0;
Rhom_vec[0] = (p_0/(R_d*Thetam_vec[0]))*std::pow(Pm_vec[0]/p_0, iGamma);
} // it
}

// Integrate from the first CC point to the top boundary.
for (int k=1; k<=kmax; k++) {

// Assignment only for downward integration in following loop
amrex::Real qvf = 1.0+(R_v/R_d)*Q_vec[k];
Thetam_vec[k] = Thetad_vec[k]*qvf;

// Vertical grid spacing
amrex::Real dz = z_vec[k]-z_vec[k-1];

// Establish known constant
amrex::Real rho_tot_lo = Rhom_vec[k-1] * (1. + Q_vec[k-1]);
amrex::Real C = -Pm_vec[k-1] + 0.5*rho_tot_lo*CONST_GRAV*dz;

// Initial guess and residual
Pm_vec[k] = Pm_vec[k-1];
Rhom_vec[k] = getRhogivenThetaPress(Thetad_vec[k],
Pm_vec[k],
R_d/Cp_d,
Q_vec[k]);
amrex::Real rho_tot_hi = Rhom_vec[k] * (1. + Q_vec[k]);
amrex::Real F = Pm_vec[k] + 0.5*rho_tot_hi*CONST_GRAV*dz + C;

// Do iterations
if (std::abs(F)>tol) HSEutils::Newton_Raphson_hse(tol, R_d/Cp_d, dz,
CONST_GRAV, C, Thetad_vec[k],
Q_vec[k], Q_vec[k],
Pm_vec[k], Rhom_vec[k], F);
} // k

// integrate from the top back down to get dry pressure and density.
Pd_vec[kmax] = Pm_vec[kmax];
Rhod_vec[kmax] = (p_0/(R_d*Thetam_vec[kmax]))*std::pow(Pd_vec[kmax]/p_0, iGamma);
for (int k=kmax-1; k>=0; k--) {
amrex::Real dz = z_vec[k+1]-z_vec[k];
Rhod_vec[k] = Rhod_vec[k+1]; // an initial guess.
for (int iter=0; iter<maxiter; iter++) {
Pd_vec[k] = Pd_vec[k+1]+0.5*dz*(Rhod_vec[k]+Rhod_vec[k+1])*CONST_GRAV;
if (Pd_vec[k] < 0.0) Pd_vec[k] = 0.0;
Rhod_vec[k] = (p_0/(R_d*Thetam_vec[k]))*std::pow(Pd_vec[k]/p_0, iGamma);
} // it
} // k

for (int k=0; k<=kmax; k++) {
p_hse_arr(i,j,k) = Pm_vec[k];
Expand Down

0 comments on commit e2cf47a

Please sign in to comment.