ERF_MakeFastCoeffs.cpp File Reference

#include <AMReX.H>
#include "ERF_Constants.H"
#include <ERF_TI_fast_headers.H>
#include <ERF_TerrainMetrics.H>

Include dependency graph for ERF_MakeFastCoeffs.cpp:

Functions
void	make_fast_coeffs (int, MultiFab &fast_coeffs, Vector< MultiFab > &S_stage_data, const MultiFab &S_stage_prim, const MultiFab &pi_stage, const amrex::Geometry geom, bool l_use_moisture, MeshType mesh_type, Real gravity, Real c_p, std::unique_ptr< MultiFab > &detJ_cc, const double dtau, Real beta_s, const MultiFab *z_phys_nd, bool use_variable_beta_s, Real beta_s_min, Real beta_s_max, amrex::GpuArray< ERF_BC, AMREX_SPACEDIM *2 > &phys_bc_type)

Function Documentation

make_fast_coeffs()

void make_fast_coeffs	(	int	,
		MultiFab &	fast_coeffs,
		Vector< MultiFab > &	S_stage_data,
		const MultiFab &	S_stage_prim,
		const MultiFab &	pi_stage,
		const amrex::Geometry	geom,
		bool	l_use_moisture,
		MeshType	mesh_type,
		Real	gravity,
		Real	c_p,
		std::unique_ptr< MultiFab > &	detJ_cc,
		const double	dtau,
		Real	beta_s,
		const MultiFab *	z_phys_nd,
		bool	use_variable_beta_s,
		Real	beta_s_min,
		Real	beta_s_max,
		amrex::GpuArray< ERF_BC, AMREX_SPACEDIM *2 > &	phys_bc_type
	)

Function for computing the coefficients for the tridiagonal solver used in the fast integrator (the acoustic substepping).

Parameters

[in]	level	level of refinement
[out]	fast_coeffs	the coefficients for the tridiagonal solver computed here, and beta_s on the w-faces (component 5)
[in]	S_stage_data	solution at the last stage
[in]	S_stage_prim	primitive variables (i.e. conserved variables divided by density) at the last stage
[in]	pi_stage	Exner function at the last stage
[in]	geom	Container for geometric information
[in]	l_use_moisture	Are we evolving moisture?
[in]	mesh_type	Do we have constant dz?
[in]	gravity	Magnitude of gravity
[in]	c_p	Coefficient at constant pressure
[in]	detJ_cc	Jacobian of the metric transformation at cell centers
[in]	dtau	Fast time step
[in]	beta_s	Coefficient which determines how implicit vs explicit the solve is
[in]	z_phys_nd	Height coordinate at nodes (null without terrain)
[in]	use_variable_beta_s	If true, beta_s = max(beta_s_min, min(beta_s_max, h_xi^2 + h_eta^2)) on each w-face in place of beta_s
[in]	beta_s_min	Lower bound of the spatially varying beta_s
[in]	beta_s_max	Upper bound of the spatially varying beta_s
[in]	phys_bc_type	Physical boundary condition types

