ERF_MakeFastCoeffs.cpp File Reference

#include <AMReX.H>
#include "ERF_Constants.H"
#include <ERF_TI_fast_headers.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, 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,
		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
[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]	phys_bc_type	Physical boundary condition types

{
    BL_PROFILE_VAR("make_fast_coeffs()",make_fast_coeffs);
 
    Real beta_2 = myhalf * (one + beta_s);  // multiplies implicit terms
 
    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);
 
 
    // *************************************************************************
    // 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();
 
        // *********************************************************************
        // *********************************************************************
        // *********************************************************************
 
        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
                Real D = beta_2 * beta_2 * dzi * static_cast<Real>(dtau * dtau);
                coeffA_a(i,j,k) = D * (one/detJ(i,j,k-1)) * ( halfg - coeff_Q * theta_t_lo );
                coeffC_a(i,j,k) = D * (one/detJ(i,j,k  )) * (-halfg + coeff_P * theta_t_hi );
 
                coeffB_a(i,j,k) = one + D * ( (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
                Real D = beta_2 * beta_2 * dzi * static_cast<Real>(dtau * dtau);
                coeffA_a(i,j,k) = D * ( halfg - coeff_Q * theta_t_lo );
                coeffC_a(i,j,k) = D * (-halfg + coeff_P * theta_t_hi );
 
                coeffB_a(i,j,k) = one + D * (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: