ERF_HSEUtils.H

Go to the documentation of this file.

#ifndef ERF_HSEUTIL_H_
#define ERF_HSEUTIL_H_
 
#include <AMReX_Math.H>
 
#include "ERF_MicrophysicsUtils.H"
#include "ERF_Constants.H"
#include "ERF_EOS.H"
 
/**
 * Utility functions for calculating a hydrostatic equilibrium (HSE) base state
*/
 
namespace HSEutils
{
    using namespace amrex;
 
    // to control Newton iterations in init_isentropic_hse*
    const int MAX_ITER = 10;
#ifdef AMREX_USE_FLOAT
    const amrex::Real TOL = Real(1.e-4);
#else
    const amrex::Real TOL = Real(1.e-8);
#endif
 
    /**
     * Function to calculate the hydrostatic density and pressure
     * from Newton-Raphson iterations at constant Th & Qv
     *
     * @param[in]  m_tol      iteration tolerance
     * @param[in]  RdOCp      Rd/Cp
     * @param[out] dz         change in vertical height
     * @param[out] g          magnitude of gravity
     * @param[in]  C          sum of known terms in HSE balance
     * @param[in]  Th         theta at the current cell center
     * @param[in] qt          total moisture (non-precip and precip)
     * @param[in] qv          water vapor
     * @param[in] P           pressure at cell center
     * @param[in] rd          dry density at cell center
     * @param[in] F           starting residual of non-linear eq
     * @param[in] maintain_Th maintain theta or temperature?
    */
    AMREX_GPU_HOST_DEVICE
    AMREX_FORCE_INLINE
    void
    Newton_Raphson_hse (const Real& m_tol,
                        const Real& RdoCp,
                        const Real& dz,
                        const Real& g,
                        const Real& C,
                        const Real& Th,
                        const Real& T,
                        const Real& qt,
                        const Real& qv,
                        Real& P,
                        Real& rd,
                        Real& F,
                        const bool& maintain_Th)
    {
        // A non-finite residual (or a non-positive initial guess) would make the
        // convergence test below evaluate to false on the very first pass, so we
        // would silently return the initial guess and leave, e.g., an all-zero
        // base state.  Trap that here rather than propagating it downstream.
        if (!(P > amrex::Real(0)) || !(rd > amrex::Real(0)) || !amrex::Math::isfinite(F)) {
            AMREX_DEVICE_PRINTF("ERROR: HSE Newton started from an invalid state: P = %e, rd = %e, F = %e\n",
                                double(P), double(rd), double(F));
            AMREX_DEVICE_PRINTF("       Check the surface pressure/temperature used to seed the integration%s\n", "");
        }
        AMREX_ALWAYS_ASSERT(P > amrex::Real(0) && rd > amrex::Real(0) && amrex::Math::isfinite(F));
 
        int iter=0;
        int max_iter=20;
        // Written as !(|F| <= tol) so that a NaN residual enters the loop instead
        // of being mistaken for convergence
        while (!(std::abs(F)<=m_tol) && iter<max_iter) {
            // Compute change in pressure
            Real dRdP = (maintain_Th) ? iGamma * rd / P : rd / P;
            Real dFdp = one + myhalf*dRdP*g*dz;
            P -= F/dFdp;
            if (P < Real(1.e3)) amrex::Warning("P < 1000 [Pa]; Domain height may be too large...");
            AMREX_ALWAYS_ASSERT(P > amrex::Real(0));
 
            // Diagnose density and residual
            rd = (maintain_Th) ? getRhogivenThetaPress(Th, P, RdoCp, qv) :
                                 getRhogivenTandPress(T , P, qv);
            AMREX_ALWAYS_ASSERT(rd > amrex::Real(0));
            Real r_tot = rd * (one + qt);
            F = P + myhalf*r_tot*g*dz + C;
            ++iter;
        }
        if (iter>=max_iter) {
            AMREX_DEVICE_PRINTF("WARNING: HSE Newton iterations did not converge to tolerance!\n%s", "");
            AMREX_DEVICE_PRINTF("HSE Newton tol: %e %e\n",F,m_tol);
        }
    }
 
    /**
     * Function to calculate the hydrostatic density and pressure with constant dz
     *
     * @param[in]  r_sfc     surface density
     * @param[in]  theta     surface potential temperature
     * @param[out] r         hydrostatically balanced density profile
     * @param[out] p         hydrostatically balanced pressure profile
     * @param[in]  dz        vertical grid spacing (constant)
     * @param[in]  z_bot     absolute height of the bottom of the domain
     * @param[in]  klo       z-index corresponding to the small end of the domain
     * @param[in]  khi       z-index corresponding to the big   end of the domain
    */
    AMREX_GPU_HOST_DEVICE
    AMREX_FORCE_INLINE
    void
    init_isentropic_hse_constant_dz (const amrex::Real& r_sfc,
                                     const amrex::Real& theta,
                                     amrex::Real* r,
                                     amrex::Real* p,
                                     const amrex::Real& dz,
                                     const amrex::Real& z_bot,
                                     const int klo, const int khi)
    {
      int kstart;
 
      // r_sfc is the initial guess for the density at klo; p_0 is the pressure at z = 0
 
      // Absolute height of the center of cell klo
      const Real hz = z_bot + myhalf*dz;
 
      // Initial guess
      r[klo] = r_sfc;
      p[klo] = p_0 - hz * r[klo] * CONST_GRAV;
 
      if (klo == 0)
      {
          kstart = 1;
 
          // We do a Newton iteration to satisfy the EOS & HSE (with constant theta)
          bool converged_hse = false;
          Real p_hse;
          Real p_eos;
 
          for (int iter = 0; iter < MAX_ITER && !converged_hse; iter++)
          {
              p_hse = p_0 - hz * r[klo] * CONST_GRAV;
              p_eos = getPgivenRTh(r[klo]*theta);
 
              Real A = p_hse - p_eos;
 
              Real dpdr = getdPdRgivenConstantTheta(r[klo],theta);
 
              Real drho = A / (dpdr + hz * CONST_GRAV);
 
              r[klo] = r[klo] + drho;
              p[klo] = getPgivenRTh(r[klo]*theta);
 
              if (std::abs(drho) < TOL)
              {
                  converged_hse = true;
                  break;
              }
          }
 
          // if (!converged_hse) amrex::Print() << "DOING ITERATIONS AT K = " << klo << std::endl;
          // if (!converged_hse) amrex::Error("Didn't converge the iterations in init");
      } else {
          kstart = klo; // because we need to solve for r[klo] here; r[klo-1] is what was passed in r_sfc
          r[klo-1] = r_sfc; // r_sfc is really the interpolated density in the ghost cell in this case
          p[klo-1] = getPgivenRTh(r[klo-1]*theta);
      }
 
      // To get values at k > 0 we do a Newton iteration to satisfy the EOS (with constant theta) and
      for (int k = kstart; k <= khi; k++)
      {
          // To get values at k > 0 we do a Newton iteration to satisfy the EOS (with constant theta) and
          // to discretely satisfy HSE -- here we assume spatial_order = 2 -- we can generalize this later if needed
          bool converged_hse = false;
 
          r[k] = r[k-1];
 
          Real p_eos = getPgivenRTh(r[k]*theta);
          Real p_hse;
 
          for (int iter = 0; iter < MAX_ITER && !converged_hse; iter++)
          {
              Real r_avg = myhalf * (r[k-1]+r[k]);
              p_hse = p[k-1] -  dz * r_avg * CONST_GRAV;
              p_eos = getPgivenRTh(r[k]*theta);
 
              Real A = p_hse - p_eos;
 
              Real dpdr = getdPdRgivenConstantTheta(r[k],theta);
              // Gamma * p_0 * std::pow( (R_d * theta / p_0), Gamma) * std::pow(r[k], Gamma-one) ;
 
              Real drho = A / (dpdr + dz * CONST_GRAV);
 
              r[k] = r[k] + drho;
              p[k] = getPgivenRTh(r[k]*theta);
 
              if (std::abs(drho) < TOL * r[k-1])
              {
                  converged_hse = true;
                  // amrex::Print() << " converged " << std::endl;
                  break;
              }
          }
 
          // if (!converged_hse) amrex::Print() << "DOING ITERATIONS AT K = " << k << std::endl;
          // if (!converged_hse) amrex::Error("Didn't converge the iterations in init");
      }
      r[khi+1] = r[khi];
    }
 
    /**
     * Function to calculate the hydrostatic density and pressure with stretched dz
     *
     * @param[in]  r_sfc     surface density
     * @param[in]  theta     surface potential temperature
     * @param[out] r         hydrostatically balanced density profile
     * @param[out] p         hydrostatically balanced pressure profile
     * @param[in]  stretched_dz vertical grid spacing (stretched)
     * @param[in]  z_bot     absolute height of the bottom of the domain
     * @param[in]  klo       z-index corresponding to the small end of the domain
     * @param[in]  khi       z-index corresponding to the big   end of the domain
    */
    AMREX_GPU_HOST_DEVICE
    AMREX_FORCE_INLINE
    void
    init_isentropic_hse_stretched_dz (const amrex::Real& r_sfc,
                                      const amrex::Real& theta,
                                      amrex::Real* r,
                                      amrex::Real* p,
                                      const amrex::Real* stretched_dz,
                                      const amrex::Real& z_bot,
                                      const int klo, const int khi)
    {
      int kstart;
 
      // r_sfc is the initial guess for the density at klo; p_0 is the pressure at z = 0
 
      // Absolute height of the center of cell klo
      const Real hz = z_bot + myhalf*stretched_dz[klo];
 
      // Initial guess
      r[klo] = r_sfc;
      p[klo] = p_0 - hz * r[klo] * CONST_GRAV;
 
      if (klo == 0)
      {
          kstart = 1;
 
          // We do a Newton iteration to satisfy the EOS & HSE (with constant theta)
          bool converged_hse = false;
          Real p_hse;
          Real p_eos;
 
          for (int iter = 0; iter < MAX_ITER && !converged_hse; iter++)
          {
              p_hse = p_0 - hz * r[klo] * CONST_GRAV;
              p_eos = getPgivenRTh(r[klo]*theta);
 
              Real A = p_hse - p_eos;
 
              Real dpdr = getdPdRgivenConstantTheta(r[klo],theta);
 
              Real drho = A / (dpdr + hz * CONST_GRAV);
 
              r[klo] = r[klo] + drho;
              p[klo] = getPgivenRTh(r[klo]*theta);
 
              if (std::abs(drho) < TOL)
              {
                  converged_hse = true;
                  break;
              }
          }
 
          // if (!converged_hse) amrex::Print() << "DOING ITERATIONS AT K = " << klo << std::endl;
          // if (!converged_hse) amrex::Error("Didn't converge the iterations in init");
      } else {
          kstart = klo; // because we need to solve for r[klo] here; r[klo-1] is what was passed in r_sfc
          r[klo-1] = r_sfc; // r_sfc is really the interpolated density in the ghost cell in this case
          p[klo-1] = getPgivenRTh(r[klo-1]*theta);
      }
 
      // To get values at k > 0 we do a Newton iteration to satisfy the EOS (with constant theta) and
      for (int k = kstart; k <= khi; k++)
      {
          // To get values at k > 0 we do a Newton iteration to satisfy the EOS (with constant theta) and
          // to discretely satisfy HSE -- here we assume spatial_order = 2 -- we can generalize this later if needed
          bool converged_hse = false;
 
          r[k] = r[k-1];
 
          Real p_eos = getPgivenRTh(r[k]*theta);
          Real p_hse;
 
          for (int iter = 0; iter < MAX_ITER && !converged_hse; iter++)
          {
              Real  r_avg = myhalf * (r[k-1]+r[k]);
              Real dz_avg = myhalf * (stretched_dz[k-1] + stretched_dz[k]);
              p_hse = p[k-1] -  dz_avg * r_avg * CONST_GRAV;
              p_eos = getPgivenRTh(r[k]*theta);
 
              Real A = p_hse - p_eos;
 
              Real dpdr = getdPdRgivenConstantTheta(r[k],theta);
              // Gamma * p_0 * std::pow( (R_d * theta / p_0), Gamma) * std::pow(r[k], Gamma-one) ;
 
              Real drho = A / (dpdr + dz_avg * CONST_GRAV);
 
              r[k] = r[k] + drho;
              p[k] = getPgivenRTh(r[k]*theta);
 
              if (std::abs(drho) < TOL * r[k-1])
              {
                  converged_hse = true;
                  // amrex::Print() << " converged " << std::endl;
                  break;
              }
          }
 
          // if (!converged_hse) amrex::Print() << "DOING ITERATIONS AT K = " << k << std::endl;
          // if (!converged_hse) amrex::Error("Didn't converge the iterations in init");
      }
      r[khi+1] = r[khi];
    }
 
 
    /**
     * Function to calculate the hydrostatic density and pressure over terrain
     *
     * @param[in]  i         x-index
     * @param[in]  j         y-index
     * @param[in]  r_sfc     surface density
     * @param[in]  theta     surface potential temperature
     * @param[out] r         hydrostatically balanced density profile
     * @param[out] p         hydrostatically balanced pressure profile
     * @param[in]  z_cc      cell-center heights
     * @param[in]  khi       z-index corresponding to the big end of the domain
    */
    AMREX_GPU_HOST_DEVICE
    AMREX_FORCE_INLINE
    void
    init_isentropic_hse_terrain (int i,
                                 int j,
                                 const amrex::Real& r_sfc,
                                 const amrex::Real& theta,
                                 amrex::Real* r,
                                 amrex::Real* p,
                                 const amrex::Array4<amrex::Real const> z_cc,
                                 const int& klo, const int& khi)
    {
        int kstart;
 
        if (klo == 0)  {
            //
            // r_sfc is the initial guess for the density in the bottom cell; p_0 is the
            // pressure at z = 0, so the bottom cell's pressure comes from its absolute height
            //
            // Initial guess
            int k0 = 0;
 
            // Where we start the lower iteration
            kstart = 1;
 
            // Absolute height of the center of the bottom cell (not its height above the ground)
            Real myhalf_dz = z_cc(i,j,k0);
            r[k0] = r_sfc;
            p[k0] = p_0 - myhalf_dz * r[k0] * CONST_GRAV;
            {
                // We do a Newton iteration to satisfy the EOS & HSE (with constant theta)
                bool converged_hse = false;
                Real p_hse;
                Real p_eos;
 
                for (int iter = 0; iter < MAX_ITER && !converged_hse; iter++)
                {
                    p_hse = p_0 - myhalf_dz * r[k0] * CONST_GRAV;
                    p_eos = getPgivenRTh(r[k0]*theta);
 
                    Real A = p_hse - p_eos;
 
                    Real dpdr = getdPdRgivenConstantTheta(r[k0],theta);
 
                    Real drho = A / (dpdr + myhalf_dz * CONST_GRAV);
 
                    r[k0] = r[k0] + drho;
                    p[k0] = getPgivenRTh(r[k0]*theta);
 
                    if (std::abs(drho) < TOL)
                    {
                        converged_hse = true;
                        break;
                    }
                }
 
                //if (!converged_hse) amrex::Print() << "DOING ITERATIONS AT K = " << k0 << std::endl;
                //if (!converged_hse) amrex::Error("Didn't converge the iterations in init");
            }
        } else {
            kstart = klo; // because we need to solve for r[klo] here; r[klo-1] is what was passed in r_sfc
            r[klo-1] = r_sfc; // r_sfc is really the interpolated density in the ghost cell in this case
            p[klo-1] = getPgivenRTh(r[klo-1]*theta);
        }
 
         // To get values at k > 0 we do a Newton iteration to satisfy the EOS (with constant theta) and
         for (int k = kstart; k <= khi; k++)
         {
             // To get values at k > 0 we do a Newton iteration to satisfy the EOS (with constant theta) and
             // to discretely satisfy HSE -- here we assume spatial_order = 2 -- we can generalize this later if needed
             bool converged_hse = false;
 
             Real dz_loc = (z_cc(i,j,k) - z_cc(i,j,k-1));
 
             r[k] = r[k-1];
 
             Real p_eos = getPgivenRTh(r[k]*theta);
             Real p_hse;
 
             for (int iter = 0; iter < MAX_ITER && !converged_hse; iter++)
             {
                 p_hse = p[k-1] -  dz_loc * myhalf * (r[k-1]+r[k]) * CONST_GRAV;
                 p_eos = getPgivenRTh(r[k]*theta);
 
                 Real A = p_hse - p_eos;
 
                 Real dpdr = getdPdRgivenConstantTheta(r[k],theta);
 
                 Real drho = A / (dpdr + dz_loc * CONST_GRAV);
 
                 r[k] = r[k] + drho;
                 p[k] = getPgivenRTh(r[k]*theta);
 
                 if (std::abs(drho) < TOL * r[k-1])
                 {
                     converged_hse = true;
                     //amrex::Print() << " converged " << std::endl;
                     break;
                 }
             }
 
             // if (!converged_hse) amrex::Print() << "DOING ITERATIONS AT K = " << k << std::endl;
             // if (!converged_hse) amrex::Error("Didn't converge the iterations in init");
        }
    }
 
AMREX_FORCE_INLINE
AMREX_GPU_HOST_DEVICE
Real compute_saturation_pressure (const Real T_b, const bool use_empirical)
{
    Real p_s = erf_esatw(T_b,use_empirical);
    return p_s * Real(100.0);
}
 
AMREX_FORCE_INLINE
AMREX_GPU_HOST_DEVICE
Real compute_relative_humidity (const Real p_b, const Real T_b, const bool use_empirical,
                                const int which_zone, const Real scaled_height)
{
    if (which_zone > 0) {
        if (which_zone == 1) { // z <= height
            Real p_s = compute_saturation_pressure(T_b, use_empirical);
            Real q_s = RdoRv*p_s/(p_b - p_s);
            return Real(0.014)/q_s;
        } else if (which_zone == 2) { // z > height and z <= z_tr; scaled_height = z/z_tr
            return one - Real(0.75)*std::pow(scaled_height,Real(1.25));
        } else { // z > z_tr
            return fourth;
        }
    } else {
        return one;
    }
}
 
AMREX_FORCE_INLINE
AMREX_GPU_HOST_DEVICE
Real vapor_mixing_ratio (const Real p_b, const Real T_b, const Real RH, const bool use_empirical,
                         int which_zone)
{
    Real p_s = compute_saturation_pressure(T_b,use_empirical);
    Real p_v = compute_vapor_pressure(p_s, RH);
    Real q_v = RdoRv*p_v/(p_b - p_v);
 
    if (which_zone == 1) { // z <= height
        return Real(0.014);
    } else {
        return q_v;
    }
}
 
AMREX_FORCE_INLINE
AMREX_GPU_HOST_DEVICE
Real compute_F_for_temp_in_zone (const Real T_b, const Real p_b, const Real q_t, const Real eq_pot_temp,
                                 const bool use_empirical, const int which_zone, const Real scaled_height)
{
    Real fac = Cp_d + Cp_l*q_t;
    Real RH  = compute_relative_humidity(p_b, T_b, use_empirical, which_zone, scaled_height);
    Real q_v = vapor_mixing_ratio(p_b, T_b, RH, use_empirical, which_zone);
    Real p_s = compute_saturation_pressure(T_b, use_empirical);
    Real p_v = compute_vapor_pressure(p_s, RH);
    // Relative-humidity factor H^(-r_v R_v / fac) of Emanuel (1994, eq. 4.5.11);
    // exactly 1 in saturated zones (RH = 1)
    Real H = amrex::max(Real(1.e-10), amrex::min(Real(1.0), RH));
    return eq_pot_temp - T_b*std::pow((p_b - p_v)/p_0, -R_d/fac)
                       *std::exp(-(q_v*R_v/fac)*std::log(H))
                       *std::exp(L_v*q_v/(fac*T_b));
}
 
AMREX_FORCE_INLINE
AMREX_GPU_HOST_DEVICE
Real compute_temperature (const Real p_b, const Real q_t, const Real eq_pot_temp, const bool use_empirical,
                          const int which_zone, const Real scaled_height)
{
    Real T_b = Real(200.0), delta_T; // Initial guess
 
#ifdef AMREX_USE_FLOAT
    Real eps = Real(1.e-4) * T_b;
#else
    Real eps = Real(1.e-10) * T_b;
#endif
    for (int iter=0; iter<20; iter++)
    {
        Real F         = compute_F_for_temp_in_zone(T_b    , p_b, q_t, eq_pot_temp, use_empirical, which_zone, scaled_height);
        Real F_plus_dF = compute_F_for_temp_in_zone(T_b+eps, p_b, q_t, eq_pot_temp, use_empirical, which_zone, scaled_height);
        Real F_prime   = (F_plus_dF - F)/eps;
        delta_T = -F/F_prime;
        T_b = T_b + delta_T;
    }
 
    if (std::fabs(delta_T) > TOL * T_b) {
         amrex::Abort("Newton Raphson for temperature could not converge");
    }
 
    return T_b;
}
 
AMREX_FORCE_INLINE
AMREX_GPU_HOST_DEVICE
Real compute_dewpoint_temperature (const Real T_b, const Real RH)
{
    Real T_dp, gamma, T;
    T = T_b - Real(273.15);
 
    Real b = Real(18.678), c = Real(257.14), d = Real(234.5);
    gamma = std::log(RH*std::exp((b - T/d)*T/(c + T)));
 
    T_dp = c*gamma/(b - gamma);
 
    return T_dp;
}
 
AMREX_FORCE_INLINE
AMREX_GPU_HOST_DEVICE
Real compute_theta (const Real scaled_height, const Real theta_0,
                    const Real theta_tr, const Real z_tr, const Real T_tr)
{
    if(scaled_height <= one) {
        return theta_0 + (theta_tr - theta_0)*std::pow(scaled_height,Real(1.25));
    } else {
        return theta_tr*std::exp(CONST_GRAV/(Cp_d*T_tr)*z_tr*(scaled_height - one));
    }
}
 
AMREX_FORCE_INLINE
AMREX_GPU_HOST_DEVICE
void compute_rho (const Real& pressure, Real& theta, Real& rho, Real& q_v, Real& T_dp, Real& T_b,
                  const Real q_t, const Real eq_pot_temp, const bool use_empirical, const int which_zone,
                  const Real scaled_height, const bool T_from_theta,
                  const Real theta_0, const Real theta_tr, const Real z_tr, const Real T_tr)
{
 
    if (T_from_theta) {
        theta   = compute_theta(scaled_height, theta_0, theta_tr, z_tr, T_tr);
        T_b     = getTgivenPandTh(pressure, theta, RdoCp);
    } else {
        T_b     = compute_temperature(pressure, q_t, eq_pot_temp, use_empirical, which_zone, scaled_height);
        theta   = getThgivenTandP(T_b, pressure, RdoCp);
    }
 
    Real RH = compute_relative_humidity(pressure, T_b, use_empirical, which_zone, scaled_height);
 
    q_v     = vapor_mixing_ratio(pressure, T_b, RH, use_empirical, which_zone);
 
    rho     = getRhogivenTandPress(T_b, pressure, q_v);
 
    if (T_from_theta) {
        rho *= (one + q_v);
    } else {
        rho *= (one + q_t);
    }
 
    T_dp    = compute_dewpoint_temperature(T_b, RH);
}
 
AMREX_FORCE_INLINE
AMREX_GPU_HOST_DEVICE
Real compute_F (const Real& p_k, const Real& p_k_minus_1, Real &theta_k, Real& rho_k, Real& q_v_k,
                Real& T_dp, Real& T_b, const Real& dz, const Real& rho_k_minus_1, const Real q_t,
                const Real eq_pot_temp, const bool use_empirical,
                const int which_zone, const Real scaled_height, const bool T_from_theta,
                const Real theta_0, const Real theta_tr, const Real z_tr, const Real T_tr)
{
    Real F;
 
    compute_rho(p_k, theta_k, rho_k, q_v_k, T_dp, T_b, q_t, eq_pot_temp, use_empirical, which_zone, scaled_height,
                T_from_theta, theta_0, theta_tr, z_tr, T_tr);
 
    if(rho_k_minus_1 == amrex::Real(0)) // This loop is for the first point above the ground
    {
        F = p_k - p_k_minus_1 + rho_k*CONST_GRAV*dz/two;
    }
    else
    {
        F = p_k - p_k_minus_1 + myhalf * (rho_k + rho_k_minus_1)*CONST_GRAV*dz;
    }
 
    return F;
}
 
AMREX_FORCE_INLINE
AMREX_GPU_HOST_DEVICE
Real compute_p_k_in_zone (Real &p_k, const Real p_k_minus_1, Real &theta_k, Real& rho_k, Real& q_v_k,
                          Real& T_dp, Real& T_b, const Real dz, const Real rho_k_minus_1, const Real q_t,
                          const Real eq_pot_temp, const bool use_empirical,
                          const int which_zone, const Real scaled_height, const bool T_from_theta,
                          const Real theta_0, const Real theta_tr, const Real z_tr, const Real T_tr)
{
    Real delta_p_k;
 
    // Perturbation used to form the difference quotient for dF/dp.  It has to
    // be scaled by p_k rather than fixed in Pa: the pressures here are O(1e5),
    // and in single precision one ulp at that magnitude is already ~1e-2 Pa,
    // so a fixed 1e-4 Pa perturbation rounds away entirely -- p_k + eps == p_k,
    // the difference quotient is zero, and the iteration returns NaN.  The
    // double-precision perturbation is left exactly as it was so that results
    // in double are unchanged.
#ifdef AMREX_USE_FLOAT
    Real eps = Real(1.e-5) * std::abs(p_k);
#else
    Real eps = Real(1e-10);
#endif
 
    for(int iter=0; iter<20; iter++)
    {
        Real F         = compute_F(p_k      , p_k_minus_1, theta_k, rho_k, q_v_k, T_dp, T_b, dz, rho_k_minus_1,
                                   q_t, eq_pot_temp, use_empirical, which_zone, scaled_height, T_from_theta,
                                   theta_0, theta_tr, z_tr, T_tr);
        Real F_plus_dF = compute_F(p_k+eps, p_k_minus_1, theta_k, rho_k, q_v_k, T_dp, T_b, dz, rho_k_minus_1,
                                   q_t, eq_pot_temp, use_empirical, which_zone, scaled_height, T_from_theta,
                                   theta_0, theta_tr, z_tr, T_tr);
        Real F_prime   = (F_plus_dF - F)/eps;
        delta_p_k = -F/F_prime;
        p_k            = p_k + delta_p_k;
    }
 
    // Convergence is measured on the last Newton step relative to p_k, for the
    // same reason: an absolute tolerance in Pa that is smaller than an ulp of
    // p_k can never be met in single precision.  The double-precision value is
    // set so that p_rel_tol * p_k reproduces the previous absolute tolerance of
    // 1e-8 at the O(1e5) Pa pressures this is used at.
#ifdef AMREX_USE_FLOAT
    const Real p_rel_tol = Real(1.e-5);
#else
    const Real p_rel_tol = Real(1.e-13);
#endif
    if (std::fabs(delta_p_k) > p_rel_tol * std::abs(p_k)) {
        amrex::Abort("Newton Raphson for pressure could not converge");
    }
 
    return p_k;
}
 
/**
 * Column integration shared by the two versions of init_isentropic_hse_no_terrain.
 *
 * @param[in] z_cc  z_cc(k) is the height of the centre of cell k above the surface.
 * @param[in] dz_cc dz_cc(k) is the distance used for the hydrostatic step to cell k:
 *                  for k = 0, twice the absolute height of the centre of cell 0 (half of
 *                  it is taken), so that the column starts from p_0 at z = 0, as in the
 *                  other HSE initializers; for k > 0, the distance between the centres
 *                  of cells k-1 and k.
 *
 * The other arguments are those of init_isentropic_hse_no_terrain.
 */
template <typename ZccFunc, typename DzFunc>
AMREX_FORCE_INLINE
AMREX_GPU_HOST_DEVICE
void
init_isentropic_hse_no_terrain_column (Real *theta, Real* r, Real* p, Real *q_v,
                                       ZccFunc const& z_cc, DzFunc const& dz_cc,
                                       const int& khi, const Real q_t,
                                       const Real eq_pot_temp, const bool use_empirical,
                                       const bool T_from_theta,
                                       const Real z_tr_1, const Real z_tr_2,
                                       const Real theta_0, const Real theta_tr, const Real T_tr)
{
    // theta_0, theta_tr and T_tr are used only when T_from_theta is set. Their
    // defaults are still physical temperatures: compute_theta is inlined into
    // the Newton loops and optimised builds evaluate its arm speculatively, so
    // zeros would raise an FPE (g/(Cp_d*0), then 0*exp(inf)) on a discarded value.
    Real T_b, T_dp;
 
    int which_zone = -1;
    Real scaled_height = amrex::Real(0);
 
    // Compute the quantities at the first cell center, integrating from p_0 at z = 0
    p[0] = p_0;
    if (z_tr_1 > amrex::Real(0)) {
        Real z = z_cc(0);
        if (z <= z_tr_1) {
            which_zone = 1;
        } else if (z <= z_tr_2) {
            which_zone = 2;
        } else {
            which_zone = 3;
        }
        scaled_height = z/z_tr_2;
    }
 
    compute_p_k_in_zone(p[0], p_0, theta[0], r[0], q_v[0], T_dp, T_b, dz_cc(0), amrex::Real(0), q_t, eq_pot_temp,
                        use_empirical, which_zone, scaled_height, T_from_theta,
                        theta_0, theta_tr, z_tr_2, T_tr);
 
    for (int k=1; k<=khi; k++)
    {
        p[k] = p[k-1];
 
        if (z_tr_1 > amrex::Real(0)) {
            Real z = z_cc(k);
            if (z <= z_tr_1) {
                which_zone = 1;
            } else if (z <= z_tr_2) {
                which_zone = 2;
            } else {
                which_zone = 3;
            }
            scaled_height = z/z_tr_2;
        }
 
        compute_p_k_in_zone(p[k], p[k-1], theta[k], r[k], q_v[k], T_dp, T_b, dz_cc(k), r[k-1], q_t, eq_pot_temp,
                            use_empirical, which_zone, scaled_height, T_from_theta,
                            theta_0, theta_tr, z_tr_2, T_tr);
    }
 
    r[khi+1] = r[khi];
}
 
/**
 * Initialize an isentropic hydrostatic equilibrium base state without terrain.
 *
 * @param[out] theta Potential temperature profile.
 * @param[out] r Density profile.
 * @param[out] p Pressure profile.
 * @param[out] q_v Water vapor mixing ratio profile.
 * @param[in] dz Vertical grid spacing.
 * @param[in] z_bot Absolute height of the bottom of the domain (geometry.prob_lo[2]).
 * @param[in] khi Top index of the domain.
 * @param[in] q_t Total water mixing ratio.
 * @param[in] eq_pot_temp Equilibrium potential temperature.
 * @param[in] use_empirical Use empirical saturation formulas.
 * @param[in] T_from_theta Whether temperature is derived from potential temperature.
 * @param[in] z_tr_1 First transition height.
 * @param[in] z_tr_2 Second transition height.
 * @param[in] theta_0 Surface potential temperature.
 * @param[in] theta_tr Potential temperature at transition height.
 * @param[in] T_tr Temperature at transition height.
 */
AMREX_FORCE_INLINE
AMREX_GPU_HOST_DEVICE
void
init_isentropic_hse_no_terrain (Real *theta, Real* r, Real* p, Real *q_v,
                                const Real& dz, const Real& z_bot, const int& khi, const Real q_t,
                                const Real eq_pot_temp, const bool use_empirical,
                                const bool T_from_theta = false,
                                const Real z_tr_1 = -one, const Real z_tr_2 = -one,
                                const Real theta_0 = amrex::Real(300), const Real theta_tr = amrex::Real(300), const Real T_tr = amrex::Real(300))
{
    init_isentropic_hse_no_terrain_column(theta, r, p, q_v,
        [=] (int k) { return (k == 0) ? myhalf * dz : (k+myhalf) * dz; },
        [=] (int k) { return (k == 0) ? two * (z_bot + myhalf * dz) : dz; },
        khi, q_t, eq_pot_temp, use_empirical, T_from_theta,
        z_tr_1, z_tr_2, theta_0, theta_tr, T_tr);
}
 
/**
 * As above, for a column whose staggered levels are given, e.g. a stretched mesh.
 *
 * @param[in] zlevels Staggered levels zlevels[0..khi+1]; heights are measured from zlevels[0],
 *                    except that the column starts from p_0 at z = 0.
 *
 * The other arguments are those of the version with a uniform dz.
 */
AMREX_FORCE_INLINE
AMREX_GPU_HOST_DEVICE
void
init_isentropic_hse_no_terrain (Real *theta, Real* r, Real* p, Real *q_v,
                                const Real* zlevels, const int& khi, const Real q_t,
                                const Real eq_pot_temp, const bool use_empirical,
                                const bool T_from_theta = false,
                                const Real z_tr_1 = -one, const Real z_tr_2 = -one,
                                const Real theta_0 = amrex::Real(300), const Real theta_tr = amrex::Real(300), const Real T_tr = amrex::Real(300))
{
    auto z_cc = [=] (int k) { return myhalf * (zlevels[k] + zlevels[k+1]) - zlevels[0]; };
    init_isentropic_hse_no_terrain_column(theta, r, p, q_v,
        z_cc,
        [=] (int k) { return (k == 0) ? zlevels[0] + zlevels[1] : z_cc(k) - z_cc(k-1); },
        khi, q_t, eq_pot_temp, use_empirical, T_from_theta,
        z_tr_1, z_tr_2, theta_0, theta_tr, T_tr);
}
 
} // namespace
 
#endif