{
    BL_PROFILE_VAR("make_fast_coeffs()",make_fast_coeffs);
 
    Real c_v = c_p - R_d;
 
    const GpuArray<Real, AMREX_SPACEDIM> dxInv = geom.InvCellSizeArray();
    Real dzi = dxInv[2];
 
    const Box &domain = geom.Domain();
 
    MultiFab coeff_A_mf(fast_coeffs, amrex::make_alias, 0, 1);
    MultiFab coeff_B_mf(fast_coeffs, amrex::make_alias, 1, 1);
    MultiFab coeff_C_mf(fast_coeffs, amrex::make_alias, 2, 1);
    MultiFab coeff_P_mf(fast_coeffs, amrex::make_alias, 3, 1);
    MultiFab coeff_Q_mf(fast_coeffs, amrex::make_alias, 4, 1);
    MultiFab  beta_s_mf(fast_coeffs, amrex::make_alias, 5, 1);
 
    const bool has_z_nd = (z_phys_nd != nullptr);
 
 
    // *************************************************************************
    // Set gravity as a vector
    const    Array<Real,AMREX_SPACEDIM> grav{zero, zero, -gravity};
    const GpuArray<Real,AMREX_SPACEDIM> grav_gpu{grav[0], grav[1], grav[2]};
 
    // *************************************************************************
    // Define updates in the current RK stage
    // *************************************************************************
#ifdef _OPENMP
#pragma omp parallel if (amrex::Gpu::notInLaunchRegion())
#endif
    {
 
    for ( MFIter mfi(S_stage_data[IntVars::cons],TileNoZ()); mfi.isValid(); ++mfi)
    {
        Box bx  = mfi.tilebox();
        Box tbz = surroundingNodes(bx,2);
 
        const Array4<const Real> & stage_cons = S_stage_data[IntVars::cons].const_array(mfi);
        const Array4<const Real> & prim       = S_stage_prim.const_array(mfi);
 
        const Array4<const Real>& detJ        = (mesh_type != MeshType::ConstantDz) ?
                                                detJ_cc->const_array(mfi) : Array4<const Real>{};
 
        const Array4<const Real>& pi_stage_ca = pi_stage.const_array(mfi);
 
        FArrayBox gam_fab; gam_fab.resize(surroundingNodes(bx,2),1,The_Async_Arena());
 
        auto const& coeffA_a  = coeff_A_mf.array(mfi);
        auto const& coeffB_a  = coeff_B_mf.array(mfi);
        auto const& coeffC_a  = coeff_C_mf.array(mfi);
        auto const& coeffP_a  = coeff_P_mf.array(mfi);
        auto const& coeffQ_a  = coeff_Q_mf.array(mfi);
        auto const&    gam_a  = gam_fab.array();
        auto const& beta_s_a  = beta_s_mf.array(mfi);
 
        const Array4<const Real>& z_nd = (has_z_nd) ? z_phys_nd->const_array(mfi) : Array4<const Real>{};
 
        // *********************************************************************
        // beta_s on the w-faces, which erf_substep_* uses for the time-weighting of the
        // vertical fluxes and of the vertical momentum equation
        // *********************************************************************
        ParallelFor(tbz, [=] AMREX_GPU_DEVICE (int i, int j, int k)
        {
            if (use_variable_beta_s) {
                Real h_xi  = (has_z_nd) ? Compute_h_xi_AtKface (i, j, k, dxInv, z_nd) : zero;
                Real h_eta = (has_z_nd) ? Compute_h_eta_AtKface(i, j, k, dxInv, z_nd) : zero;
                beta_s_a(i,j,k) = amrex::max(beta_s_min, amrex::min(beta_s_max, h_xi*h_xi + h_eta*h_eta));
            } else {
                beta_s_a(i,j,k) = beta_s;
            }
        });
 
        // *********************************************************************
        // *********************************************************************
        // *********************************************************************
 
        Box bx_shrunk_in_k = bx;
        int klo = tbz.smallEnd(2);
        int khi = tbz.bigEnd(2);
        bx_shrunk_in_k.setSmall(2,klo+1);
        bx_shrunk_in_k.setBig(2,khi-1);
 
        // Note that the notes use "g" to mean the magnitude of gravity, so it is positive
        // We set grav_gpu[2] to be the vector component which is negative
        // We define halfg to match the notes (which is why we take the absolute value)
        Real halfg = std::abs(myhalf * grav_gpu[2]);
 
        //Note we don't act on the bottom or top boundaries of the domain
        if (mesh_type != MeshType::ConstantDz)
        {
            ParallelFor(bx_shrunk_in_k, [=] AMREX_GPU_DEVICE (int i, int j, int k)
            {
                Real pi_c =  myhalf * (pi_stage_ca(i,j,k-1) + pi_stage_ca(i,j,k));
 
                Real     detJ_on_kface = myhalf * (detJ(i,j,k) + detJ(i,j,k-1));
                Real inv_detJ_on_kface = one / detJ_on_kface;
 
                Real qv_p = (l_use_moisture) ? prim(i,j,k  ,PrimQ1_comp) : zero;
                Real qv_q = (l_use_moisture) ? prim(i,j,k-1,PrimQ1_comp) : zero;
 
                Real Thm_hi = stage_cons(i,j,k  ,RhoTheta_comp) * (one + RvoRd*qv_p);
                Real Thm_lo = stage_cons(i,j,k-1,RhoTheta_comp) * (one + RvoRd*qv_q);
                Real Thm_grad = dzi * inv_detJ_on_kface * ( Thm_hi - Thm_lo );
 
                Real coeff_P = -Gamma * R_d * dzi * inv_detJ_on_kface * pi_c * (one + RvoRd*qv_p)
                             -  Gamma * R_d * R_d * Thm_grad * myhalf * pi_stage_ca(i,j,k  ) /
                             (  c_v * stage_cons(i,j,k  ,RhoTheta_comp) );
 
                Real coeff_Q =  Gamma * R_d * dzi * inv_detJ_on_kface * pi_c * (one + RvoRd*qv_q)
                             -  Gamma * R_d * R_d * Thm_grad * myhalf * pi_stage_ca(i,j,k-1) /
                             (  c_v * stage_cons(i,j,k-1,RhoTheta_comp) );
 
                if (l_use_moisture) {
                    Real q = myhalf * ( prim(i,j,k,PrimQ1_comp) + prim(i,j,k-1,PrimQ1_comp)
                                      + prim(i,j,k,PrimQ2_comp) + prim(i,j,k-1,PrimQ2_comp) );
                    coeff_P /= (one + q);
                    coeff_Q /= (one + q);
                }
 
                // NOTE: we store the moisture-normalized coefficients so that the explicit
                //       RHS assembled in erf_substep_* uses exactly the same vertical fast
                //       pressure gradient as the implicit tridiagonal system built below
                coeffP_a(i,j,k) = coeff_P;
                coeffQ_a(i,j,k) = coeff_Q;
 
                Real theta_t_lo  = myhalf * ( prim(i,j,k-2,PrimTheta_comp) + prim(i,j,k-1,PrimTheta_comp) );
                Real theta_t_mid = myhalf * ( prim(i,j,k-1,PrimTheta_comp) + prim(i,j,k  ,PrimTheta_comp) );
                Real theta_t_hi  = myhalf * ( prim(i,j,k  ,PrimTheta_comp) + prim(i,j,k+1,PrimTheta_comp) );
 
                // LHS for tri-diagonal system
                // NOTE: the w-equation at face k is weighted by beta_2(k), and the vertical
                //       flux at face k' by beta_2(k'), so each entry carries both
                //       (multiplied first, so a uniform beta_s rounds as beta_2 * beta_2 did)
                Real beta_2_km1 = myhalf * (one + beta_s_a(i,j,k-1));
                Real beta_2_k   = myhalf * (one + beta_s_a(i,j,k  ));
                Real beta_2_kp1 = myhalf * (one + beta_s_a(i,j,k+1));
                Real D_A = beta_2_km1 * beta_2_k   * dzi * static_cast<Real>(dtau * dtau);
                Real D_B = beta_2_k   * beta_2_k   * dzi * static_cast<Real>(dtau * dtau);
                Real D_C = beta_2_k   * beta_2_kp1 * dzi * static_cast<Real>(dtau * dtau);
                coeffA_a(i,j,k) = D_A * (one/detJ(i,j,k-1)) * ( halfg - coeff_Q * theta_t_lo );
                coeffC_a(i,j,k) = D_C * (one/detJ(i,j,k  )) * (-halfg + coeff_P * theta_t_hi );
 
                coeffB_a(i,j,k) = one + D_B * ( (coeff_Q/detJ(i,j,k-1) - coeff_P/detJ(i,j,k)) * theta_t_mid
                                      + halfg * (one/detJ(i,j,k) - one/detJ(i,j,k-1)) );
            });
 
        } else {
 
            ParallelFor(bx_shrunk_in_k, [=] AMREX_GPU_DEVICE (int i, int j, int k)
            {
                Real pi_c =  myhalf * (pi_stage_ca(i,j,k-1) + pi_stage_ca(i,j,k));
 
                Real qv_p = (l_use_moisture) ? prim(i,j,k  ,PrimQ1_comp) : zero;
                Real qv_q = (l_use_moisture) ? prim(i,j,k-1,PrimQ1_comp) : zero;
 
                Real Thm_hi = stage_cons(i,j,k  ,RhoTheta_comp) * (one + RvoRd*qv_p);
                Real Thm_lo = stage_cons(i,j,k-1,RhoTheta_comp) * (one + RvoRd*qv_q);
                Real Thm_grad = dzi * ( Thm_hi - Thm_lo );
 
                Real coeff_P = -Gamma * R_d * dzi * pi_c * (one + RvoRd*qv_p)
                             -  Gamma * R_d * R_d * Thm_grad * myhalf * pi_stage_ca(i,j,k  ) /
                             (  c_v * stage_cons(i,j,k  ,RhoTheta_comp) );
 
                Real coeff_Q =  Gamma * R_d * dzi * pi_c * (one + RvoRd*qv_q)
                             -  Gamma * R_d * R_d * Thm_grad * myhalf * pi_stage_ca(i,j,k-1) /
                             (  c_v * stage_cons(i,j,k-1,RhoTheta_comp) );
 
                if (l_use_moisture) {
                    Real q = myhalf * ( prim(i,j,k,PrimQ1_comp) + prim(i,j,k-1,PrimQ1_comp)
                                      + prim(i,j,k,PrimQ2_comp) + prim(i,j,k-1,PrimQ2_comp) );
                    coeff_P /= (one + q);
                    coeff_Q /= (one + q);
                }
 
                // NOTE: we store the moisture-normalized coefficients so that the explicit
                //       RHS assembled in erf_substep_* uses exactly the same vertical fast
                //       pressure gradient as the implicit tridiagonal system built below
                coeffP_a(i,j,k) = coeff_P;
                coeffQ_a(i,j,k) = coeff_Q;
 
                Real theta_t_lo  = myhalf * ( prim(i,j,k-2,PrimTheta_comp) + prim(i,j,k-1,PrimTheta_comp) );
                Real theta_t_mid = myhalf * ( prim(i,j,k-1,PrimTheta_comp) + prim(i,j,k  ,PrimTheta_comp) );
                Real theta_t_hi  = myhalf * ( prim(i,j,k  ,PrimTheta_comp) + prim(i,j,k+1,PrimTheta_comp) );
 
                // LHS for tri-diagonal system
                // NOTE: the w-equation at face k is weighted by beta_2(k), and the vertical
                //       flux at face k' by beta_2(k'), so each entry carries both
                //       (multiplied first, so a uniform beta_s rounds as beta_2 * beta_2 did)
                Real beta_2_km1 = myhalf * (one + beta_s_a(i,j,k-1));
                Real beta_2_k   = myhalf * (one + beta_s_a(i,j,k  ));
                Real beta_2_kp1 = myhalf * (one + beta_s_a(i,j,k+1));
                Real D_A = beta_2_km1 * beta_2_k   * dzi * static_cast<Real>(dtau * dtau);
                Real D_B = beta_2_k   * beta_2_k   * dzi * static_cast<Real>(dtau * dtau);
                Real D_C = beta_2_k   * beta_2_kp1 * dzi * static_cast<Real>(dtau * dtau);
                coeffA_a(i,j,k) = D_A * ( halfg - coeff_Q * theta_t_lo );
                coeffC_a(i,j,k) = D_C * (-halfg + coeff_P * theta_t_hi );
 
                coeffB_a(i,j,k) = one + D_B * (coeff_Q - coeff_P) * theta_t_mid;
            });
        }
 
        amrex::Box b2d = tbz; // Copy constructor
        b2d.setRange(2,0);
 
        auto const lo = amrex::lbound(bx);
        auto const hi = amrex::ubound(bx);
 
        auto const domhi = amrex::ubound(domain);
 
        {
        BL_PROFILE("make_coeffs_b2d_loop");
#ifdef AMREX_USE_GPU
        ParallelFor(b2d, [=] AMREX_GPU_DEVICE (int i, int j, int) {
 
            // If at the bottom of the grid, we will set w to a specified Dirichlet value
            coeffA_a(i,j,lo.z) =  zero;
            coeffB_a(i,j,lo.z) =  one;
            coeffC_a(i,j,lo.z) =  zero;
 
            // If at the top of the grid, we will set w to a specified Dirichlet value
            coeffA_a(i,j,hi.z+1) =  zero;
            coeffB_a(i,j,hi.z+1) =  one;
            coeffC_a(i,j,hi.z+1) =  zero;
 
            // UNLESS if at the top of the domain and the boundary is outflow,
            //     we will use a homogeneous Neumann condition
            if ( (hi.z == domhi.z) &&
                 (phys_bc_type[5] == ERF_BC::outflow or phys_bc_type[5] == ERF_BC::ho_outflow) )
            {
                coeffA_a(i,j,hi.z+1) =  -one;
            }
 
            // w = specified Dirichlet value at k = lo.z
            gam_a(i,j,lo.z) = coeffC_a(i,j,lo.z) / coeffB_a(i,j,lo.z);
            for (int k = lo.z+1; k <= hi.z+1; k++) {
                coeffB_a(i,j,k) = one / ( coeffB_a(i,j,k) - coeffA_a(i,j,k)*gam_a(i,j,k-1) );
                gam_a(i,j,k)    = coeffC_a(i,j,k) * coeffB_a(i,j,k);
            }
        });
#else
        // If at the bottom of the grid, we will set w to a specified Dirichlet value
        for (int j = lo.y; j <= hi.y; ++j) {
            AMREX_PRAGMA_SIMD
            for (int i = lo.x; i <= hi.x; ++i) {
                coeffA_a(i,j,lo.z) =  zero;
                coeffB_a(i,j,lo.z) =  one;
                coeffC_a(i,j,lo.z) =  zero;
                gam_a(i,j,lo.z)    = coeffC_a(i,j,lo.z) / coeffB_a(i,j,lo.z);
            }
        }
        for (int j = lo.y; j <= hi.y; ++j) {
            AMREX_PRAGMA_SIMD
            for (int i = lo.x; i <= hi.x; ++i) {
 
                // If at the top of the grid, we will set w to a specified Dirichlet value
                coeffA_a(i,j,hi.z+1) =  zero;
                coeffB_a(i,j,hi.z+1) =  one;
                coeffC_a(i,j,hi.z+1) =  zero;
 
                // UNLESS if at the top of the domain and the boundary is outflow,
                //     we will use a homogeneous Neumann condition
                if ( (hi.z == domhi.z) &&
                     (phys_bc_type[5] == ERF_BC::outflow or phys_bc_type[5] == ERF_BC::ho_outflow) )
                {
                    coeffA_a(i,j,hi.z+1) =  -one;
                }
            }
        }
        for (int k = lo.z+1; k <= hi.z+1; ++k) {
            for (int j = lo.y; j <= hi.y; ++j) {
                AMREX_PRAGMA_SIMD
                for (int i = lo.x; i <= hi.x; ++i) {
                    coeffB_a(i,j,k) = one / ( coeffB_a(i,j,k) - coeffA_a(i,j,k)*gam_a(i,j,k-1) );
                    gam_a(i,j,k)    = coeffC_a(i,j,k) * coeffB_a(i,j,k);
                }
            }
        }
#endif
        } // end profile
    } // mfi
    } // omp
}

Here is the call graph for this function: