From 6fc689039babe9ffeac0ebff153137da8568c422 Mon Sep 17 00:00:00 2001 From: Matt Norman Date: Wed, 20 Sep 2023 15:43:32 -0400 Subject: [PATCH 1/2] Adding -DPAMA_DYCORE when A-grid dycore is specified in CMake. New A-grid dycore that uses SSPRK3, hydrostatically balanced gravity, and better boundary vertical conditions (isentropically constant). More informative error messages for the coupler, DataManager, and Options classes. Removing hydrostasis from the coupler. Adding gcm_cloud_water and gcm_cloud_ice to the coupler's default allocated state and setting to zero in the mmf_simplified driver. Removing duplicate -DPAM_RAD= from A-grid cmakescript.sh. Adding dycore.declare_current_profile_as_hydrostatic( coupler , true ) to the mmf_simplified driver for the A-grid dycore --- dynamics/awfl/CMakeLists.txt | 1 + dynamics/awfl/Dycore.h | 1532 +++++++++++- dynamics/awfl/Spatial_expl.h | 2228 ----------------- dynamics/awfl/Temporal_ader.h | 193 -- dynamics/awfl/Temporal_ssprk3_needs_update.h | 243 -- pam_core/DataManager.h | 64 +- pam_core/Options.h | 11 +- pam_core/pam_coupler.h | 109 +- .../mmf_simplified/build/cmakescript_pama.sh | 1 - standalone/mmf_simplified/driver.cpp | 9 +- .../mmf_simplified/inputs/input_pama.yaml | 4 +- 11 files changed, 1611 insertions(+), 2784 deletions(-) delete mode 100644 dynamics/awfl/Spatial_expl.h delete mode 100644 dynamics/awfl/Temporal_ader.h delete mode 100644 dynamics/awfl/Temporal_ssprk3_needs_update.h diff --git a/dynamics/awfl/CMakeLists.txt b/dynamics/awfl/CMakeLists.txt index 9065ddb8..16fce0f5 100644 --- a/dynamics/awfl/CMakeLists.txt +++ b/dynamics/awfl/CMakeLists.txt @@ -1,4 +1,5 @@ add_library(dycore INTERFACE) target_include_directories(dycore INTERFACE ${CMAKE_CURRENT_SOURCE_DIR}) +target_compile_definitions(dycore INTERFACE PAMA_DYCORE) diff --git a/dynamics/awfl/Dycore.h b/dynamics/awfl/Dycore.h index ee76eba2..8c2a7e43 100644 --- a/dynamics/awfl/Dycore.h +++ b/dynamics/awfl/Dycore.h @@ -1,10 +1,1530 @@ -#include "Spatial_expl.h" -#include "Temporal_ader.h" +#pragma once -// Define the Spatial operator based on constants from the Temporal operatora header file -typedef Spatial_operator Spatial; +#include "awfl_const.h" +#include "TransformMatrices.h" +#include "TransformMatrices_variable.h" +#include "WenoLimiter.h" +#include "idealized_profiles.h" +#include "MultipleFields.h" +#include "pam_coupler.h" -// Define the Temporal operator based on the Spatial operator -typedef Temporal_operator Dycore; + +class Dycore { + public: + + // Order of accuracy (numerical convergence for smooth flows) for the dynamical core + #ifndef MW_ORD + int static constexpr ord = 5; + #else + int static constexpr ord = MW_ORD; + #endif + // This is one value higher than normal to facilitate removing the Riemann solver phase + int static constexpr hs = (ord+1)/2; // Number of halo cells ("hs" == "halo size") + int static constexpr num_state = 5; + + // IDs for the variables in the state vector + int static constexpr idR = 0; // Density + int static constexpr idU = 1; // u-momentum + int static constexpr idV = 2; // v-momentum + int static constexpr idW = 3; // w-momentum + int static constexpr idT = 4; // Density * potential temperature + + SArray weno_recon_lower; + + + + real2d compute_mass( pam::PamCoupler const &coupler , realConst5d state , realConst5d tracers ) const { + using yakl::c::parallel_for; + using yakl::c::SimpleBounds; + auto num_tracers = coupler.get_num_tracers(); + auto nens = coupler.get_nens(); + auto nx = coupler.get_nx(); + auto ny = coupler.get_ny(); + auto nz = coupler.get_nz(); + auto dz = coupler.get_data_manager_device_readonly().get("vertical_cell_dz"); + int num_vars = num_tracers+2; + realHost2d mass("mass",num_vars,nens); + for (int ivar=0; ivar < num_vars; ivar++) { + for (int iens=0; iens < nens; iens++) { + real3d tmp("tmp",nz,ny,nx); + parallel_for( YAKL_AUTO_LABEL() , SimpleBounds<3>(nz,ny,nx) , YAKL_LAMBDA (int k, int j, int i) { + if (ivar < num_tracers ) { tmp(k,j,i) = tracers(ivar,hs+k,hs+j,hs+i,iens)*dz(k,iens); } + else if (ivar == num_tracers ) { tmp(k,j,i) = state (idR ,hs+k,hs+j,hs+i,iens)*dz(k,iens); } + else if (ivar == num_tracers+1) { tmp(k,j,i) = state (idT ,hs+k,hs+j,hs+i,iens)*dz(k,iens); } + }); + mass(ivar,iens) = yakl::intrinsics::sum(tmp)/tmp.size(); + } + } + return mass.createDeviceCopy(); + } + + + + // Compute the maximum stable time step using very conservative assumptions about max wind speed + real compute_time_step( pam::PamCoupler const &coupler , real cfl = 0.8 ) const { + using yakl::c::parallel_for; + using yakl::c::SimpleBounds; + auto nens = coupler.get_nens(); + auto nx = coupler.get_nx(); + auto ny = coupler.get_ny(); + auto nz = coupler.get_nz(); + auto dx = coupler.get_dx(); + auto dy = coupler.get_dy(); + auto R_d = coupler.get_option("R_d"); + auto R_v = coupler.get_option("R_v"); + auto gamma_d = coupler.get_option("gamma_d"); + auto &dm = coupler.get_data_manager_device_readonly(); + auto dm_rho_d = dm.get("density_dry"); + auto dm_uvel = dm.get("uvel" ); + auto dm_vvel = dm.get("vvel" ); + auto dm_wvel = dm.get("wvel" ); + auto dm_temp = dm.get("temp" ); + auto dm_rho_v = dm.get("water_vapor"); + auto dz = dm.get("vertical_cell_dz"); + real4d dt4d("dt4d",nz,ny,nx,nens); + parallel_for( YAKL_AUTO_LABEL() , SimpleBounds<4>(nz,ny,nx,nens) , YAKL_LAMBDA (int k, int j, int i, int iens) { + real rho_d = dm_rho_d(k,j,i,iens); + real u = dm_uvel (k,j,i,iens); + real v = dm_vvel (k,j,i,iens); + real w = dm_wvel (k,j,i,iens); + real temp = dm_temp (k,j,i,iens); + real rho_v = dm_rho_v(k,j,i,iens); + real rho = rho_d + rho_v; + real p = ( rho_d*R_d + rho_v*R_v ) * temp; + real cs = sqrt(gamma_d*p/rho); + real dtx = cfl * dx / (std::abs(u)+cs); + real dty = cfl * dy / (std::abs(v)+cs); + real dtz = cfl * dz(k,iens) / (std::abs(w)+cs); + dt4d(k,j,i,iens) = std::min( std::min( dtx , dty ) , dtz ); + }); + return yakl::intrinsics::minval( dt4d ); + } + + + + // Perform a single time step using SSPRK3 time stepping + void timeStep(pam::PamCoupler &coupler) { + using yakl::c::parallel_for; + using yakl::c::SimpleBounds; + using yakl::intrinsics::maxval; + using yakl::intrinsics::abs; + + auto num_tracers = coupler.get_num_tracers(); + auto nens = coupler.get_nens(); + auto nx = coupler.get_nx(); + auto ny = coupler.get_ny(); + auto nz = coupler.get_nz(); + auto tracer_positive = coupler.get_data_manager_device_readonly().get("tracer_positive"); + + real dt_phys = coupler.get_option("crm_dt"); + + // Create arrays to hold state and tracers with halos on the left and right of the domain + // Cells [0:hs-1] are the left halos, and cells [nx+hs:nx+2*hs-1] are the right halos + real5d state ("state" ,num_state ,nz+2*hs,ny+2*hs,nx+2*hs,nens); + real5d tracers("tracers",num_tracers,nz+2*hs,ny+2*hs,nx+2*hs,nens); + + // Populate the state and tracers arrays using data from the coupler, convert to the dycore's desired state + convert_coupler_to_dynamics( coupler , state , tracers ); + + parallel_for( YAKL_AUTO_LABEL() , SimpleBounds<5>(num_tracers,nz,ny,nx,nens) , YAKL_LAMBDA (int l, int k, int j, int i, int iens) { + if (tracer_positive(l)) { + tracers(l,hs+k,hs+j,hs+i,iens) = std::max( 0._fp , tracers(l,hs+k,hs+j,hs+i,iens) ); + } + }); + + #ifdef PAM_DEBUG + auto mass_init = compute_mass( coupler , state , tracers ); + #endif + + // Get the max stable time step for the dynamics. dt_phys might be > dt_dyn, meaning we would need to sub-cycle + real dt_dyn = compute_time_step( coupler ); + + // Get the number of sub-cycles we need, and set the dynamics time step accordingly + int ncycles = (int) std::ceil( dt_phys / dt_dyn ); + dt_dyn = dt_phys / ncycles; + + for (int icycle = 0; icycle < ncycles; icycle++) { + // SSPRK3 requires temporary arrays to hold intermediate state and tracers arrays + real5d state_tmp ("state_tmp" ,num_state ,nz+2*hs,ny+2*hs,nx+2*hs,nens); + real5d state_tend ("state_tend" ,num_state ,nz ,ny ,nx ,nens); + real5d tracers_tmp ("tracers_tmp" ,num_tracers,nz+2*hs,ny+2*hs,nx+2*hs,nens); + real5d tracers_tend("tracers_tend",num_tracers,nz ,ny ,nx ,nens); + ////////////// + // Stage 1 + ////////////// + parallel_for( YAKL_AUTO_LABEL() , SimpleBounds<5>(num_tracers,nz,ny,nx,nens) , YAKL_LAMBDA (int l, int k, int j, int i, int iens) { + // Store the starting point for FCT positivity in the next stage + tracers_tend(l,k,j,i,iens) = tracers(l,hs+k,hs+j,hs+i,iens); + }); + compute_tendencies( coupler , state , state_tend , tracers , tracers_tend , dt_dyn ); + // Apply tendencies + parallel_for( YAKL_AUTO_LABEL() , SimpleBounds<4>(nz,ny,nx,nens) , YAKL_LAMBDA (int k, int j, int i, int iens) { + for (int l = 0; l < num_state ; l++) { + state_tmp (l,hs+k,hs+j,hs+i,iens) = state (l,hs+k,hs+j,hs+i,iens) + dt_dyn * state_tend (l,k,j,i,iens); + } + for (int l = 0; l < num_tracers; l++) { + tracers_tmp(l,hs+k,hs+j,hs+i,iens) = tracers(l,hs+k,hs+j,hs+i,iens) + dt_dyn * tracers_tend(l,k,j,i,iens); + // For machine precision negative values after FCT-enforced positivity application + if (tracer_positive(l)) { + tracers_tmp(l,hs+k,hs+j,hs+i,iens) = std::max( 0._fp , tracers_tmp(l,hs+k,hs+j,hs+i,iens) ); + } + // Store the starting point for FCT positivity in the next stage + tracers_tend(l,k,j,i,iens) = (3._fp/4._fp) * tracers (l,hs+k,hs+j,hs+i,iens) + + (1._fp/4._fp) * tracers_tmp(l,hs+k,hs+j,hs+i,iens); + } + }); + ////////////// + // Stage 2 + ////////////// + compute_tendencies( coupler , state_tmp , state_tend , tracers_tmp , tracers_tend , (1._fp/4._fp) * dt_dyn ); + // Apply tendencies + parallel_for( YAKL_AUTO_LABEL() , SimpleBounds<4>(nz,ny,nx,nens) , YAKL_LAMBDA (int k, int j, int i, int iens) { + for (int l = 0; l < num_state ; l++) { + state_tmp (l,hs+k,hs+j,hs+i,iens) = (3._fp/4._fp) * state (l,hs+k,hs+j,hs+i,iens) + + (1._fp/4._fp) * state_tmp (l,hs+k,hs+j,hs+i,iens) + + (1._fp/4._fp) * dt_dyn * state_tend (l,k,j,i,iens); + } + for (int l = 0; l < num_tracers; l++) { + tracers_tmp(l,hs+k,hs+j,hs+i,iens) = (3._fp/4._fp) * tracers (l,hs+k,hs+j,hs+i,iens) + + (1._fp/4._fp) * tracers_tmp(l,hs+k,hs+j,hs+i,iens) + + (1._fp/4._fp) * dt_dyn * tracers_tend(l,k,j,i,iens); + // For machine precision negative values after FCT-enforced positivity application + if (tracer_positive(l)) { + tracers_tmp(l,hs+k,hs+j,hs+i,iens) = std::max( 0._fp , tracers_tmp(l,hs+k,hs+j,hs+i,iens) ); + } + // Store the starting point for FCT positivity in the next stage + tracers_tend(l,k,j,i,iens) = (1._fp/3._fp) * tracers (l,hs+k,hs+j,hs+i,iens) + + (2._fp/3._fp) * tracers_tmp(l,hs+k,hs+j,hs+i,iens); + } + }); + ////////////// + // Stage 3 + ////////////// + compute_tendencies( coupler , state_tmp , state_tend , tracers_tmp , tracers_tend , (2._fp/3._fp) * dt_dyn ); + // Apply tendencies + parallel_for( YAKL_AUTO_LABEL() , SimpleBounds<4>(nz,ny,nx,nens) , YAKL_LAMBDA (int k, int j, int i, int iens) { + for (int l = 0; l < num_state ; l++) { + state (l,hs+k,hs+j,hs+i,iens) = (1._fp/3._fp) * state (l,hs+k,hs+j,hs+i,iens) + + (2._fp/3._fp) * state_tmp (l,hs+k,hs+j,hs+i,iens) + + (2._fp/3._fp) * dt_dyn * state_tend (l,k,j,i,iens); + } + for (int l = 0; l < num_tracers; l++) { + tracers (l,hs+k,hs+j,hs+i,iens) = (1._fp/3._fp) * tracers (l,hs+k,hs+j,hs+i,iens) + + (2._fp/3._fp) * tracers_tmp(l,hs+k,hs+j,hs+i,iens) + + (2._fp/3._fp) * dt_dyn * tracers_tend(l,k,j,i,iens); + // For machine precision negative values after FCT-enforced positivity application + if (tracer_positive(l)) { + tracers (l,hs+k,hs+j,hs+i,iens) = std::max( 0._fp , tracers (l,hs+k,hs+j,hs+i,iens) ); + } + } + }); + } + + #ifdef PAM_DEBUG + auto mass_final = compute_mass( coupler , state , tracers ); + using yakl::componentwise::operator-; + using yakl::componentwise::operator+; + using yakl::componentwise::operator/; + using yakl::intrinsics::abs; + auto abs_mass_diff = abs(mass_final - mass_init).createHostCopy(); + auto rel_mass_diff = ( abs(mass_final - mass_init) / (abs(mass_init) + 1.e-20) ).createHostCopy(); + auto mass_diff = (mass_final - mass_init).createHostCopy(); + auto mass_init_host = mass_init.createHostCopy(); + auto mass_final_host = mass_final.createHostCopy(); + int num_vars = rel_mass_diff.extent(0); + for (int ivar=0; ivar < num_vars; ivar++) { + for (int iens=0; iens < nens; iens++) { + if (rel_mass_diff(ivar,iens) > 1.e-10) { + if (abs_mass_diff(ivar,iens) > 1.e-10) { + std::cout << "WARNING: conservation violated variable,ensemble,rel_diff,mass_diff,init,final: " + << ivar << " , " + << iens << " , " + << std::scientific << std::setw(10) << rel_mass_diff (ivar,iens) << " , " + << std::scientific << std::setw(10) << mass_diff (ivar,iens) << " , " + << std::scientific << std::setw(10) << mass_init_host (ivar,iens) << " , " + << std::scientific << std::setw(10) << mass_final_host(ivar,iens) << std::endl; + } + } + } + } + #endif + + // Convert the dycore's state back to the coupler's state + convert_dynamics_to_coupler( coupler , state , tracers ); + } + + + + // Compute the tendencies for state and tracers for one semi-discretized step inside the RK integrator + // Tendencies are the time rate of change for a quantity + // Coupler is non-const because we are writing to the flux variables + void compute_tendencies( pam::PamCoupler &coupler , + real5d const &state , + real5d const &state_tend , + real5d const &tracers , + real5d const &tracers_tend , + real dt ) const { + using yakl::c::parallel_for; + using yakl::c::SimpleBounds; + using std::min; + using std::max; + + auto nens = coupler.get_nens(); + auto nx = coupler.get_nx(); + auto ny = coupler.get_ny(); + auto nz = coupler.get_nz(); + auto dx = coupler.get_dx(); + auto dy = coupler.get_dy(); + auto sim2d = ny == 1; + auto C0 = coupler.get_option("C0" ); + auto gamma_d = coupler.get_option("gamma_d"); + auto grav = coupler.get_option("grav" ); + auto num_tracers = coupler.get_num_tracers(); + auto grav_balance = coupler.get_option("balance_hydrostasis_with_gravity"); + + // The store a single values flux at cell edges + auto &dm = coupler.get_data_manager_device_readwrite(); + auto tracer_positive = dm.get("tracer_positive"); + auto dz = dm.get("vertical_cell_dz"); + auto vert_weno_recon_lower = dm.get("vert_weno_recon_lower"); + auto vert_sten_to_coefs = dm.get("vert_sten_to_coefs"); + auto hy_dens_cells = dm.get("hy_dens_cells"); + auto hy_pressure_cells = dm.get("hy_pressure_cells"); + auto grav_var = dm.get("variable_gravity"); + + YAKL_SCOPE( weno_recon_lower , this->weno_recon_lower ); + + // Use TransformMatrices class to create matrices & GLL points to convert degrees of freedom as needed + SArray sten_to_coefs; + SArray coefs_to_gll; + TransformMatrices::coefs_to_gll_lower(coefs_to_gll ); + TransformMatrices::sten_to_coefs (sten_to_coefs); + SArray idl; + real sigma; + weno::wenoSetIdealSigma(idl,sigma); + + real4d pressure("pressure",nz+2*hs,ny+2*hs,nx+2*hs,nens); + + // Compute pressure perturbation, density perturbation, and divide density from all other quantities before interpolation + parallel_for( YAKL_AUTO_LABEL() , SimpleBounds<4>(nz,ny,nx,nens) , YAKL_LAMBDA (int k, int j, int i, int iens) { + if (grav_balance) { + pressure (hs+k,hs+j,hs+i,iens) = C0*std::pow(state(idT,hs+k,hs+j,hs+i,iens),gamma_d); + } else { + pressure (hs+k,hs+j,hs+i,iens) = C0*std::pow(state(idT,hs+k,hs+j,hs+i,iens),gamma_d) - hy_pressure_cells(k,iens); + } + state(idU,hs+k,hs+j,hs+i,iens) /= state(idR,hs+k,hs+j,hs+i,iens); + state(idV,hs+k,hs+j,hs+i,iens) /= state(idR,hs+k,hs+j,hs+i,iens); + state(idW,hs+k,hs+j,hs+i,iens) /= state(idR,hs+k,hs+j,hs+i,iens); + state(idT,hs+k,hs+j,hs+i,iens) /= state(idR,hs+k,hs+j,hs+i,iens); + for (int tr=0; tr < num_tracers; tr++) { tracers(tr,hs+k,hs+j,hs+i,iens) /= state(idR,hs+k,hs+j,hs+i,iens); } + }); + + halo_exchange( coupler , state , tracers , pressure ); + + real5d state_flux_x ("state_flux_x" ,num_state ,nz,ny,nx+1,nens); + real5d state_flux_y ("state_flux_y" ,num_state ,nz,ny+1,nx,nens); + real5d state_flux_z ("state_flux_z" ,num_state ,nz+1,ny,nx,nens); + real5d tracers_flux_x("tracers_flux_x",num_tracers,nz,ny,nx+1,nens); + real5d tracers_flux_y("tracers_flux_y",num_tracers,nz,ny+1,nx,nens); + real5d tracers_flux_z("tracers_flux_z",num_tracers,nz+1,ny,nx,nens); + + // Compute samples of state and tracers at cell edges using cell-centered reconstructions at high-order with WENO + // At the end of this, we will have two samples per cell edge in each dimension, one from each adjacent cell. + parallel_for( YAKL_AUTO_LABEL() , SimpleBounds<4>(nz+1,ny+1,nx+1,nens) , YAKL_LAMBDA (int k, int j, int i, int iens) { + real constexpr cs = 350; + //////////////////////////////////////////////////////// + // X-direction + //////////////////////////////////////////////////////// + if (j < ny && k < nz) { + SArray stencil; + // ACOUSTIC + real ru, pp; + { + // rho*u (left estimate) + int i_upw = 0; + for (int s=0; s < ord; s++) { stencil(s) = state(idR,hs+k,hs+j,i+i_upw+s,iens)*state(idU,hs+k,hs+j,i+i_upw+s,iens); } + real ru_L = reconstruct(stencil,coefs_to_gll,sten_to_coefs,weno_recon_lower,idl,sigma,1-i_upw); + // rho*u (right estimate) + i_upw = 1; + for (int s=0; s < ord; s++) { stencil(s) = state(idR,hs+k,hs+j,i+i_upw+s,iens)*state(idU,hs+k,hs+j,i+i_upw+s,iens); } + real ru_R = reconstruct(stencil,coefs_to_gll,sten_to_coefs,weno_recon_lower,idl,sigma,1-i_upw); + // pressure perturbation (left estimate) + i_upw = 0; + for (int s=0; s < ord; s++) { stencil(s) = pressure(hs+k,hs+j,i+i_upw+s,iens); } + real pp_L = reconstruct(stencil,coefs_to_gll,sten_to_coefs,weno_recon_lower,idl,sigma,1-i_upw); + // pressure perturbation (right estimate) + i_upw = 1; + for (int s=0; s < ord; s++) { stencil(s) = pressure(hs+k,hs+j,i+i_upw+s,iens); } + real pp_R = reconstruct(stencil,coefs_to_gll,sten_to_coefs,weno_recon_lower,idl,sigma,1-i_upw); + // Characteristics & upwind values + real w1 = 0.5_fp * (pp_R-cs*ru_R); + real w2 = 0.5_fp * (pp_L+cs*ru_L); + pp = w1+w2; + ru = (w2-w1)/cs; + state_flux_x(idR,k,j,i,iens) = ru; + } + // ADVECTIVE + int i_upw = ru > 0 ? 0 : 1; + // u-velocity + for (int s=0; s < ord; s++) { stencil(s) = state(idU,hs+k,hs+j,i+i_upw+s,iens); } + state_flux_x(idU,k,j,i,iens) = ru * reconstruct(stencil,coefs_to_gll,sten_to_coefs,weno_recon_lower,idl,sigma,1-i_upw) + pp; + // v-velocity + for (int s=0; s < ord; s++) { stencil(s) = state(idV,hs+k,hs+j,i+i_upw+s,iens); } + state_flux_x(idV,k,j,i,iens) = ru * reconstruct(stencil,coefs_to_gll,sten_to_coefs,weno_recon_lower,idl,sigma,1-i_upw); + // w-velocity + for (int s=0; s < ord; s++) { stencil(s) = state(idW,hs+k,hs+j,i+i_upw+s,iens); } + state_flux_x(idW,k,j,i,iens) = ru * reconstruct(stencil,coefs_to_gll,sten_to_coefs,weno_recon_lower,idl,sigma,1-i_upw); + // theta + for (int s=0; s < ord; s++) { stencil(s) = state(idT,hs+k,hs+j,i+i_upw+s,iens); } + state_flux_x(idT,k,j,i,iens) = ru * reconstruct(stencil,coefs_to_gll,sten_to_coefs,weno_recon_lower,idl,sigma,1-i_upw); + // tracers + for (int tr=0; tr < num_tracers; tr++) { + for (int s=0; s < ord; s++) { stencil(s) = tracers(tr,hs+k,hs+j,i+i_upw+s,iens); } + tracers_flux_x(tr,k,j,i,iens) = ru * reconstruct(stencil,coefs_to_gll,sten_to_coefs,weno_recon_lower,idl,sigma,1-i_upw); + } + } + + //////////////////////////////////////////////////////// + // Y-direction + //////////////////////////////////////////////////////// + if (i < nx && k < nz) { + if (! sim2d) { + SArray stencil; + // ACOUSTIC + real rv, pp; + { + // rho*v (left estimate) + int j_upw = 0; + for (int s=0; s < ord; s++) { stencil(s) = state(idR,hs+k,j+j_upw+s,hs+i,iens)*state(idV,hs+k,j+j_upw+s,hs+i,iens); } + real rv_L = reconstruct(stencil,coefs_to_gll,sten_to_coefs,weno_recon_lower,idl,sigma,1-j_upw); + // rho*v (right estimate) + j_upw = 1; + for (int s=0; s < ord; s++) { stencil(s) = state(idR,hs+k,j+j_upw+s,hs+i,iens)*state(idV,hs+k,j+j_upw+s,hs+i,iens); } + real rv_R = reconstruct(stencil,coefs_to_gll,sten_to_coefs,weno_recon_lower,idl,sigma,1-j_upw); + // pressure perturbation (left estimate) + j_upw = 0; + for (int s=0; s < ord; s++) { stencil(s) = pressure(hs+k,j+j_upw+s,hs+i,iens); } + real pp_L = reconstruct(stencil,coefs_to_gll,sten_to_coefs,weno_recon_lower,idl,sigma,1-j_upw); + // pressure perturbation (right estimate) + j_upw = 1; + for (int s=0; s < ord; s++) { stencil(s) = pressure(hs+k,j+j_upw+s,hs+i,iens); } + real pp_R = reconstruct(stencil,coefs_to_gll,sten_to_coefs,weno_recon_lower,idl,sigma,1-j_upw); + // Characteristics & upwind values + real w1 = 0.5_fp * (pp_R-cs*rv_R); + real w2 = 0.5_fp * (pp_L+cs*rv_L); + pp = w1+w2; + rv = (w2-w1)/cs; + state_flux_y(idR,k,j,i,iens) = rv; + } + // ADVECTIVE + int j_upw = rv > 0 ? 0 : 1; + // u-velocity + for (int s=0; s < ord; s++) { stencil(s) = state(idU,hs+k,j+j_upw+s,hs+i,iens); } + state_flux_y(idU,k,j,i,iens) = rv * reconstruct(stencil,coefs_to_gll,sten_to_coefs,weno_recon_lower,idl,sigma,1-j_upw); + // v-velocity + for (int s=0; s < ord; s++) { stencil(s) = state(idV,hs+k,j+j_upw+s,hs+i,iens); } + state_flux_y(idV,k,j,i,iens) = rv * reconstruct(stencil,coefs_to_gll,sten_to_coefs,weno_recon_lower,idl,sigma,1-j_upw) + pp; + // w-velocity + for (int s=0; s < ord; s++) { stencil(s) = state(idW,hs+k,j+j_upw+s,hs+i,iens); } + state_flux_y(idW,k,j,i,iens) = rv * reconstruct(stencil,coefs_to_gll,sten_to_coefs,weno_recon_lower,idl,sigma,1-j_upw); + // theta + for (int s=0; s < ord; s++) { stencil(s) = state(idT,hs+k,j+j_upw+s,hs+i,iens); } + state_flux_y(idT,k,j,i,iens) = rv * reconstruct(stencil,coefs_to_gll,sten_to_coefs,weno_recon_lower,idl,sigma,1-j_upw); + // tracers + for (int tr=0; tr < num_tracers; tr++) { + for (int s=0; s < ord; s++) { stencil(s) = tracers(tr,hs+k,j+j_upw+s,hs+i,iens); } + tracers_flux_y(tr,k,j,i,iens) = rv * reconstruct(stencil,coefs_to_gll,sten_to_coefs,weno_recon_lower,idl,sigma,1-j_upw); + } + } else { + state_flux_y(idR,k,j,i,iens) = 0; + state_flux_y(idU,k,j,i,iens) = 0; + state_flux_y(idV,k,j,i,iens) = 0; + state_flux_y(idW,k,j,i,iens) = 0; + state_flux_y(idT,k,j,i,iens) = 0; + for (int tr=0; tr < num_tracers; tr++) { tracers_flux_y(tr,k,j,i,iens) = 0; } + } + } + + //////////////////////////////////////////////////////// + // Z-direction + //////////////////////////////////////////////////////// + if (i < nx && j < ny) { + SArray stencil; + SArray s2c_loc[2]; + SArray wrl_loc[2]; + for (int i1=0; i1 < ord; i1++) { + for (int i2=0; i2 < ord; i2++) { + s2c_loc[0](i1,i2) = vert_sten_to_coefs(k ,i1,i2,iens); + s2c_loc[1](i1,i2) = vert_sten_to_coefs(k+1,i1,i2,iens); + } + } + for (int i1=0; i1 < hs; i1++) { + for (int i2=0; i2 < hs; i2++) { + for (int i3=0; i3 < hs; i3++) { + wrl_loc[0](i1,i2,i3) = vert_weno_recon_lower(k ,i1,i2,i3,iens); + wrl_loc[1](i1,i2,i3) = vert_weno_recon_lower(k+1,i1,i2,i3,iens); + } + } + } + // ACOUSTIC + real rw, pp; + { + // rho*w (left estimate) + int k_upw = 0; + for (int s=0; s < ord; s++) { stencil(s) = state(idR,k+k_upw+s,hs+j,hs+i,iens)*state(idW,k+k_upw+s,hs+j,hs+i,iens); } + real rw_L = reconstruct(stencil,coefs_to_gll,s2c_loc[k_upw],wrl_loc[k_upw],idl,sigma,1-k_upw); + if ((k == 0 || k == nz)) rw_L = 0; + // rho*w (right estimate) + k_upw = 1; + for (int s=0; s < ord; s++) { stencil(s) = state(idR,k+k_upw+s,hs+j,hs+i,iens)*state(idW,k+k_upw+s,hs+j,hs+i,iens); } + real rw_R = reconstruct(stencil,coefs_to_gll,s2c_loc[k_upw],wrl_loc[k_upw],idl,sigma,1-k_upw); + if ((k == 0 || k == nz)) rw_R = 0; + // pressure perturbation (left estimate) + k_upw = 0; + for (int s=0; s < ord; s++) { stencil(s) = pressure(k+k_upw+s,hs+j,hs+i,iens); } + real pp_L = reconstruct(stencil,coefs_to_gll,s2c_loc[k_upw],wrl_loc[k_upw],idl,sigma,1-k_upw); + // pressure perturbation (right estimate) + k_upw = 1; + for (int s=0; s < ord; s++) { stencil(s) = pressure(k+k_upw+s,hs+j,hs+i,iens); } + real pp_R = reconstruct(stencil,coefs_to_gll,s2c_loc[k_upw],wrl_loc[k_upw],idl,sigma,1-k_upw); + // Characteristics & upwind values + real w1 = 0.5_fp * (pp_R-cs*rw_R); + real w2 = 0.5_fp * (pp_L+cs*rw_L); + pp = w1+w2; + rw = (w2-w1)/cs; + if (k == 0 || k == nz) rw = 0; + state_flux_z(idR,k,j,i,iens) = rw; + } + // ADVECTIVE + int k_upw = rw > 0 ? 0 : 1; + // u-velocity + for (int s=0; s < ord; s++) { stencil(s) = state(idU,k+k_upw+s,hs+j,hs+i,iens); } + state_flux_z(idU,k,j,i,iens) = rw * reconstruct(stencil,coefs_to_gll,s2c_loc[k_upw],wrl_loc[k_upw],idl,sigma,1-k_upw); + // v-velocity + for (int s=0; s < ord; s++) { stencil(s) = state(idV,k+k_upw+s,hs+j,hs+i,iens); } + state_flux_z(idV,k,j,i,iens) = rw * reconstruct(stencil,coefs_to_gll,s2c_loc[k_upw],wrl_loc[k_upw],idl,sigma,1-k_upw); + // w-velocity + for (int s=0; s < ord; s++) { stencil(s) = state(idW,k+k_upw+s,hs+j,hs+i,iens); } + state_flux_z(idW,k,j,i,iens) = rw * reconstruct(stencil,coefs_to_gll,s2c_loc[k_upw],wrl_loc[k_upw],idl,sigma,1-k_upw) + pp; + // theta + for (int s=0; s < ord; s++) { stencil(s) = state(idT,k+k_upw+s,hs+j,hs+i,iens); } + state_flux_z(idT,k,j,i,iens) = rw * reconstruct(stencil,coefs_to_gll,s2c_loc[k_upw],wrl_loc[k_upw],idl,sigma,1-k_upw); + // tracers + for (int tr=0; tr < num_tracers; tr++) { + for (int s=0; s < ord; s++) { stencil(s) = tracers(tr,k+k_upw+s,hs+j,hs+i,iens); } + tracers_flux_z(tr,k,j,i,iens) = rw * reconstruct(stencil,coefs_to_gll,s2c_loc[k_upw],wrl_loc[k_upw],idl,sigma,1-k_upw); + } + } + }); + + // Flux Corrected Transport to enforce positivity for tracer species that must remain non-negative + // This looks like it has a race condition, but it does not. Only one of the adjacent cells can ever change + // a given edge flux because it's only changed if its sign oriented outward from a cell. + // Also, multiply density back onto the state and tracers + parallel_for( YAKL_AUTO_LABEL() , SimpleBounds<4>(nz,ny,nx,nens) , + YAKL_LAMBDA (int k, int j, int i, int iens) { + state(idU,hs+k,hs+j,hs+i,iens) *= state(idR,hs+k,hs+j,hs+i,iens); + state(idV,hs+k,hs+j,hs+i,iens) *= state(idR,hs+k,hs+j,hs+i,iens); + state(idW,hs+k,hs+j,hs+i,iens) *= state(idR,hs+k,hs+j,hs+i,iens); + state(idT,hs+k,hs+j,hs+i,iens) *= state(idR,hs+k,hs+j,hs+i,iens); + for (int tr=0; tr < num_tracers; tr++) { + tracers(tr,hs+k,hs+j,hs+i,iens) *= state(idR,hs+k,hs+j,hs+i,iens); + if (tracer_positive(tr)) { + real mass_available = max(tracers_tend(tr,k,j,i,iens),0._fp) * dx * dy * dz(k,iens); + real flux_out_x = ( max(tracers_flux_x(tr,k,j,i+1,iens),0._fp) - min(tracers_flux_x(tr,k,j,i,iens),0._fp) ) / dx; + real flux_out_y = ( max(tracers_flux_y(tr,k,j+1,i,iens),0._fp) - min(tracers_flux_y(tr,k,j,i,iens),0._fp) ) / dy; + real flux_out_z = ( max(tracers_flux_z(tr,k+1,j,i,iens),0._fp) - min(tracers_flux_z(tr,k,j,i,iens),0._fp) ) / dz(k,iens); + real mass_out = (flux_out_x + flux_out_y + flux_out_z) * dt * dx * dy * dz(k,iens); + if (mass_out > mass_available) { + real mult = mass_available / mass_out; + if (tracers_flux_x(tr,k,j,i+1,iens) > 0) tracers_flux_x(tr,k,j,i+1,iens) *= mult; + if (tracers_flux_x(tr,k,j,i ,iens) < 0) tracers_flux_x(tr,k,j,i ,iens) *= mult; + if (tracers_flux_y(tr,k,j+1,i,iens) > 0) tracers_flux_y(tr,k,j+1,i,iens) *= mult; + if (tracers_flux_y(tr,k,j ,i,iens) < 0) tracers_flux_y(tr,k,j ,i,iens) *= mult; + if (tracers_flux_z(tr,k+1,j,i,iens) > 0) tracers_flux_z(tr,k+1,j,i,iens) *= mult; + if (tracers_flux_z(tr,k ,j,i,iens) < 0) tracers_flux_z(tr,k ,j,i,iens) *= mult; + } + } + } + }); + + // Compute tendencies as the flux divergence + gravity source term + parallel_for( YAKL_AUTO_LABEL() , SimpleBounds<4>(nz,ny,nx,nens) , YAKL_LAMBDA (int k, int j, int i, int iens) { + for (int l = 0; l < num_state; l++) { + state_tend (l,k,j,i,iens) = -( state_flux_x(l,k ,j ,i+1,iens) - state_flux_x(l,k,j,i,iens) ) / dx + -( state_flux_y(l,k ,j+1,i ,iens) - state_flux_y(l,k,j,i,iens) ) / dy + -( state_flux_z(l,k+1,j ,i ,iens) - state_flux_z(l,k,j,i,iens) ) / dz(k,iens); + if (l == idW) { + if (grav_balance) { + state_tend(l,k,j,i,iens) += -grav_var(k,iens) * state(idR,hs+k,hs+j,hs+i,iens); + } else { + state_tend(l,k,j,i,iens) += -grav * ( state(idR,hs+k,hs+j,hs+i,iens) - hy_dens_cells(k,iens) ); + } + } + if (l == idV && sim2d) state_tend(l,k,j,i,iens) = 0; + } + for (int l = 0; l < num_tracers; l++) { + real fx = tracers_flux_x(l,k ,j ,i ,iens); + real fxp1 = tracers_flux_x(l,k ,j ,i+1,iens); + real fy = tracers_flux_y(l,k ,j ,i ,iens); + real fyp1 = tracers_flux_y(l,k ,j+1,i ,iens); + real fz = tracers_flux_z(l,k ,j ,i ,iens); + real fzp1 = tracers_flux_z(l,k+1,j ,i ,iens); + if (i == 0 ) { + fx = std::min( fx , tracers_flux_x(l,k,j,nx,iens) ); + } + if (i == nx-1) { fxp1 = std::min( fxp1 , tracers_flux_x(l,k,j,0 ,iens) ); } + if (j == 0 ) { fy = std::min( fy , tracers_flux_y(l,k,ny,i,iens) ); } + if (j == ny-1) { fyp1 = std::min( fyp1 , tracers_flux_y(l,k,0 ,i,iens) ); } + tracers_tend(l,k,j,i,iens) = -( fxp1 - fx ) / dx + -( fyp1 - fy ) / dy + -( fzp1 - fz ) / dz(k,iens); + } + }); + + } + + + + // ord stencil cell averages to two GLL point values via high-order reconstruction and WENO limiting + YAKL_INLINE static real reconstruct( SArray const &stencil , + SArray const &coefs_to_gll , + SArray const &sten_to_coefs , + SArray const &weno_recon_lower , + SArray const &idl , + real sigma , + int ind ) { + // Reconstruct values + SArray wenoCoefs; + weno::compute_weno_coefs( weno_recon_lower , sten_to_coefs , stencil , wenoCoefs , idl , sigma ); + real tmp = 0; + for (int s=0; s < ord; s++) { tmp += coefs_to_gll(s,ind) * wenoCoefs(s); } + return tmp; + } + + + + void halo_exchange( pam::PamCoupler const &coupler , + real5d const &state , + real5d const &tracers , + real4d const &pressure ) const { + using yakl::c::parallel_for; + using yakl::c::SimpleBounds; + + auto nens = coupler.get_nens(); + auto nx = coupler.get_nx(); + auto ny = coupler.get_ny(); + auto nz = coupler.get_nz(); + auto num_tracers = coupler.get_num_tracers(); + auto sim2d = ny == 1; + auto gamma = coupler.get_option("gamma_d"); + auto C0 = coupler.get_option("C0"); + auto grav = coupler.get_option("grav"); + auto grav_balance = coupler.get_option("balance_hydrostasis_with_gravity"); + auto dz = coupler.get_data_manager_device_readonly().get("vertical_cell_dz"); + + int npack = num_state + num_tracers + 1; + + parallel_for( YAKL_AUTO_LABEL() , SimpleBounds<5>(npack,nz,ny,hs,nens) , + YAKL_LAMBDA (int v, int k, int j, int ii, int iens) { + if (v < num_state) { + state (v ,hs+k,hs+j,nx+hs+ii,iens) = state (v ,hs+k,hs+j,hs+ii,iens); + state (v ,hs+k,hs+j, ii,iens) = state (v ,hs+k,hs+j,nx+ii,iens); + } else if (v < num_state + num_tracers) { + tracers(v-num_state,hs+k,hs+j,nx+hs+ii,iens) = tracers(v-num_state,hs+k,hs+j,hs+ii,iens); + tracers(v-num_state,hs+k,hs+j, ii,iens) = tracers(v-num_state,hs+k,hs+j,nx+ii,iens); + } else { + pressure(hs+k,hs+j,nx+hs+ii,iens) = pressure(hs+k,hs+j,hs+ii,iens); + pressure(hs+k,hs+j, ii,iens) = pressure(hs+k,hs+j,nx+ii,iens); + } + }); + + if (!sim2d) { + parallel_for( YAKL_AUTO_LABEL() , SimpleBounds<5>(npack,nz,hs,nx,nens) , + YAKL_LAMBDA (int v, int k, int jj, int i, int iens) { + if (v < num_state) { + state (v ,hs+k,ny+hs+jj,hs+i,iens) = state (v ,hs+k,hs+jj,hs+i,iens); + state (v ,hs+k, jj,hs+i,iens) = state (v ,hs+k,ny+jj,hs+i,iens); + } else if (v < num_state + num_tracers) { + tracers(v-num_state,hs+k,ny+hs+jj,hs+i,iens) = tracers(v-num_state,hs+k,hs+jj,hs+i,iens); + tracers(v-num_state,hs+k, jj,hs+i,iens) = tracers(v-num_state,hs+k,ny+jj,hs+i,iens); + } else { + pressure(hs+k,ny+hs+jj,hs+i,iens) = pressure(hs+k,hs+jj,hs+i,iens); + pressure(hs+k, jj,hs+i,iens) = pressure(hs+k,ny+jj,hs+i,iens); + } + }); + } + + //////////////////////////////////// + // Begin boundary conditions + //////////////////////////////////// + parallel_for( YAKL_AUTO_LABEL() , SimpleBounds<4>(hs,ny,nx,nens) , + YAKL_LAMBDA (int kk, int j, int i, int iens) { + // Exclude density (idR == 0) because that's calculated lower down + for (int l=1; l < num_state; l++) { + if (l == idW) { + state(l, kk,hs+j,hs+i,iens) = 0; + state(l,hs+nz+kk,hs+j,hs+i,iens) = 0; + } else { + state(l, kk,hs+j,hs+i,iens) = state(l,hs+0 ,hs+j,hs+i,iens); + state(l,hs+nz+kk,hs+j,hs+i,iens) = state(l,hs+nz-1,hs+j,hs+i,iens); + } + } + for (int l=0; l < num_tracers; l++) { + tracers(l, kk,hs+j,hs+i,iens) = tracers(l,hs+0 ,hs+j,hs+i,iens); + tracers(l,hs+nz+kk,hs+j,hs+i,iens) = tracers(l,hs+nz-1,hs+j,hs+i,iens); + } + if (! grav_balance) { + pressure( kk,hs+j,hs+i,iens) = pressure(hs+0 ,hs+j,hs+i,iens); + pressure(hs+nz+kk,hs+j,hs+i,iens) = pressure(hs+nz-1,hs+j,hs+i,iens); + } + { + int k0 = hs; + int k = k0-1-kk; + real rho0 = state(idR,k0,hs+j,hs+i,iens); + real theta0 = state(idT,k0,hs+j,hs+i,iens); + real rho0_gm1 = std::pow(rho0 ,gamma-1); + real theta0_g = std::pow(theta0,gamma ); + state(idR,k,hs+j,hs+i,iens) = std::pow( rho0_gm1 + grav*(gamma-1)*dz(k0-hs,iens)*(kk+1)/(gamma*C0*theta0_g) , + 1._fp/(gamma-1) ); + state(idT,k,hs+j,hs+i,iens) = theta0; + if (grav_balance) { + real rt = state(idR,k,hs+j,hs+i,iens) * state(idT,k,hs+j,hs+i,iens); + pressure(k,hs+j,hs+i,iens) = C0*std::pow( rt , gamma ); + } + } + { + int k0 = hs+nz-1; + int k = k0+1+kk; + real rho0 = state(idR,k0,hs+j,hs+i,iens); + real theta0 = state(idT,k0,hs+j,hs+i,iens); + real rho0_gm1 = std::pow(rho0 ,gamma-1); + real theta0_g = std::pow(theta0,gamma ); + state(idR,k,hs+j,hs+i,iens) = std::pow( rho0_gm1 - grav*(gamma-1)*dz(k0-hs,iens)*(kk+1)/(gamma*C0*theta0_g) , + 1._fp/(gamma-1) ); + state(idT,k,hs+j,hs+i,iens) = theta0; + if (grav_balance) { + real rt = state(idR,k,hs+j,hs+i,iens) * state(idT,k,hs+j,hs+i,iens); + pressure(k,hs+j,hs+i,iens) = C0*std::pow( rt , gamma ); + } + } + }); + } + + + + // Creates initial data at a point in space for the rising moist thermal test case + YAKL_INLINE static void thermal(real x, real y, real z, real xlen, real ylen, real grav, real C0, real gamma, + real cp, real p0, real R_d, real R_v, real &rho, real &u, real &v, real &w, + real &theta, real &rho_v, real &hr, real &ht) { + hydro_const_theta(z,grav,C0,cp,p0,gamma,R_d,hr,ht); + real rho_d = hr; + u = 0.; + v = 0.; + w = 0.; + real theta_d = ht + sample_ellipse_cosine(2._fp , x,y,z , xlen/2,ylen/2,2000. , 2000.,2000.,2000.); + real p_d = C0 * pow( rho_d*theta_d , gamma ); + real temp = p_d / rho_d / R_d; + real sat_pv = saturation_vapor_pressure(temp); + real sat_rv = sat_pv / R_v / temp; + rho_v = sample_ellipse_cosine(0.8_fp , x,y,z , xlen/2,ylen/2,2000. , 2000.,2000.,2000.) * sat_rv; + real p = rho_d * R_d * temp + rho_v * R_v * temp; + rho = rho_d + rho_v; + theta = std::pow( p / C0 , 1._fp / gamma ) / rho; + } + + + + // Computes a hydrostatic background density and potential temperature using c constant potential temperature + // backgrounda for a single vertical location + YAKL_INLINE static void hydro_const_theta( real z, real grav, real C0, real cp, real p0, real gamma, real rd, + real &r, real &t ) { + const real theta0 = 300.; //Background potential temperature + const real exner0 = 1.; //Surface-level Exner pressure + t = theta0; //Potential Temperature at z + real exner = exner0 - grav * z / (cp * theta0); //Exner pressure at z + real p = p0 * std::pow(exner,(cp/rd)); //Pressure at z + real rt = std::pow((p / C0),(1._fp / gamma)); //rho*theta at z + r = rt / t; //Density at z + } + + + + // Samples a 3-D ellipsoid at a point in space + YAKL_INLINE static real sample_ellipse_cosine(real amp, real x , real y , real z , + real x0 , real y0 , real z0 , + real xrad, real yrad, real zrad) { + //Compute distance from bubble center + real dist = sqrt( ((x-x0)/xrad)*((x-x0)/xrad) + + ((y-y0)/yrad)*((y-y0)/yrad) + + ((z-z0)/zrad)*((z-z0)/zrad) ) * M_PI / 2.; + //If the distance from bubble center is less than the radius, create a cos**2 profile + if (dist <= M_PI / 2.) { + return amp * std::pow(cos(dist),2._fp); + } else { + return 0.; + } + } + + + + YAKL_INLINE static real saturation_vapor_pressure(real temp) { + real tc = temp - 273.15; + return 610.94 * std::exp( 17.625*tc / (243.04+tc) ); + } + + + + // Compute supercell temperature profile at a vertical location + YAKL_INLINE static real init_supercell_temperature(real z, real z_0, real z_trop, real z_top, + real T_0, real T_trop, real T_top) { + if (z <= z_trop) { + real lapse = - (T_trop - T_0) / (z_trop - z_0); + return T_0 - lapse * (z - z_0); + } else { + real lapse = - (T_top - T_trop) / (z_top - z_trop); + return T_trop - lapse * (z - z_trop); + } + } + + + + // Compute supercell dry pressure profile at a vertical location + YAKL_INLINE static real init_supercell_pressure_dry(real z, real z_0, real z_trop, real z_top, + real T_0, real T_trop, real T_top, + real p_0, real R_d, real grav) { + if (z <= z_trop) { + real lapse = - (T_trop - T_0) / (z_trop - z_0); + real T = init_supercell_temperature(z, z_0, z_trop, z_top, T_0, T_trop, T_top); + return p_0 * pow( T / T_0 , grav/(R_d*lapse) ); + } else { + // Get pressure at the tropopause + real lapse = - (T_trop - T_0) / (z_trop - z_0); + real p_trop = p_0 * pow( T_trop / T_0 , grav/(R_d*lapse) ); + // Get pressure at requested height + lapse = - (T_top - T_trop) / (z_top - z_trop); + if (lapse != 0) { + real T = init_supercell_temperature(z, z_0, z_trop, z_top, T_0, T_trop, T_top); + return p_trop * pow( T / T_trop , grav/(R_d*lapse) ); + } else { + return p_trop * exp(-grav*(z-z_trop)/(R_d*T_trop)); + } + } + } + + + + // Compute supercell relative humidity profile at a vertical location + YAKL_INLINE static real init_supercell_relhum(real z, real z_0, real z_trop) { + if (z <= z_trop) { + return 1._fp - 0.75_fp * pow(z / z_trop , 1.25_fp ); + } else { + return 0.25_fp; + } + } + + + + // Computes dry saturation mixing ratio + YAKL_INLINE static real init_supercell_sat_mix_dry( real press , real T ) { + return 380/(press) * exp( 17.27_fp * (T-273)/(T-36) ); + } + + + + // Initialize the class data as well as the state and tracers arrays and convert them back into the coupler state + void init(pam::PamCoupler &coupler) { + using yakl::c::parallel_for; + using yakl::c::SimpleBounds; + + auto nens = coupler.get_nens(); + auto nx = coupler.get_nx(); + auto ny = coupler.get_ny(); + auto nz = coupler.get_nz(); + auto dx = coupler.get_dx(); + auto dy = coupler.get_dy(); + auto xlen = coupler.get_xlen(); + auto ylen = coupler.get_ylen(); + auto sim2d = ny == 1; + auto num_tracers = coupler.get_num_tracers(); + + coupler.set_option("balance_hydrostasis_with_gravity",true); + if (coupler.get_option("balance_hydrostasis_with_gravity")) { + coupler.get_data_manager_device_readwrite().register_and_allocate("variable_gravity","",{nz,nens}); + } + + if (! coupler.option_exists("R_d" )) coupler.set_option("R_d" ,287. ); + if (! coupler.option_exists("cp_d" )) coupler.set_option("cp_d" ,1003. ); + if (! coupler.option_exists("R_v" )) coupler.set_option("R_v" ,461. ); + if (! coupler.option_exists("cp_v" )) coupler.set_option("cp_v" ,1859 ); + if (! coupler.option_exists("p0" )) coupler.set_option("p0" ,1.e5 ); + if (! coupler.option_exists("grav" )) coupler.set_option("grav" ,9.81 ); + auto R_d = coupler.get_option("R_d" ); + auto cp_d = coupler.get_option("cp_d"); + auto R_v = coupler.get_option("R_v" ); + auto cp_v = coupler.get_option("cp_v"); + auto p0 = coupler.get_option("p0" ); + auto grav = coupler.get_option("grav"); + if (! coupler.option_exists("cv_d" )) coupler.set_option("cv_d" ,cp_d - R_d ); + auto cv_d = coupler.get_option("cv_d"); + if (! coupler.option_exists("gamma_d")) coupler.set_option("gamma_d",cp_d / cv_d); + if (! coupler.option_exists("kappa_d")) coupler.set_option("kappa_d",R_d / cp_d); + if (! coupler.option_exists("cv_v" )) coupler.set_option("cv_v" ,R_v - cp_v ); + auto gamma = coupler.get_option("gamma_d"); + auto kappa = coupler.get_option("kappa_d"); + if (! coupler.option_exists("C0")) coupler.set_option("C0" , pow( R_d * pow( p0 , -kappa ) , gamma )); + auto C0 = coupler.get_option("C0"); + + auto &dm = coupler.get_data_manager_device_readwrite(); + auto dz = dm.get("vertical_cell_dz"); + + // Create WENO reconstruction matrices + dm.register_and_allocate("vert_weno_recon_lower","",{nz+2,hs,hs,hs,nens}); + dm.register_and_allocate("vert_sten_to_coefs" ,"",{nz+2,ord,ord ,nens}); + auto vert_weno_recon_lower = dm.get("vert_weno_recon_lower"); + auto vert_sten_to_coefs = dm.get("vert_sten_to_coefs" ); + auto dz_host = dz .createHostCopy(); + auto vert_weno_host = vert_weno_recon_lower.createHostCopy(); + auto vert_s2c_host = vert_sten_to_coefs .createHostCopy(); + for (int k=0; k < nz+2; k++) { + for (int iens = 0; iens < nens; iens++) { + // Load normalize grid spacings + SArray dzloc; + for (int kk=0; kk < ord; kk++) { + int ind1 = std::min(nz-1,std::max(0,-1+k+kk)); + int ind2 = std::min(nz-1,std::max(0,-1+k )); + dzloc(kk) = dz_host(ind1,iens) / dz_host(ind2,iens); + } + // Compute normalized locations of cell edges + SArray locs; + locs(0) = 0; + for (int kk=1; kk < ord+1; kk++) { locs(kk) = locs(kk-1) + dzloc(kk-1); } + real midloc = ( locs((ord-1)/2) + locs((ord+1)/2) ) / 2; + for (int kk=0; kk < ord+1; kk++) { locs(kk) = locs(kk) - midloc; } + // Compute WENO reconstruction matrices for each level + SArray s2c_var; + SArray weno_recon_lower_var; + TransformMatrices_variable::sten_to_coefs_variable (locs,s2c_var ); + TransformMatrices_variable::weno_lower_sten_to_coefs(locs,weno_recon_lower_var); + for (int jj=0; jj < ord; jj++) { + for (int ii=0; ii < ord; ii++) { + vert_s2c_host(k,jj,ii,iens) = s2c_var(jj,ii); + } + } + for (int kk=0; kk < hs; kk++) { + for (int jj=0; jj < hs; jj++) { + for (int ii=0; ii < hs; ii++) { + vert_weno_host(k,kk,jj,ii,iens) = weno_recon_lower_var(kk,jj,ii); + } + } + } + } + } + vert_s2c_host .deep_copy_to(vert_sten_to_coefs ); + vert_weno_host.deep_copy_to(vert_weno_recon_lower); + TransformMatrices::weno_lower_sten_to_coefs(this->weno_recon_lower); + + R_d = coupler.get_option("R_d" ); + R_v = coupler.get_option("R_v" ); + cp_d = coupler.get_option("cp_d" ); + cp_v = coupler.get_option("cp_v" ); + p0 = coupler.get_option("p0" ); + grav = coupler.get_option("grav" ); + kappa = coupler.get_option("kappa_d"); + gamma = coupler.get_option("gamma_d"); + C0 = coupler.get_option("C0" ); + + // Create arrays to determine whether we should add mass for a tracer or whether it should remain non-negative + num_tracers = coupler.get_num_tracers(); + bool1d tracer_adds_mass("tracer_adds_mass",num_tracers); + bool1d tracer_positive ("tracer_positive" ,num_tracers); + + // Must assign on the host to avoid segfaults + auto tracer_adds_mass_host = tracer_adds_mass.createHostCopy(); + auto tracer_positive_host = tracer_positive .createHostCopy(); + + auto tracer_names = coupler.get_tracer_names(); // Get a list of tracer names + int idWV; + for (int tr=0; tr < num_tracers; tr++) { + std::string tracer_desc; + bool tracer_found, positive, adds_mass; + coupler.get_tracer_info( tracer_names[tr] , tracer_desc, tracer_found , positive , adds_mass); + tracer_positive_host (tr) = positive; + tracer_adds_mass_host(tr) = adds_mass; + if (tracer_names[tr] == "water_vapor") idWV = tr; // Be sure to track which index belongs to water vapor + } + tracer_positive_host .deep_copy_to(tracer_positive ); + tracer_adds_mass_host.deep_copy_to(tracer_adds_mass); + + coupler.set_option("idWV",idWV); + dm.register_and_allocate("tracer_adds_mass","",{num_tracers}); + auto dm_tracer_adds_mass = dm.get("tracer_adds_mass"); + tracer_adds_mass.deep_copy_to(dm_tracer_adds_mass); + + dm.register_and_allocate("tracer_positive","",{num_tracers}); + auto dm_tracer_positive = dm.get("tracer_positive"); + tracer_positive.deep_copy_to(dm_tracer_positive); + + dm.register_and_allocate("hy_dens_cells" ,"",{nz,nens}); + dm.register_and_allocate("hy_pressure_cells","",{nz,nens}); + + int static constexpr DATA_SPEC_THERMAL = 0; + int static constexpr DATA_SPEC_SUPERCELL = 1; + int static constexpr DATA_SPEC_EXTERNAL = 2; + int data_spec; + if (coupler.option_exists("standalone_input_file")) { + #ifdef PAM_STANDALONE + std::string inFile = coupler.get_option( "standalone_input_file" ); + YAML::Node config = YAML::LoadFile(inFile); + std::string dataStr = config["initData"].as(); + if (dataStr == "thermal") { + data_spec = DATA_SPEC_THERMAL; + } else if (dataStr == "supercell") { + data_spec = DATA_SPEC_SUPERCELL; + } else if (dataStr == "external") { + data_spec = DATA_SPEC_EXTERNAL; + } else { + endrun("ERROR: Invalid data_spec"); + } + #endif + } else { + data_spec = DATA_SPEC_EXTERNAL; + } + + if (data_spec == DATA_SPEC_EXTERNAL) { + + } else { + real5d state ("state" ,num_state ,nz+2*hs,ny+2*hs,nx+2*hs,nens); + real5d tracers("tracers",num_tracers,nz+2*hs,ny+2*hs,nx+2*hs,nens); + + auto zmid = dm.get("vertical_midpoint_height"); + + if (data_spec == DATA_SPEC_SUPERCELL) { + + init_supercell( coupler , state , tracers ); + + } else if (data_spec == DATA_SPEC_THERMAL) { + + real2d hy_dens_cells ("hy_dens_cells" ,nz,nens); + real2d hy_pressure_cells("hy_pressure_cells",nz,nens); + + // Define quadrature weights and points for 3-point rules + const int nqpoints = 9; + SArray qpoints; + SArray qweights; + + TransformMatrices::get_gll_points ( qpoints ); + TransformMatrices::get_gll_weights( qweights ); + + // Compute hydrostatic background cell averages using quadrature + parallel_for( YAKL_AUTO_LABEL() , SimpleBounds<2>(nz,nens) , YAKL_LAMBDA (int k, int iens) { + hy_dens_cells (k,iens) = 0.; + hy_pressure_cells(k,iens) = 0.; + for (int kk=0; kk(nz,ny,nx,nens) , YAKL_LAMBDA (int k, int j, int i, int iens) { + for (int l=0; l < num_state ; l++) { state (l,hs+k,hs+j,hs+i,iens) = 0.; } + for (int l=0; l < num_tracers; l++) { tracers(l,hs+k,hs+j,hs+i,iens) = 0.; } + //Use Gauss-Legendre quadrature + for (int kk=0; kk gll_pts, gll_wts; + TransformMatrices::get_gll_points (gll_pts); + TransformMatrices::get_gll_weights(gll_wts); + + auto nens = coupler.get_nens(); + auto nx = coupler.get_nx(); + auto ny = coupler.get_ny(); + auto nz = coupler.get_nz(); + auto dx = coupler.get_dx(); + auto dy = coupler.get_dy(); + auto xlen = coupler.get_xlen(); + auto ylen = coupler.get_ylen(); + auto sim2d = ny == 1; + auto R_d = coupler.get_option("R_d" ); + auto R_v = coupler.get_option("R_v" ); + auto grav = coupler.get_option("grav" ); + auto gamma = coupler.get_option("gamma_d"); + auto C0 = coupler.get_option("C0" ); + auto idWV = coupler.get_option("idWV"); + auto num_tracers = coupler.get_num_tracers(); + + real2d hy_dens_cells ("hy_dens_cells" ,nz,nens); + real2d hy_pressure_cells("hy_pressure_cells",nz,nens); + + auto &dm = coupler.get_data_manager_device_readwrite(); + auto dz = dm.get("vertical_cell_dz"); + auto zmid = dm.get("vertical_midpoint_height"); + auto zint = dm.get("vertical_interface_height"); + + // Temporary arrays used to compute the initial state for high-CAPE supercell conditions + real4d quad_temp ("quad_temp" ,nz,ngll-1,ngll,nens); + real3d hyDensGLL ("hyDensGLL" ,nz,ngll,nens); + real3d hyDensThetaGLL ("hyDensThetaGLL" ,nz,ngll,nens); + real3d hyDensVapGLL ("hyDensVapGLL" ,nz,ngll,nens); + real3d hyPressureGLL ("hyPressureGLL" ,nz,ngll,nens); + real2d hyDensCells ("hyDensCells" ,nz,nens); + real2d hyDensThetaCells("hyDensThetaCells",nz,nens); + + // Compute quadrature term to integrate to get pressure at GLL points + parallel_for( YAKL_AUTO_LABEL() , SimpleBounds<4>(nz,ngll-1,ngll,nens) , + YAKL_LAMBDA (int k, int kk, int kkk, int iens) { + // Middle of this cell + real cellmid = zmid(k,iens); + // Bottom, top, and middle of the space between these two ngll GLL points + real ngll_b = cellmid + gll_pts(kk )*dz(k,iens); + real ngll_t = cellmid + gll_pts(kk+1)*dz(k,iens); + real ngll_m = 0.5_fp * (ngll_b + ngll_t); + // Compute grid spacing between these ngll GLL points + real ngll_dz = dz(k,iens) * ( gll_pts(kk+1) - gll_pts(kk) ); + // Compute the locate of this GLL point within the ngll GLL points + real zloc = ngll_m + ngll_dz * gll_pts(kkk); + // Compute full density at this location + real temp = init_supercell_temperature (zloc, z_0, z_trop, zint(nz,iens), T_0, T_trop, T_top); + real press_dry = init_supercell_pressure_dry(zloc, z_0, z_trop, zint(nz,iens), T_0, T_trop, T_top, p_0, R_d, grav); + real qvs = init_supercell_sat_mix_dry(press_dry, temp); + real relhum = init_supercell_relhum(zloc, z_0, z_trop); + if (relhum * qvs > 0.014_fp) relhum = 0.014_fp / qvs; + real qv = std::min( 0.014_fp , qvs*relhum ); + quad_temp(k,kk,kkk,iens) = -(1+qv)*grav/(R_d+qv*R_v)/temp; + }); + + // Compute pressure at GLL points + parallel_for( YAKL_AUTO_LABEL() , nens , YAKL_LAMBDA (int iens) { + hyPressureGLL(0,0,iens) = p_0; + for (int k=0; k < nz; k++) { + for (int kk=0; kk < ngll-1; kk++) { + real tot = 0; + for (int kkk=0; kkk < ngll; kkk++) { + tot += quad_temp(k,kk,kkk,iens) * gll_wts(kkk); + } + tot *= dz(k,iens) * ( gll_pts(kk+1) - gll_pts(kk) ); + hyPressureGLL(k,kk+1,iens) = hyPressureGLL(k,kk,iens) * exp( tot ); + if (kk == ngll-2 && k < nz-1) { + hyPressureGLL(k+1,0,iens) = hyPressureGLL(k,ngll-1,iens); + } + } + } + }); + + // Compute hydrostatic background state at GLL points + parallel_for( YAKL_AUTO_LABEL() , SimpleBounds<3>(nz,ngll,nens) , YAKL_LAMBDA (int k, int kk, int iens) { + real zloc = zmid(k,iens) + gll_pts(kk)*dz(k,iens); + real temp = init_supercell_temperature (zloc, z_0, z_trop, zint(nz,iens), T_0, T_trop, T_top); + real press_tmp = init_supercell_pressure_dry(zloc, z_0, z_trop, zint(nz,iens), T_0, T_trop, T_top, p_0, R_d, grav); + real qvs = init_supercell_sat_mix_dry(press_tmp, temp); + real relhum = init_supercell_relhum(zloc, z_0, z_trop); + if (relhum * qvs > 0.014_fp) relhum = 0.014_fp / qvs; + real qv = std::min( 0.014_fp , qvs*relhum ); + real press = hyPressureGLL(k,kk,iens); + real dens_dry = press / (R_d+qv*R_v) / temp; + real dens_vap = qv * dens_dry; + real dens = dens_dry + dens_vap; + real dens_theta = pow( press / C0 , 1._fp / gamma ); + hyDensGLL (k,kk,iens) = dens; + hyDensThetaGLL(k,kk,iens) = dens_theta; + hyDensVapGLL (k,kk,iens) = dens_vap; + }); + + // Compute hydrostatic background state over cells + parallel_for( YAKL_AUTO_LABEL() , SimpleBounds<2>(nz,nens) , YAKL_LAMBDA (int k, int iens) { + real press_tot = 0; + real dens_tot = 0; + real dens_vap_tot = 0; + real dens_theta_tot = 0; + for (int kk=0; kk < ngll; kk++) { + press_tot += hyPressureGLL (k,kk,iens) * gll_wts(kk); + dens_tot += hyDensGLL (k,kk,iens) * gll_wts(kk); + dens_vap_tot += hyDensVapGLL (k,kk,iens) * gll_wts(kk); + dens_theta_tot += hyDensThetaGLL(k,kk,iens) * gll_wts(kk); + } + real press = press_tot; + real dens = dens_tot; + real dens_vap = dens_vap_tot; + real dens_dry = dens - dens_vap; + real R = dens_dry / dens * R_d + dens_vap / dens * R_v; + real temp = press / (dens * R); + real qv = dens_vap / dens_dry; + real zloc = (k+0.5_fp)*dz(k,iens); + real press_tmp = init_supercell_pressure_dry(zloc, z_0, z_trop, zint(nz,iens), T_0, T_trop, T_top, p_0, R_d, grav); + real qvs = init_supercell_sat_mix_dry(press_tmp, temp); + real relhum = qv / qvs; + real T = temp - 273; + real a = 17.27; + real b = 237.7; + real tdew = b * ( a*T / (b + T) + log(relhum) ) / ( a - ( a*T / (b+T) + log(relhum) ) ); + // These are used in the rest of the model + for (int iens=0; iens < nens; iens++) { + hy_dens_cells (k,iens) = dens; + hy_pressure_cells(k,iens) = press; + } + }); + + // Initialize the state + parallel_for( YAKL_AUTO_LABEL() , SimpleBounds<4>(nz,ny,nx,nens) , YAKL_LAMBDA (int k, int j, int i, int iens) { + state(idR,hs+k,hs+j,hs+i,iens) = 0; + state(idU,hs+k,hs+j,hs+i,iens) = 0; + state(idV,hs+k,hs+j,hs+i,iens) = 0; + state(idW,hs+k,hs+j,hs+i,iens) = 0; + state(idT,hs+k,hs+j,hs+i,iens) = 0; + for (int tr=0; tr < num_tracers; tr++) { tracers(tr,hs+k,hs+j,hs+i,iens) = 0; } + real pres = hy_pressure_cells(k,iens); + state(idR,hs+k,hs+j,hs+i,iens) = hy_dens_cells(k,iens); + state(idT,hs+k,hs+j,hs+i,iens) = std::pow( pres/C0 , 1._fp/gamma ); + for (int kk=0; kk < ngll; kk++) { + for (int jj=0; jj < ngll; jj++) { + for (int ii=0; ii < ngll; ii++) { + real xloc = (i+0.5_fp)*dx + gll_pts(ii)*dx; + real yloc = (j+0.5_fp)*dy + gll_pts(jj)*dy; + real zloc = zmid(k,iens) + gll_pts(kk)*dz(k,iens); + + if (sim2d) yloc = ylen/2; + + real constexpr zs = 5000; + real constexpr us = 30; + real constexpr uc = 15; + real uvel = zloc < zs ? us*(zloc/zs)-uc : us-uc; + real vvel = 0; + real wvel = 0; + real dens_vap = hyDensVapGLL(k,kk,iens); + + real factor = gll_wts(ii) * gll_wts(jj) * gll_wts(kk); + state (idU ,hs+k,hs+j,hs+i,iens) += state(idR,hs+k,hs+j,hs+i,iens) * uvel * factor; + state (idV ,hs+k,hs+j,hs+i,iens) += state(idR,hs+k,hs+j,hs+i,iens) * vvel * factor; + state (idW ,hs+k,hs+j,hs+i,iens) += state(idR,hs+k,hs+j,hs+i,iens) * wvel * factor; + tracers(idWV,hs+k,hs+j,hs+i,iens) += dens_vap * factor; + } + } + } + }); + } + + + + // Convert dynamics state and tracers arrays to the coupler state and write to the coupler's data + void convert_dynamics_to_coupler( pam::PamCoupler &coupler , + realConst5d state , + realConst5d tracers ) const { + using yakl::c::parallel_for; + using yakl::c::SimpleBounds; + auto nens = coupler.get_nens(); + auto nx = coupler.get_nx(); + auto ny = coupler.get_ny(); + auto nz = coupler.get_nz(); + auto R_d = coupler.get_option("R_d" ); + auto R_v = coupler.get_option("R_v" ); + auto gamma = coupler.get_option("gamma_d"); + auto C0 = coupler.get_option("C0" ); + auto idWV = coupler.get_option("idWV"); + auto num_tracers = coupler.get_num_tracers(); + + auto &dm = coupler.get_data_manager_device_readwrite(); + auto tracer_adds_mass = dm.get("tracer_adds_mass"); + + // Get state from the coupler + auto dm_rho_d = dm.get("density_dry"); + auto dm_uvel = dm.get("uvel" ); + auto dm_vvel = dm.get("vvel" ); + auto dm_wvel = dm.get("wvel" ); + auto dm_temp = dm.get("temp" ); + + // Get tracers from the coupler + pam::MultiField dm_tracers; + auto tracer_names = coupler.get_tracer_names(); + for (int tr=0; tr < num_tracers; tr++) { dm_tracers.add_field( dm.get(tracer_names[tr]) ); } + + // Convert from state and tracers arrays to the coupler's data + parallel_for( YAKL_AUTO_LABEL() , SimpleBounds<4>(nz,ny,nx,nens) , YAKL_LAMBDA (int k, int j, int i, int iens) { + real rho = state(idR,hs+k,hs+j,hs+i,iens); + real u = state(idU,hs+k,hs+j,hs+i,iens) / rho; + real v = state(idV,hs+k,hs+j,hs+i,iens) / rho; + real w = state(idW,hs+k,hs+j,hs+i,iens) / rho; + real theta = state(idT,hs+k,hs+j,hs+i,iens) / rho; + real press = C0 * pow( rho*theta , gamma ); + real rho_v = tracers(idWV,hs+k,hs+j,hs+i,iens); + real rho_d = rho; + for (int tr=0; tr < num_tracers; tr++) { if (tracer_adds_mass(tr)) rho_d -= tracers(tr,hs+k,hs+j,hs+i,iens); } + real temp = press / ( rho_d * R_d + rho_v * R_v ); + dm_rho_d(k,j,i,iens) = rho_d; + dm_uvel (k,j,i,iens) = u; + dm_vvel (k,j,i,iens) = v; + dm_wvel (k,j,i,iens) = w; + dm_temp (k,j,i,iens) = temp; + for (int tr=0; tr < num_tracers; tr++) { dm_tracers(tr,k,j,i,iens) = tracers(tr,hs+k,hs+j,hs+i,iens); } + }); + } + + + + // Convert coupler's data to state and tracers arrays + void convert_coupler_to_dynamics( pam::PamCoupler &coupler , + real5d &state , + real5d &tracers ) const { + using yakl::c::parallel_for; + using yakl::c::SimpleBounds; + auto nens = coupler.get_nens(); + auto nx = coupler.get_nx(); + auto ny = coupler.get_ny(); + auto nz = coupler.get_nz(); + auto R_d = coupler.get_option("R_d" ); + auto R_v = coupler.get_option("R_v" ); + auto gamma_d = coupler.get_option("gamma_d"); + auto C0 = coupler.get_option("C0" ); + auto idWV = coupler.get_option("idWV"); + auto grav = coupler.get_option("grav"); + auto num_tracers = coupler.get_num_tracers(); + + auto &dm = coupler.get_data_manager_device_readwrite(); + auto tracer_adds_mass = dm.get("tracer_adds_mass"); + + // Get the coupler's state (as const because it's read-only) + auto dm_rho_d = dm.get("density_dry"); + auto dm_uvel = dm.get("uvel" ); + auto dm_vvel = dm.get("vvel" ); + auto dm_wvel = dm.get("wvel" ); + auto dm_temp = dm.get("temp" ); + + // Get the coupler's tracers (as const because it's read-only) + pam::MultiField dm_tracers; + auto tracer_names = coupler.get_tracer_names(); + for (int tr=0; tr < num_tracers; tr++) { dm_tracers.add_field( dm.get(tracer_names[tr]) ); } + + // Convert from the coupler's state to the dycore's state and tracers arrays. + // Compute domain-averaged pressure column + parallel_for( YAKL_AUTO_LABEL() , SimpleBounds<4>(nz,ny,nx,nens) , YAKL_LAMBDA (int k, int j, int i, int iens) { + real rho_d = dm_rho_d(k,j,i,iens); + real u = dm_uvel (k,j,i,iens); + real v = dm_vvel (k,j,i,iens); + real w = dm_wvel (k,j,i,iens); + real temp = dm_temp (k,j,i,iens); + real rho_v = dm_tracers(idWV,k,j,i,iens); + real press = rho_d * R_d * temp + rho_v * R_v * temp; + real rho = rho_d; + for (int tr=0; tr < num_tracers; tr++) { if (tracer_adds_mass(tr)) rho += dm_tracers(tr,k,j,i,iens); } + real theta = pow( press/C0 , 1._fp / gamma_d ) / rho; + state(idR,hs+k,hs+j,hs+i,iens) = rho; + state(idU,hs+k,hs+j,hs+i,iens) = rho * u; + state(idV,hs+k,hs+j,hs+i,iens) = rho * v; + state(idW,hs+k,hs+j,hs+i,iens) = rho * w; + state(idT,hs+k,hs+j,hs+i,iens) = rho * theta; + for (int tr=0; tr < num_tracers; tr++) { tracers(tr,hs+k,hs+j,hs+i,iens) = dm_tracers(tr,k,j,i,iens); } + }); + } + + + + void declare_current_profile_as_hydrostatic( pam::PamCoupler &coupler , bool use_gcm_data = false ) const { + using yakl::c::parallel_for; + using yakl::c::SimpleBounds; + auto nens = coupler.get_nens(); + auto nx = coupler.get_nx(); + auto ny = coupler.get_ny(); + auto nz = coupler.get_nz(); + auto gamma_d = coupler.get_option("gamma_d"); + auto C0 = coupler.get_option("C0" ); + auto idWV = coupler.get_option("idWV" ); + auto R_d = coupler.get_option("R_d"); + auto R_v = coupler.get_option("R_v"); + auto num_tracers = coupler.get_num_tracers(); + auto &dm = coupler.get_data_manager_device_readwrite(); + auto dz = dm.get("vertical_cell_dz"); + auto vert_weno_recon_lower = dm.get("vert_weno_recon_lower"); + auto vert_sten_to_coefs = dm.get("vert_sten_to_coefs"); + + auto grav_balance = coupler.get_option("balance_hydrostasis_with_gravity"); + + real5d state ("state" ,num_state ,nz+2*hs,ny+2*hs,nx+2*hs,nens); + real5d tracers("tracers",num_tracers,nz+2*hs,ny+2*hs,nx+2*hs,nens); + + if (use_gcm_data) { + auto gcm_rho_d = dm.get("gcm_density_dry"); + auto gcm_temp = dm.get("gcm_temp" ); + auto gcm_rho_v = dm.get("gcm_water_vapor"); + auto gcm_rho_c = dm.get("gcm_cloud_water"); + auto gcm_rho_i = dm.get("gcm_cloud_ice" ); + parallel_for( YAKL_AUTO_LABEL() , SimpleBounds<4>(nz,ny,nx,nens) , YAKL_LAMBDA (int k, int j, int i, int iens) { + real rho_d = gcm_rho_d(k,iens); + real rho_v = gcm_rho_v(k,iens); + real rho = gcm_rho_d(k,iens) + gcm_rho_v(k,iens) + gcm_rho_c(k,iens) + gcm_rho_i(k,iens); + real temp = gcm_temp(k,iens); + real p = (rho_d*R_d + rho_v*R_v)*temp; + real rho_theta = std::pow( p/C0 , 1._fp/gamma_d ); + state(idR,hs+k,hs+j,hs+i,iens) = rho; + state(idU,hs+k,hs+j,hs+i,iens) = 0; + state(idV,hs+k,hs+j,hs+i,iens) = 0; + state(idW,hs+k,hs+j,hs+i,iens) = 0; + state(idT,hs+k,hs+j,hs+i,iens) = rho_theta; + for (int tr=0; tr < num_tracers; tr++) { tracers(tr,hs+k,hs+j,hs+i,iens) = 0; } + tracers(idWV,hs+k,hs+j,hs+i,iens) = rho_v; + }); + } else { + convert_coupler_to_dynamics( coupler , state , tracers ); + } + + if (grav_balance) { + + SArray coefs_to_gll; + SArray idl; + real sigma; + TransformMatrices::coefs_to_gll_lower(coefs_to_gll ); + weno::wenoSetIdealSigma(idl,sigma); + + auto grav_var = coupler.get_data_manager_device_readwrite().get("variable_gravity"); + // Discretize interface pressure spatially exactly as we would discretize it in the solver + real4d pressure("pressure",nz+2*hs,ny+2*hs,nx+2*hs,nens); + parallel_for( YAKL_AUTO_LABEL() , SimpleBounds<4>(nz,ny,nx,nens) , YAKL_LAMBDA (int k, int j, int i, int iens) { + pressure(hs+k,hs+j,hs+i,iens) = C0*std::pow(state(idT,hs+k,hs+j,hs+i,iens),gamma_d); + state(idT,hs+k,hs+j,hs+i,iens) /= state(idR,hs+k,hs+j,hs+i,iens); // In prep for halo_exchange + if (j == 0 && i == 0) grav_var(k,iens) = 0; + }); + halo_exchange( coupler , state , tracers , pressure ); + real4d pint("pint",nz+1,ny,nx,nens); + parallel_for( YAKL_AUTO_LABEL() , SimpleBounds<4>(nz+1,ny,nx,nens) , YAKL_LAMBDA (int k, int j, int i, int iens) { + SArray stencil; + SArray s2c_loc[2]; + SArray wrl_loc[2]; + for (int i1=0; i1 < ord; i1++) { + for (int i2=0; i2 < ord; i2++) { + s2c_loc[0](i1,i2) = vert_sten_to_coefs(k ,i1,i2,iens); + s2c_loc[1](i1,i2) = vert_sten_to_coefs(k+1,i1,i2,iens); + } + } + for (int i1=0; i1 < hs; i1++) { + for (int i2=0; i2 < hs; i2++) { + for (int i3=0; i3 < hs; i3++) { + wrl_loc[0](i1,i2,i3) = vert_weno_recon_lower(k ,i1,i2,i3,iens); + wrl_loc[1](i1,i2,i3) = vert_weno_recon_lower(k+1,i1,i2,i3,iens); + } + } + } + int k_upw = 0; + for (int s=0; s < ord; s++) { stencil(s) = pressure(k+k_upw+s,hs+j,hs+i,iens); } + real pp_L = reconstruct(stencil,coefs_to_gll,s2c_loc[k_upw],wrl_loc[k_upw],idl,sigma,1-k_upw); + k_upw = 1; + for (int s=0; s < ord; s++) { stencil(s) = pressure(k+k_upw+s,hs+j,hs+i,iens); } + real pp_R = reconstruct(stencil,coefs_to_gll,s2c_loc[k_upw],wrl_loc[k_upw],idl,sigma,1-k_upw); + pint(k,j,i,iens) = 0.5_fp * (pp_L + pp_R); + }); + // Compute average column of variable gravity that provides exact hydrostatic balance + real r_nx_ny = 1./(nx*ny); + parallel_for( YAKL_AUTO_LABEL() , SimpleBounds<4>(nz,ny,nx,nens) , YAKL_LAMBDA (int k, int j, int i, int iens) { + real dens = state(idR,hs+k,hs+j,hs+i,iens); + yakl::atomicAdd( grav_var(k,iens) , -(pint(k+1,j,i,iens)-pint(k,j,i,iens))/(dens*dz(k,iens))*r_nx_ny ); + }); + + std::cout << grav_var; + + } else { + + auto hy_dens_cells = coupler.get_data_manager_device_readwrite().get("hy_dens_cells"); + auto hy_pressure_cells = coupler.get_data_manager_device_readwrite().get("hy_pressure_cells"); + hy_dens_cells = 0; + hy_pressure_cells = 0; + real r_nx_ny = 1./(nx*ny); + parallel_for( YAKL_AUTO_LABEL() , SimpleBounds<4>(nz,ny,nx,nens) , YAKL_LAMBDA (int k, int j, int i, int iens) { + real press = C0 * std::pow( state(idT,hs+k,hs+j,hs+i,iens) , gamma_d ); + yakl::atomicAdd( hy_pressure_cells(k,iens) , press *r_nx_ny ); + yakl::atomicAdd( hy_dens_cells (k,iens) , state(idR,hs+k,hs+j,hs+i,iens)*r_nx_ny ); + }); + + } + } + + + + // real2d solve_banded( realConst3d diags_in , realConst2d rhs ) const { + // int nbands = diags.extent(0); + // int n = diags.extent(1); + // int nens = diags.extent(2); + // if (nbands % 2 != 1) yakl::yakl_throw("ERROR: Number of bands must be odd"); + // if (n != rhs.extent(0)) yakl::yakl_throw("ERROR: Number of elements for diags and rhs must match"); + // if (nens != rhs.extent(1)) yakl::yakl_throw("ERROR: Number of ensembles must match between diags and rhs"); + // int h = (n-1)/2; + // auto diags_work = diags.createDeviceObject(); + // auto rhs_work = rhs .createDeviceObject(); + // parallel_for( YAKL_AUTO_LABEL , nens , YAKL_LAMBDA (int iens) { + // // Clear out lower values first + // for (int irow = 0 ; irow < n; irow++) { + // // Grab the value from the center band + // real r_diag_val = 1._fp/diags(h,irow,iens); + // for (int iband=h-1; iband >= 0; iband--) { + // real mult = -diags(iband,irow,iens)*r_diag_val; + // for (int iband2=0; iband2 <= n; iband2++) { + // } + // } + // } + // }); + // } + + + + char const * dycore_name() const { return "SSPRK3+WENO+FV A-grid"; } + + + + void finalize( pam::PamCoupler const &coupler ) const { } + +}; diff --git a/dynamics/awfl/Spatial_expl.h b/dynamics/awfl/Spatial_expl.h deleted file mode 100644 index 4ccbfd1d..00000000 --- a/dynamics/awfl/Spatial_expl.h +++ /dev/null @@ -1,2228 +0,0 @@ - -#pragma once - -#include "awfl_const.h" -#include "TransformMatrices.h" -#include "TransformMatrices_variable.h" -#include "WenoLimiter.h" -#include "idealized_profiles.h" -#include "MultipleFields.h" -#include "pam_coupler.h" - - -template -class Spatial_operator { -public: - - static_assert(nTimeDerivs == 1 , "ERROR: This Spatial class isn't setup to use nTimeDerivs > 1"); - - int static constexpr hs = (ord-1)/2; - int static constexpr num_state = 5; - int static constexpr max_tracers = 50; - - real Rd ; - real cp ; - real gamma; - real p0 ; - real C0 ; - real Rv ; - real grav ; - - int idWV; // Tracer index for water vapor (set in add_tracer by testing the tracer name against "water_vapor" - - real hydrostasis_parameters_sum; - - typedef real5d StateArr; // Array of state variables (rho, rho*u, rho*v, rho*w, and rho*theta) - typedef real5d TracerArr; // Array of tracers (total tracer mass) - - typedef real5d StateTendArr; // State tendencies - typedef real5d TracerTendArr; // Tracer tendencies - - // Hydrostatically balanced values for density, potential temperature, and pressure (cell-averages) - real3d hyDensSten; - real3d hyDensThetaSten; - - // Hydrostatically balanced values for density, potential temperature, and pressure (GLL points) - real3d hyDensGLL; - real3d hyPressureGLL; - real3d hyDensThetaGLL; - - // Matrices to transform DOFs from one form to another - SArray coefs_to_gll; - SArray coefs_to_deriv_gll; - SArray sten_to_gll; - SArray sten_to_coefs; - SArray sten_to_deriv_gll; - // WENO reconstruction matrices - SArray weno_recon_lower; - SArray idl; // Ideal weights for WENO - real sigma; // WENO sigma parameter (handicap high-order TV estimate) - // For ADER spatial derivative computation (ngll points -> coefs -> deriv -> ngll points) - SArray derivMatrix; - // For quadrature - SArray gllWts_ord; - SArray gllPts_ord; - SArray gllWts_ngll; - SArray gllPts_ngll; - - real2d vert_interface; - real2d vert_interface_ghost; - real3d vert_locs_normalized; - real2d dz; - real2d dz_ghost; - real4d vert_sten_to_gll; - real4d vert_sten_to_coefs; - real5d vert_weno_recon_lower; - - // For indexing into the state and state tendency arrays - int static constexpr idR = 0; // density perturbation - int static constexpr idU = 1; // u - int static constexpr idV = 2; // v - int static constexpr idW = 3; // w - int static constexpr idT = 4; // potential temperature perturbation - - // The two boundary condition options for each direction - int static constexpr BC_PERIODIC = 0; - int static constexpr BC_WALL = 1; - - // Options for initializing the data - int static constexpr DATA_SPEC_EXTERNAL = 0; - int static constexpr DATA_SPEC_THERMAL = 1; - int static constexpr DATA_SPEC_SUPERCELL = 2; - - bool sim2d; // Whether we're simulating in 2-D - - // Grid spacing in each dimension - real dx; - real dy; - - // Initial time step (used throughout the simulation) - real dtInit; - - // Which direction we're passing through for a given time step (x,y,z) or (z,y,x) - // For Strang splitting - bool dimSwitch; - - int num_tracers; // Number of tracers - std::vector tracer_name; // Name of each tracer - std::vector tracer_desc; // Description of each tracer - bool1d tracer_pos; // Whether each tracer is positive-definite - bool1d tracer_adds_mass; // Whether each tracer adds mass (otherwise it's passive) - - ////////////////////////////////////// - // Values read from input file - ////////////////////////////////////// - // Number of ensembles to simulate at the same time - int nens; - // Number of cells in each direction - int nx; - int ny; - int nz; - // Length of the domain in each direction (m) - real xlen; - real ylen; - real1d zbot; - real1d ztop; - // Whether to use WENO for scalars and also for winds - bool weno_scalars; - bool weno_winds; - // How to initialize the data - int data_spec; - - - // When this class is created, initialize num_tracers to zero - Spatial_operator() { - num_tracers = 0; - idWV = -1; - } - - - - // Make sure it's odd-order-accurate - static_assert(ord%2 == 1,"ERROR: ord must be an odd integer"); - - void convert_dynamics_to_coupler_state( pam::PamCoupler &coupler , realConst5d state , realConst5d tracers ) const { - using yakl::c::parallel_for; - using yakl::c::SimpleBounds; - - auto &dm = coupler.get_data_manager_device_readwrite(); - - real4d dm_dens_dry = dm.get( "density_dry" ); - real4d dm_uvel = dm.get( "uvel" ); - real4d dm_vvel = dm.get( "vvel" ); - real4d dm_wvel = dm.get( "wvel" ); - real4d dm_temp = dm.get( "temp" ); - - YAKL_SCOPE( C0 , this->C0 ); - YAKL_SCOPE( gamma , this->gamma ); - YAKL_SCOPE( num_tracers , this->num_tracers ); - YAKL_SCOPE( Rd , this->Rd ); - YAKL_SCOPE( Rv , this->Rv ); - YAKL_SCOPE( tracer_adds_mass , this->tracer_adds_mass ); - YAKL_SCOPE( idWV , this->idWV ); - - pam::MultipleFields dm_tracers; - for (int tr = 0; tr < num_tracers; tr++) { - auto trac = dm.get( tracer_name[tr] ); - dm_tracers.add_field( trac ); - } - - parallel_for( "Spatial.h d2c" , SimpleBounds<4>(nz,ny,nx,nens) , YAKL_LAMBDA (int k, int j, int i, int iens) { - real dens = state(idR,hs+k,hs+j,hs+i,iens); - real uvel = state(idU,hs+k,hs+j,hs+i,iens) / dens; - real vvel = state(idV,hs+k,hs+j,hs+i,iens) / dens; - real wvel = state(idW,hs+k,hs+j,hs+i,iens) / dens; - real theta = state(idT,hs+k,hs+j,hs+i,iens) / dens; - real pressure = C0 * pow( dens*theta , gamma ); - real dens_vap = tracers(idWV,hs+k,hs+j,hs+i,iens); - real dens_dry = dens; - for (int tr=0; tr < num_tracers; tr++) { - if (tracer_adds_mass(tr)) dens_dry -= tracers(tr,hs+k,hs+j,hs+i,iens); - } - real temp = pressure / ( dens_dry * Rd + dens_vap * Rv ); - dm_dens_dry(k,j,i,iens) = dens_dry; - dm_uvel (k,j,i,iens) = uvel; - dm_vvel (k,j,i,iens) = vvel; - dm_wvel (k,j,i,iens) = wvel; - dm_temp (k,j,i,iens) = temp; - for (int tr=0; tr < num_tracers; tr++) { - dm_tracers(tr,k,j,i,iens) = tracers(tr,hs+k,hs+j,hs+i,iens); - } - }); - } - - - - YAKL_INLINE static real hydrostatic_dens_theta( realConst3d hy_params , real z , real z0 , real dz , - int k, int iens , real C0 , real gamma ) { - real p = pam::hydrostatic_pressure( hy_params , z , z0 , dz , k , iens ); - // p = C0*(rho*theta)^gamma - return pow(p/C0,1._fp/gamma); - } - - - - void convert_coupler_state_to_dynamics( pam::PamCoupler const &coupler , real5d const &state , real5d const &tracers ) { - using yakl::c::parallel_for; - using yakl::c::SimpleBounds; - - auto &dm = coupler.get_data_manager_device_readonly(); - auto hy_params = dm.get("hydrostasis_parameters"); - - YAKL_SCOPE( hyDensSten , this->hyDensSten ); - YAKL_SCOPE( hyDensThetaSten , this->hyDensThetaSten ); - YAKL_SCOPE( hyPressureGLL , this->hyPressureGLL ); - YAKL_SCOPE( hyDensGLL , this->hyDensGLL ); - YAKL_SCOPE( hyDensThetaGLL , this->hyDensThetaGLL ); - YAKL_SCOPE( C0 , this->C0 ); - YAKL_SCOPE( gamma , this->gamma ); - YAKL_SCOPE( num_tracers , this->num_tracers ); - YAKL_SCOPE( Rd , this->Rd ); - YAKL_SCOPE( Rv , this->Rv ); - YAKL_SCOPE( tracer_adds_mass , this->tracer_adds_mass ); - YAKL_SCOPE( zbot , this->zbot ); - YAKL_SCOPE( ztop , this->ztop ); - YAKL_SCOPE( gllPts_ngll , this->gllPts_ngll ); - YAKL_SCOPE( dz , this->dz ); - YAKL_SCOPE( vert_interface , this->vert_interface ); - YAKL_SCOPE( idWV , this->idWV ); - YAKL_SCOPE( grav , this->grav ); - YAKL_SCOPE( vert_interface_ghost , this->vert_interface_ghost ); - YAKL_SCOPE( dz_ghost , this->dz_ghost ); - - auto dm_dens_dry = dm.get( "density_dry" ); - auto dm_uvel = dm.get( "uvel" ); - auto dm_vvel = dm.get( "vvel" ); - auto dm_wvel = dm.get( "wvel" ); - auto dm_temp = dm.get( "temp" ); - - pam::MultipleFields dm_tracers; - for (int tr = 0; tr < num_tracers; tr++) { - auto trac = dm.get( tracer_name[tr] ); - dm_tracers.add_field( trac ); - } - - // If hydrostasis in the coupler has changed, then we need to re-compute - // hydrostatically balanced cells and GLL points for the dycore's time step - real tmp = yakl::intrinsics::sum(hy_params); - if (tmp != hydrostasis_parameters_sum) { - SArray gll_pts, gll_wts; - TransformMatrices::get_gll_points ( gll_pts ); - TransformMatrices::get_gll_weights( gll_wts ); - - // Compute new cell averages and GLL point values for hydrostasis - hydrostasis_parameters_sum = tmp; - parallel_for( "Spatial.h new hydrostasis" , SimpleBounds<3>(nz,ord,nens) , YAKL_LAMBDA (int k, int kk, int iens) { - real r = 0; - real rt = 0; - for (int l=0; l < 9; l++) { - real zloc = vert_interface_ghost(k+kk,iens) + 0.5_fp*dz_ghost(k+kk,iens) + gll_pts(l)*dz_ghost(k+kk,iens); - real z0 = vert_interface(k,iens) + 0.5_fp*dz(k,iens); - real wt = gll_wts(l); - r += pam::hydrostatic_density (hy_params,zloc,z0,dz(k,iens),k,iens ,grav) * wt; - rt += hydrostatic_dens_theta(hy_params,zloc,z0,dz(k,iens),k,iens,C0,gamma ) * wt; - } - hyDensSten (k,kk,iens) = r; - hyDensThetaSten(k,kk,iens) = rt; - }); - - parallel_for( "Spatial.h new hydrostasis" , SimpleBounds<3>(nz,ngll,nens) , YAKL_LAMBDA (int k, int kk, int iens) { - real zloc = vert_interface(k,iens) + 0.5_fp*dz(k,iens) + gllPts_ngll(kk)*dz(k,iens); - real z0 = vert_interface(k,iens) + 0.5_fp*dz(k,iens); - hyPressureGLL (k,kk,iens) = pam::hydrostatic_pressure (hy_params,zloc,z0,dz(k,iens),k,iens ); - hyDensGLL (k,kk,iens) = pam::hydrostatic_density (hy_params,zloc,z0,dz(k,iens),k,iens ,grav); - hyDensThetaGLL(k,kk,iens) = hydrostatic_dens_theta(hy_params,zloc,z0,dz(k,iens),k,iens,C0,gamma ); - }); - } - - parallel_for( "Spatial.h c2d" , SimpleBounds<4>(nz,ny,nx,nens) , YAKL_LAMBDA (int k, int j, int i, int iens) { - for (int tr=0; tr < num_tracers; tr++) { - tracers(tr,hs+k,hs+j,hs+i,iens) = dm_tracers(tr,k,j,i,iens); - } - real dens_dry = dm_dens_dry (k,j,i,iens); - real uvel = dm_uvel (k,j,i,iens); - real vvel = dm_vvel (k,j,i,iens); - real wvel = dm_wvel (k,j,i,iens); - real temp = dm_temp (k,j,i,iens); - real dens_vap = tracers(idWV,hs+k,hs+j,hs+i,iens); - real dens = dens_dry; - for (int tr=0; tr < num_tracers; tr++) { - if (tracer_adds_mass(tr)) dens += tracers(tr,hs+k,hs+j,hs+i,iens); - } - real pressure = dens_dry * Rd * temp + dens_vap * Rv * temp; - real theta = pow( pressure / C0 , 1._fp / gamma ) / dens; - state(idR,hs+k,hs+j,hs+i,iens) = dens; - state(idU,hs+k,hs+j,hs+i,iens) = dens * uvel; - state(idV,hs+k,hs+j,hs+i,iens) = dens * vvel; - state(idW,hs+k,hs+j,hs+i,iens) = dens * wvel; - state(idT,hs+k,hs+j,hs+i,iens) = dens * theta; - }); - } - - - - real5d createStateArr() const { - return real5d("stateArr",num_state,nz+2*hs,ny+2*hs,nx+2*hs,nens); - } - - - - real5d createTracerArr() const { - return real5d("tracerArr",num_tracers,nz+2*hs,ny+2*hs,nx+2*hs,nens); - } - - - - real5d createStateTendArr() const { - return real5d("stateTendArr",num_state,nz,ny,nx,nens); - } - - - - real5d createTracerTendArr() const { - return real5d("tracerTendArr",num_tracers,nz,ny,nx,nens); - } - - - - // Number of operator splittinng steps to use - // Normally this would be 3, but the z-directly CFL is reduced because of how the fluxes are - // handled in the presence of a solid wall boundary condition. I'm looking into how to fix this - int numSplit() const { - if (sim2d) { - return 2; - } else { - return 3; - } - } - - - - // Given the model data and CFL value, compute the maximum stable time step - real compute_time_step(pam::PamCoupler const &coupler, real cfl_in = -1) { - using yakl::c::parallel_for; - using yakl::c::SimpleBounds; - - real cfl = cfl_in; - if (cfl < 0) cfl = 0.75; - - // If we've already computed the time step, then don't compute it again - if (dtInit <= 0) { - YAKL_SCOPE( dx , this->dx ); - YAKL_SCOPE( dy , this->dy ); - YAKL_SCOPE( dz , this->dz ); - YAKL_SCOPE( gamma , this->gamma ); - YAKL_SCOPE( Rd , this->Rd ); - YAKL_SCOPE( Rv , this->Rv ); - YAKL_SCOPE( sim2d , this->sim2d ); - - auto &dm = coupler.get_data_manager_device_readonly(); - - // Convert data from DataManager to state and tracers array for convenience - auto dm_dens_dry = dm.get( "density_dry" ); - auto dm_uvel = dm.get( "uvel" ); - auto dm_vvel = dm.get( "vvel" ); - auto dm_wvel = dm.get( "wvel" ); - auto dm_temp = dm.get( "temp" ); - auto dm_dens_vap = dm.get( "water_vapor" ); - - // Allocate a 3-D array for the max stable time steps (we'll use this for a reduction later) - real4d dt3d("dt3d",nz,ny,nx,nens); - - // Loop through the cells, calculate the max stable time step for each cell - parallel_for( "Spatial.h compute_time_step" , SimpleBounds<4>(nz,ny,nx,nens) , - YAKL_LAMBDA (int k, int j, int i, int iens) { - real rho_d = dm_dens_dry(k,j,i,iens); - real u = dm_uvel (k,j,i,iens); - real v = dm_vvel (k,j,i,iens); - real w = dm_wvel (k,j,i,iens); - real temp = dm_temp (k,j,i,iens); - real rho_v = dm_dens_vap(k,j,i,iens); - real p = Rd * rho_d * temp + Rv * rho_v * temp; - // This neglects non-wv mass-adding tracers, but these are small, and their lack only increases cs - // Thus the resulting time step is conservative w/r to these missing masses, which is more stable - real cs = sqrt(gamma*p/(rho_v+rho_d)); - - // Compute the maximum stable time step in each direction - real udt = cfl * dx / max( abs(u-cs) , abs(u+cs) ); - real vdt = cfl * dy / max( abs(v-cs) , abs(v+cs) ); - if (sim2d) vdt = std::numeric_limits::max(); - real wdt = cfl * dz(k,iens) / max( abs(w-cs) , abs(w+cs) ); - - // Compute the min of the max stable time steps - dt3d(k,j,i,iens) = min( min(udt,vdt) , wdt ); - }); - // Store to dtInit so we don't have to compute this again - dtInit = yakl::intrinsics::minval( dt3d ); - } - - return dtInit; - } - - - - // Initialize crap needed by recon() - void init(pam::PamCoupler &coupler) { - using yakl::c::parallel_for; - using yakl::c::SimpleBounds; - - using yakl::intrinsics::matmul_cr; - - this->nens = coupler.get_nens(); - this->nx = coupler.get_nx(); - this->ny = coupler.get_ny(); - this->xlen = coupler.get_xlen(); - this->ylen = coupler.get_ylen(); - this->num_tracers = coupler.get_num_tracers(); - - this->hydrostasis_parameters_sum = 0; - - this->Rd = coupler.get_option("R_d" ); - this->cp = coupler.get_option("cp_d"); - this->p0 = coupler.get_option("p0" ); - this->Rv = coupler.get_option("R_v" ); - this->grav = coupler.get_option("grav"); - this->gamma = cp / (cp-Rd); - real kappa = Rd/cp; - this->C0 = pow( Rd * pow( p0 , -kappa ) , gamma ); - - // Allocate device arrays for whether tracers are positive-definite or add mass - tracer_pos = bool1d("tracer_pos" ,num_tracers); - tracer_adds_mass = bool1d("tracer_adds_mass",num_tracers); - boolHost1d tracer_pos_host ("tracer_pos_host" ,num_tracers); - boolHost1d tracer_adds_mass_host("tracer_adds_mass_host",num_tracers); - - std::vector tracer_names_loc = coupler.get_tracer_names(); - bool water_vapor_found = false; - for (int tr=0; tr < num_tracers; tr++) { - bool found, positive, adds_mass; - std::string desc; - coupler.get_tracer_info( tracer_names_loc[tr] , desc , found , positive , adds_mass ); - tracer_name.push_back(tracer_names_loc[tr]); - tracer_desc.push_back(desc); - tracer_pos_host (tr) = positive ; - tracer_adds_mass_host(tr) = adds_mass; - if (tracer_names_loc[tr] == std::string("water_vapor")) { - idWV = tr; - water_vapor_found = true; - } - } - if (! water_vapor_found) endrun("ERROR: processed registered tracers, and water_vapor was not found"); - - tracer_pos_host .deep_copy_to(tracer_pos ); - tracer_adds_mass_host.deep_copy_to(tracer_adds_mass); - fence(); - - // Inialize time step to zero, and dimensional splitting switch - dtInit = 0; - dimSwitch = true; - - // If inFile is empty, then we aren't reading in an input file - if (coupler.option_exists("standalone_input_file")) { - #ifdef PAM_STANDALONE - std::string inFile = coupler.get_option( "standalone_input_file" ); - - // Read the YAML input file - YAML::Node config = YAML::LoadFile(inFile); - - // Read whether we're doing WENO limiting on scalars and winds - weno_scalars = config["weno_scalars"].as(); - weno_winds = config["weno_winds"].as(); - - // Read the data initialization option - std::string dataStr = config["initData"].as(); - if (dataStr == "thermal") { - data_spec = DATA_SPEC_THERMAL; - } else if (dataStr == "supercell") { - data_spec = DATA_SPEC_SUPERCELL; - } else if (dataStr == "external") { - data_spec = DATA_SPEC_EXTERNAL; - } else { - endrun("ERROR: Invalid data_spec"); - } - #endif - } else { - weno_scalars = true; - weno_winds = true; - data_spec = DATA_SPEC_EXTERNAL; - } - - // Determine whether this is a 2-D simulation - sim2d = ny == 1; - - // Store vertical cell interface heights in the data manager - auto &dm = coupler.get_data_manager_device_readonly(); - auto zint = dm.get("vertical_interface_height"); - - nz = coupler.get_nz(); - - // Get the height of the z-dimension - zbot = real1d("zbot",nens); - ztop = real1d("ztop",nens); - YAKL_SCOPE( zbot , this->zbot ); - YAKL_SCOPE( ztop , this->ztop ); - YAKL_SCOPE( nz , this->nz ); - parallel_for( "Spatial.h init 1" , nens , YAKL_LAMBDA (int iens) { - zbot(iens) = zint(0 ,iens); - ztop(iens) = zint(nz,iens); - }); - - vert_interface = real2d("vert_interface" ,nz+1 ,nens); - vert_interface_ghost = real2d("vert_interface_ghost",nz+2*hs+1 ,nens); - vert_locs_normalized = real3d("vert_locs_normalized",nz,ord+1 ,nens); - dz = real2d("dz" ,nz ,nens); - dz_ghost = real2d("dz_ghost" ,nz+2*hs ,nens); - vert_sten_to_gll = real4d("vert_sten_to_gll" ,nz,ord,ngll,nens); - vert_sten_to_coefs = real4d("vert_sten_to_coefs" ,nz,ord,ord ,nens); - vert_weno_recon_lower = real5d("vert_weno_recon_lower",nz,hs+1,hs+1,hs+1,nens); - - YAKL_SCOPE( vert_interface , this->vert_interface ); - YAKL_SCOPE( vert_interface_ghost , this->vert_interface_ghost ); - YAKL_SCOPE( vert_locs_normalized , this->vert_locs_normalized ); - YAKL_SCOPE( dz , this->dz ); - YAKL_SCOPE( dz_ghost , this->dz_ghost ); - YAKL_SCOPE( vert_sten_to_gll , this->vert_sten_to_gll ); - YAKL_SCOPE( vert_sten_to_coefs , this->vert_sten_to_coefs ); - YAKL_SCOPE( vert_weno_recon_lower , this->vert_weno_recon_lower ); - - zint.deep_copy_to(vert_interface); - - parallel_for( "Spatial.h init 1" , SimpleBounds<2>(nz,nens) , YAKL_LAMBDA (int k, int iens) { - dz(k,iens) = vert_interface(k+1,iens) - vert_interface(k,iens); - }); - - parallel_for( "Spatial.h init 2" , SimpleBounds<2>(nz+2*hs,nens) , YAKL_LAMBDA (int k, int iens) { - if (k >= hs && k < hs+nz) { - dz_ghost(k,iens) = dz(k-hs,iens); - } else if (k < hs) { - dz_ghost(k,iens) = dz(0,iens); - } else if (k >= hs+nz) { - dz_ghost(k,iens) = dz(nz-1,iens); - } - }); - - parallel_for( "Spatial.h init 3" , nens , YAKL_LAMBDA (int iens) { - vert_interface_ghost(0,iens) = vert_interface(0,iens) - hs*dz(0,iens); - for (int k=1; k < nz+2*hs+1; k++) { - vert_interface_ghost(k,iens) = vert_interface_ghost(k-1,iens) + dz_ghost(k-1,iens); - } - }); - - auto vint_host = vert_interface_ghost .createHostCopy(); - auto vert_s2g_host = vert_sten_to_gll .createHostCopy(); - auto vert_s2c_host = vert_sten_to_coefs .createHostCopy(); - auto vert_weno_host = vert_weno_recon_lower.createHostCopy(); - auto vert_locs_host = vert_locs_normalized .createHostCopy(); - - SArray c2g; - TransformMatrices::coefs_to_gll_lower(c2g); - - for (int k=0; k < nz; k++) { - for (int iens = 0; iens < nens; iens++) { - // Store stencil locations - SArray locs; - for (int kk=0; kk < ord+1; kk++) { - locs(kk) = vint_host(k+kk,iens); - } - - // Normalize stencil locations - double zmid = ( locs(hs+1) + locs(hs) ) / 2; - double dzmid = locs(hs+1) - locs(hs); - for (int kk=0; kk < ord+1; kk++) { - locs(kk) = ( locs(kk) - zmid ) / dzmid; - vert_locs_host(k,kk,iens) = locs(kk); - } - - // Compute reconstruction matrices - SArray s2c_var_in; - SArray weno_recon_lower_var; - TransformMatrices_variable::sten_to_coefs_variable(locs,s2c_var_in); - TransformMatrices_variable::weno_lower_sten_to_coefs(locs,weno_recon_lower_var); - SArray s2c_var; - for (int jj=0; jj < ord; jj++) { - for (int ii=0; ii < ord; ii++) { - s2c_var(jj,ii) = s2c_var_in(jj,ii); - } - } - auto s2g_var = matmul_cr( c2g , s2c_var ); - - // Store reconstruction matrices - for (int jj=0; jj < ord; jj++) { - for (int ii=0; ii < ord; ii++) { - vert_s2c_host(k,jj,ii,iens) = s2c_var(jj,ii); - } - } - - for (int jj=0; jj < ord; jj++) { - for (int ii=0; ii < ngll; ii++) { - vert_s2g_host(k,jj,ii,iens) = s2g_var(jj,ii); - } - } - - for (int kk=0; kk < hs+1; kk++) { - for (int jj=0; jj < hs+1; jj++) { - for (int ii=0; ii < hs+1; ii++) { - vert_weno_host(k,kk,jj,ii,iens) = weno_recon_lower_var(kk,jj,ii); - } - } - } - - } - } - - vert_s2g_host .deep_copy_to(vert_sten_to_gll ); - vert_s2c_host .deep_copy_to(vert_sten_to_coefs ); - vert_weno_host.deep_copy_to(vert_weno_recon_lower); - vert_locs_host.deep_copy_to(vert_locs_normalized ); - - // Compute the grid spacing in each dimension - dx = xlen/nx; - dy = ylen/ny; - - // Store the WENO reconstruction matrices - TransformMatrices::weno_lower_sten_to_coefs(this->weno_recon_lower); - - // Block exists to avoid name mangling stufff - { - SArray s2c; // Converts ord stencil cell averages to ord coefficients - SArray c2g_lower; // Converts ord coefficients to ngll GLL points - SArray c2d; // Converts ord coefficients to order differentiated coefficients - - TransformMatrices::sten_to_coefs (s2c ); - TransformMatrices::coefs_to_gll_lower(c2g_lower); - TransformMatrices::coefs_to_deriv (c2d ); - - this->coefs_to_gll = c2g_lower; - this->coefs_to_deriv_gll = matmul_cr( c2g_lower , c2d ); - this->sten_to_coefs = s2c; - this->sten_to_gll = matmul_cr( c2g_lower , s2c ); - this->sten_to_deriv_gll = matmul_cr( c2g_lower , matmul_cr( c2d , s2c ) ); - } - // Store ader derivMatrix - { - SArray g2c; // Converts ngll GLL points to ngll coefficients - SArray c2d; // Converts ngll coefficients to ngll differentiated coefficients - SArray c2g; // Converts ngll coefficients to ngll GLL points - - TransformMatrices::gll_to_coefs (g2c); - TransformMatrices::coefs_to_deriv(c2d); - TransformMatrices::coefs_to_gll (c2g); - - this->derivMatrix = matmul_cr( c2g , matmul_cr( c2d , g2c ) ); - } - // Store quadrature weights using ord GLL points - TransformMatrices::get_gll_points (this->gllPts_ord); - TransformMatrices::get_gll_weights(this->gllWts_ord); - // Store quadrature weights using ngll GLL points - TransformMatrices::get_gll_points (this->gllPts_ngll); - TransformMatrices::get_gll_weights(this->gllWts_ngll); - - // Store WENO ideal weights and sigma value - weno::wenoSetIdealSigma(this->idl,this->sigma); - - // Allocate data - hyDensSten = real3d("hyDensSten ",nz,ord,nens); - hyDensThetaSten = real3d("hyDensThetaSten ",nz,ord,nens); - hyDensGLL = real3d("hyDensGLL ",nz,ngll,nens); - hyPressureGLL = real3d("hyPressureGLL ",nz,ngll,nens); - hyDensThetaGLL = real3d("hyDensThetaGLL ",nz,ngll,nens); - - init_idealized_state_and_tracers( coupler ); - - } - - - - // Initialize the state - void init_idealized_state_and_tracers( pam::PamCoupler &coupler ) { - using yakl::c::parallel_for; - using yakl::c::SimpleBounds; - - YAKL_SCOPE( nx , this->nx ); - YAKL_SCOPE( ny , this->ny ); - YAKL_SCOPE( nz , this->nz ); - YAKL_SCOPE( dx , this->dx ); - YAKL_SCOPE( dy , this->dy ); - YAKL_SCOPE( dz , this->dz ); - YAKL_SCOPE( gllPts_ord , this->gllPts_ord ); - YAKL_SCOPE( gllWts_ord , this->gllWts_ord ); - YAKL_SCOPE( gllPts_ngll , this->gllPts_ngll ); - YAKL_SCOPE( gllWts_ngll , this->gllWts_ngll ); - YAKL_SCOPE( data_spec , this->data_spec ); - YAKL_SCOPE( sim2d , this->sim2d ); - YAKL_SCOPE( xlen , this->xlen ); - YAKL_SCOPE( ylen , this->ylen ); - YAKL_SCOPE( Rd , this->Rd ); - YAKL_SCOPE( Rv , this->Rv ); - YAKL_SCOPE( cp , this->cp ); - YAKL_SCOPE( gamma , this->gamma ); - YAKL_SCOPE( p0 , this->p0 ); - YAKL_SCOPE( C0 , this->C0 ); - YAKL_SCOPE( vert_interface , this->vert_interface ); - YAKL_SCOPE( num_tracers , this->num_tracers ); - YAKL_SCOPE( idWV , this->idWV ); - YAKL_SCOPE( grav , this->grav ); - - // If data's being specified by the driver externally, then there's nothing to do here - if (data_spec == DATA_SPEC_EXTERNAL) return; - - real5d state = createStateArr(); - real5d tracers = createTracerArr(); - - // If the data_spec is thermal or ..., then initialize the domain with Exner pressure-based hydrostasis - // This is mostly to make plotting potential temperature perturbation easier for publications - if (data_spec == DATA_SPEC_THERMAL) { - - // Compute the state - parallel_for( "Spatial.h init_state 3" , SimpleBounds<4>(nz,ny,nx,nens) , - YAKL_LAMBDA (int k, int j, int i, int iens) { - state(idR,hs+k,hs+j,hs+i,iens) = 0; - state(idU,hs+k,hs+j,hs+i,iens) = 0; - state(idV,hs+k,hs+j,hs+i,iens) = 0; - state(idW,hs+k,hs+j,hs+i,iens) = 0; - state(idT,hs+k,hs+j,hs+i,iens) = 0; - for (int kk=0; kk(nz,ny,nx,nens) , - YAKL_LAMBDA (int k, int j, int i, int iens) { - for (int tr=0; tr < num_tracers; tr++) { tracers(tr,hs+k,hs+j,hs+i,iens) = 0; } - // Loop over quadrature points - for (int kk=0; kk(nz,ny,nx,nens) , - YAKL_LAMBDA (int k, int j, int i, int iens) { - // Add tracer density to dry density if it adds mass - real rho_dry = state(idR,hs+k,hs+j,hs+i,iens); - state(idR,hs+k,hs+j,hs+i,iens) += tracers(idWV,hs+k,hs+j,hs+i,iens); - real rho_moist = state(idR,hs+k,hs+j,hs+i,iens); - - // Adjust momenta for moist density - state(idU,hs+k,hs+j,hs+i,iens) = state(idU,hs+k,hs+j,hs+i,iens) / rho_dry * rho_moist; - state(idV,hs+k,hs+j,hs+i,iens) = state(idV,hs+k,hs+j,hs+i,iens) / rho_dry * rho_moist; - state(idW,hs+k,hs+j,hs+i,iens) = state(idW,hs+k,hs+j,hs+i,iens) / rho_dry * rho_moist; - - // Compute the dry temperature (same as the moist temperature) - real rho_theta_dry = state(idT,hs+k,hs+j,hs+i,iens); - real press = C0*pow(rho_theta_dry,gamma); // Dry pressure - real temp = press / Rd / rho_dry; // Temp (same dry or moist) - - // Compute moist theta - real rho_v = tracers(idWV,hs+k,hs+j,hs+i,iens); - real R_moist = Rd * (rho_dry / rho_moist) + Rv * (rho_v / rho_moist); - real press_moist = rho_moist * R_moist * temp; - real rho_theta_moist = pow( press_moist / C0 , 1._fp/gamma ); - - // Compute moist rho*theta - state(idT,hs+k,hs+j,hs+i,iens) = rho_theta_moist; - - for (int tr = 0 ; tr < num_tracers ; tr++) { - tracers(tr,hs+k,hs+j,hs+i,iens) = tracers(tr,hs+k,hs+j,hs+i,iens) / rho_dry * rho_moist; - } - }); - - } // if (data_spec == DATA_SPEC_THERMAL) - - convert_dynamics_to_coupler_state( coupler , state , tracers ); - } - - - - // Compute state and tendency time derivatives from the state - void computeTendencies( real5d const &state , real5d const &stateTend , - real5d const &tracers , real5d const &tracerTend , - real &dt , int splitIndex ) const { - if (sim2d) { - if (dimSwitch) { - if (splitIndex == 0) { - computeTendenciesX( state , stateTend , tracers , tracerTend , dt ); - } else if (splitIndex == 1) { - computeTendenciesZ( state , stateTend , tracers , tracerTend , dt ); - } - } else { - if (splitIndex == 0) { - computeTendenciesZ( state , stateTend , tracers , tracerTend , dt ); - } else if (splitIndex == 1) { - computeTendenciesX( state , stateTend , tracers , tracerTend , dt ); - } - } - } else { - if (dimSwitch) { - if (splitIndex == 0) { - computeTendenciesX( state , stateTend , tracers , tracerTend , dt ); - } else if (splitIndex == 1) { - computeTendenciesY( state , stateTend , tracers , tracerTend , dt ); - } else if (splitIndex == 2) { - computeTendenciesZ( state , stateTend , tracers , tracerTend , dt ); - } - } else { - if (splitIndex == 0) { - computeTendenciesZ( state , stateTend , tracers , tracerTend , dt ); - } else if (splitIndex == 1) { - computeTendenciesY( state , stateTend , tracers , tracerTend , dt ); - } else if (splitIndex == 2) { - computeTendenciesX( state , stateTend , tracers , tracerTend , dt ); - } - } - } - } // computeTendencies - - - - void switch_directions() { - dimSwitch = ! dimSwitch; - } - - - - YAKL_INLINE static int wrapx(int i, int ii, int nx) { - int ret = i+ii; - if (ret < hs+0 ) ret += nx; - if (ret > hs+nx-1) ret -= nx; - return ret; - } - - - - void computeTendenciesX( real5d const &state , real5d const &stateTend , - real5d const &tracers , real5d const &tracerTend , - real &dt ) const { - using yakl::c::parallel_for; - using yakl::c::SimpleBounds; - - YAKL_SCOPE( nx , this->nx ); - YAKL_SCOPE( weno_scalars , this->weno_scalars ); - YAKL_SCOPE( weno_winds , this->weno_winds ); - YAKL_SCOPE( c2g , this->coefs_to_gll ); - YAKL_SCOPE( s2g , this->sten_to_gll ); - YAKL_SCOPE( s2c , this->sten_to_coefs ); - YAKL_SCOPE( weno_recon_lower , this->weno_recon_lower ); - YAKL_SCOPE( idl , this->idl ); - YAKL_SCOPE( sigma , this->sigma ); - YAKL_SCOPE( sim2d , this->sim2d ); - YAKL_SCOPE( derivMatrix , this->derivMatrix ); - YAKL_SCOPE( dx , this->dx ); - YAKL_SCOPE( tracer_pos , this->tracer_pos ); - YAKL_SCOPE( num_tracers , this->num_tracers ); - YAKL_SCOPE( gamma , this->gamma ); - YAKL_SCOPE( C0 , this->C0 ); - - real6d stateLimits ("stateLimits" ,num_state ,2,nz,ny,nx+1,nens); - real6d tracerLimits("tracerLimits",num_tracers,2,nz,ny,nx+1,nens); - - // Loop through all cells, reconstruct in x-direction, compute centered tendencies, store cell-edge state estimates - parallel_for( "Spatial.h X recon" , SimpleBounds<4>(nz,ny,nx,nens) , - YAKL_LAMBDA (int k, int j, int i, int iens) { - SArray r_DTs , ru_DTs; - - { // State - SArray rv_DTs , rw_DTs , rt_DTs, ruu_DTs , ruv_DTs , ruw_DTs , rut_DTs , rt_gamma_DTs; - - { // Recon - SArray stencil; - SArray gll; - - // Density - for (int ii=0; ii < ord; ii++) { stencil(ii) = state(idR,hs+k,hs+j,wrapx(i,ii,nx),iens); } - reconstruct_gll_values( stencil , gll , c2g , s2g , s2c , weno_recon_lower , idl , sigma , weno_scalars ); - // Add hydrostasis back on - for (int ii=0; ii < ngll; ii++) { r_DTs(0,ii) = gll(ii); } - - // u - for (int ii=0; ii < ord; ii++) { stencil(ii) = state(idU,hs+k,hs+j,wrapx(i,ii,nx),iens); } - reconstruct_gll_values( stencil , gll , c2g , s2g , s2c , weno_recon_lower , idl , sigma , weno_winds ); - for (int ii=0; ii < ngll; ii++) { ru_DTs(0,ii) = gll(ii); } - - // v - for (int ii=0; ii < ord; ii++) { stencil(ii) = state(idV,hs+k,hs+j,wrapx(i,ii,nx),iens); } - reconstruct_gll_values( stencil , gll , c2g , s2g , s2c , weno_recon_lower , idl , sigma , weno_winds ); - for (int ii=0; ii < ngll; ii++) { rv_DTs(0,ii) = gll(ii); } - - // w - for (int ii=0; ii < ord; ii++) { stencil(ii) = state(idW,hs+k,hs+j,wrapx(i,ii,nx),iens); } - reconstruct_gll_values( stencil , gll , c2g , s2g , s2c , weno_recon_lower , idl , sigma , weno_winds ); - for (int ii=0; ii < ngll; ii++) { rw_DTs(0,ii) = gll(ii); } - - // theta - for (int ii=0; ii < ord; ii++) { stencil(ii) = state(idT,hs+k,hs+j,wrapx(i,ii,nx),iens); } - reconstruct_gll_values( stencil , gll , c2g , s2g , s2c , weno_recon_lower , idl , sigma , weno_scalars ); - // Add hydrostasis back on - for (int ii=0; ii < ngll; ii++) { rt_DTs(0,ii) = gll(ii); } - } - - for (int ii=0; ii < ngll; ii++) { - real r = r_DTs (0,ii); - real u = ru_DTs(0,ii) / r; - real v = rv_DTs(0,ii) / r; - real w = rw_DTs(0,ii) / r; - real t = rt_DTs(0,ii) / r; - ruu_DTs (0,ii) = r*u*u; - ruv_DTs (0,ii) = r*u*v; - ruw_DTs (0,ii) = r*u*w; - rut_DTs (0,ii) = r*u*t; - rt_gamma_DTs(0,ii) = pow(r*t,gamma); - } - - if (nAder > 1) { - diffTransformEulerConsX( r_DTs , ru_DTs , rv_DTs , rw_DTs , rt_DTs , ruu_DTs , ruv_DTs , ruw_DTs , - rut_DTs , rt_gamma_DTs , derivMatrix , C0 , gamma , dx ); - } - - SArray r_tavg, ru_tavg; - if (timeAvg) { - compute_timeAvg( r_DTs , r_tavg , dt ); - compute_timeAvg( ru_DTs , ru_tavg , dt ); - compute_timeAvg( rv_DTs , dt ); - compute_timeAvg( rw_DTs , dt ); - compute_timeAvg( rt_DTs , dt ); - } else { - for (int ii=0; ii < ngll; ii++) { - r_tavg (ii) = r_DTs (0,ii); - ru_tavg(ii) = ru_DTs(0,ii); - } - } - - // Left interface - stateLimits(idR,1,k,j,i ,iens) = r_tavg (0 ); - stateLimits(idU,1,k,j,i ,iens) = ru_tavg (0 ); - stateLimits(idV,1,k,j,i ,iens) = rv_DTs(0,0 ); - stateLimits(idW,1,k,j,i ,iens) = rw_DTs(0,0 ); - stateLimits(idT,1,k,j,i ,iens) = rt_DTs(0,0 ); - // Right interface - stateLimits(idR,0,k,j,i+1,iens) = r_tavg (ngll-1); - stateLimits(idU,0,k,j,i+1,iens) = ru_tavg (ngll-1); - stateLimits(idV,0,k,j,i+1,iens) = rv_DTs(0,ngll-1); - stateLimits(idW,0,k,j,i+1,iens) = rw_DTs(0,ngll-1); - stateLimits(idT,0,k,j,i+1,iens) = rt_DTs(0,ngll-1); - } - - { // Tracers - for (int tr=0; tr < num_tracers; tr++) { - SArray rt_DTs, rut_DTs; - - { // Recon - SArray stencil; - SArray gll; - - for (int ii=0; ii < ord; ii++) { stencil(ii) = tracers(tr,hs+k,hs+j,wrapx(i,ii,nx),iens); } - reconstruct_gll_values( stencil , gll , c2g , s2g , s2c , weno_recon_lower , idl , sigma , weno_scalars ); - if (tracer_pos(tr)) { - for (int ii=0; ii < ngll; ii++) { gll(ii) = max( 0._fp , gll(ii) ); } - } - for (int ii=0; ii < ngll; ii++) { rt_DTs(0,ii) = gll(ii); } - } - - for (int ii=0; ii < ngll; ii++) { - rut_DTs(0,ii) = rt_DTs(0,ii) * ru_DTs(0,ii) / r_DTs(0,ii); - } - - if (nAder > 1) { - diffTransformTracer( r_DTs , ru_DTs , rt_DTs , rut_DTs , derivMatrix , dx ); - } - - if (timeAvg) { - compute_timeAvg( rt_DTs , dt ); - } - - if (tracer_pos(tr)) { - for (int ii=0; ii < ngll; ii++) { rt_DTs(0,ii) = max( 0._fp , rt_DTs(0,ii) ); } - } - - tracerLimits(tr,1,k,j,i ,iens) = rt_DTs(0,0 ); // Left interface - tracerLimits(tr,0,k,j,i+1,iens) = rt_DTs(0,ngll-1); // Right interface - } - } - }); - - real5d state_flux ("state_flux" ,num_state ,nz,ny,nx+1,nens); - real5d tracer_flux("tracer_flux",num_tracers,nz,ny,nx+1,nens); - - ////////////////////////////////////////////////////////// - // Compute the upwind fluxes - ////////////////////////////////////////////////////////// - parallel_for( "Spatial.h X Riemann" , SimpleBounds<4>(nz,ny,nx+1,nens) , - YAKL_LAMBDA (int k, int j, int i, int iens) { - if (i == 0 ) { - for (int l=0; l < num_state ; l++) { stateLimits (l,0,k,j,0 ,iens) = stateLimits (l,0,k,j,nx,iens); } - for (int l=0; l < num_tracers; l++) { tracerLimits(l,0,k,j,0 ,iens) = tracerLimits(l,0,k,j,nx,iens); } - } - if (i == nx) { - for (int l=0; l < num_state ; l++) { stateLimits (l,1,k,j,nx,iens) = stateLimits (l,1,k,j,0 ,iens); } - for (int l=0; l < num_tracers; l++) { tracerLimits(l,1,k,j,nx,iens) = tracerLimits(l,1,k,j,0 ,iens); } - } - // Get left and right state - real r_L = stateLimits(idR,0,k,j,i,iens) ; real r_R = stateLimits(idR,1,k,j,i,iens) ; - real u_L = stateLimits(idU,0,k,j,i,iens)/r_L; real u_R = stateLimits(idU,1,k,j,i,iens)/r_R; - real v_L = stateLimits(idV,0,k,j,i,iens)/r_L; real v_R = stateLimits(idV,1,k,j,i,iens)/r_R; - real w_L = stateLimits(idW,0,k,j,i,iens)/r_L; real w_R = stateLimits(idW,1,k,j,i,iens)/r_R; - real t_L = stateLimits(idT,0,k,j,i,iens)/r_L; real t_R = stateLimits(idT,1,k,j,i,iens)/r_R; - // Compute average state - real r = 0.5_fp * (r_L + r_R); - real u = 0.5_fp * (u_L + u_R); - real v = 0.5_fp * (v_L + v_R); - real w = 0.5_fp * (w_L + w_R); - real t = 0.5_fp * (t_L + t_R); - real p = C0 * pow(r*t,gamma); - real cs2 = gamma*p/r; - real cs = sqrt(cs2); - - // COMPUTE UPWIND STATE FLUXES - // Get left and right fluxes - real q1_L = stateLimits(idR,0,k,j,i,iens); real q1_R = stateLimits(idR,1,k,j,i,iens); - real q2_L = stateLimits(idU,0,k,j,i,iens); real q2_R = stateLimits(idU,1,k,j,i,iens); - real q3_L = stateLimits(idV,0,k,j,i,iens); real q3_R = stateLimits(idV,1,k,j,i,iens); - real q4_L = stateLimits(idW,0,k,j,i,iens); real q4_R = stateLimits(idW,1,k,j,i,iens); - real q5_L = stateLimits(idT,0,k,j,i,iens); real q5_R = stateLimits(idT,1,k,j,i,iens); - // Compute upwind characteristics - // Waves 1-3, velocity: u - real w1, w2, w3; - if (u > 0) { - w1 = q1_L - q5_L/t; - w2 = q3_L - v*q5_L/t; - w3 = q4_L - w*q5_L/t; - } else { - w1 = q1_R - q5_R/t; - w2 = q3_R - v*q5_R/t; - w3 = q4_R - w*q5_R/t; - } - // Wave 5, velocity: u-cs - real w5 = u*q1_R/(2*cs) - q2_R/(2*cs) + q5_R/(2*t); - // Wave 6, velocity: u+cs - real w6 = -u*q1_L/(2*cs) + q2_L/(2*cs) + q5_L/(2*t); - // Use right eigenmatrix to compute upwind flux - real q1 = w1 + w5 + w6; - real q2 = u*w1 + (u-cs)*w5 + (u+cs)*w6; - real q3 = w2 + v*w5 + v*w6; - real q4 = w3 + w*w5 + w*w6; - real q5 = t*w5 + t*w6; - - state_flux(idR,k,j,i,iens) = q2; - state_flux(idU,k,j,i,iens) = q2*q2/q1 + C0*pow(q5,gamma); - state_flux(idV,k,j,i,iens) = q2*q3/q1; - state_flux(idW,k,j,i,iens) = q2*q4/q1; - state_flux(idT,k,j,i,iens) = q2*q5/q1; - - // COMPUTE UPWIND TRACER FLUXES - // Handle it one tracer at a time - for (int tr=0; tr < num_tracers; tr++) { - if (u > 0) { - tracer_flux(tr,k,j,i,iens) = q2 * tracerLimits(tr,0,k,j,i,iens) / r_L; - } else { - tracer_flux(tr,k,j,i,iens) = q2 * tracerLimits(tr,1,k,j,i,iens) / r_R; - } - } - }); - - stateLimits = real6d(); - tracerLimits = real6d(); - - ////////////////////////////////////////////////////////// - // Limit the tracer fluxes for positivity - ////////////////////////////////////////////////////////// - real5d fct_mult("fct_mult",num_tracers,nz,ny,nx+1,nens); - parallel_for( "Spatial.h X FCT" , SimpleBounds<5>(num_tracers,nz,ny,nx+1,nens) , - YAKL_LAMBDA (int tr, int k, int j, int i, int iens) { - fct_mult(tr,k,j,i,iens) = 1.; - // Solid wall BCs mean u == 0 at boundaries, so we assume periodic if u != 0 - if (tracer_pos(tr)) { - // Compute and apply the flux reduction factor of the upwind cell - if (tracer_flux(tr,k,j,i,iens) > 0) { - // if u > 0, then it pulls mass out of the left cell - int ind_i = i-1; - // TODO: Relax the periodic assumption here - if (ind_i == -1) ind_i = nx-1; - real f1 = min( tracer_flux(tr,k,j,ind_i ,iens) , 0._fp ); - real f2 = max( tracer_flux(tr,k,j,ind_i+1,iens) , 0._fp ); - real fluxOut = dt*(f2-f1)/dx; - real mass = tracers(tr,hs+k,hs+j,hs+ind_i,iens); - if (fluxOut > 0) { - fct_mult(tr,k,j,i,iens) = min( 1._fp , mass / fluxOut ); - } - } else if (tracer_flux(tr,k,j,i,iens) < 0) { - // upwind is to the right of this interface - int ind_i = i; - // TODO: Relax the periodic assumption here - if (ind_i == nx) ind_i = 0; - real f1 = min( tracer_flux(tr,k,j,ind_i ,iens) , 0._fp ); - real f2 = max( tracer_flux(tr,k,j,ind_i+1,iens) , 0._fp ); - real fluxOut = dt*(f2-f1)/dx; - real mass = tracers(tr,hs+k,hs+j,hs+ind_i,iens); - if (fluxOut > 0) { - fct_mult(tr,k,j,i,iens) = min( 1._fp , mass / fluxOut ); - } - } - } - }); - - ////////////////////////////////////////////////////////// - // Compute the tendencies - ////////////////////////////////////////////////////////// - parallel_for( "Spatial.h X tendencies" , SimpleBounds<4>(nz,ny,nx,nens) , - YAKL_LAMBDA(int k, int j, int i, int iens) { - for (int l = 0; l < num_state; l++) { - if (sim2d && l == idV) { - stateTend(l,k,j,i,iens) = 0; - } else { - stateTend(l,k,j,i,iens) = - ( state_flux(l,k,j,i+1,iens) - state_flux(l,k,j,i,iens) ) / dx; - } - } - for (int l = 0; l < num_tracers; l++) { - // Compute tracer tendency - tracerTend(l,k,j,i,iens) = - ( tracer_flux(l,k,j,i+1,iens)*fct_mult(l,k,j,i+1,iens) - - tracer_flux(l,k,j,i ,iens)*fct_mult(l,k,j,i ,iens) ) / dx; - } - }); - } - - - - YAKL_INLINE static int wrapy(int j, int jj, int ny) { - int ret = j+jj; - if (ret < hs+0 ) ret += ny; - if (ret > hs+ny-1) ret -= ny; - return ret; - } - - - - void computeTendenciesY( real5d const &state , real5d const &stateTend , - real5d const &tracers , real5d const &tracerTend , - real &dt ) const { - using yakl::c::parallel_for; - using yakl::c::SimpleBounds; - - YAKL_SCOPE( ny , this->ny ); - YAKL_SCOPE( weno_scalars , this->weno_scalars ); - YAKL_SCOPE( weno_winds , this->weno_winds ); - YAKL_SCOPE( c2g , this->coefs_to_gll ); - YAKL_SCOPE( s2g , this->sten_to_gll ); - YAKL_SCOPE( s2c , this->sten_to_coefs ); - YAKL_SCOPE( weno_recon_lower , this->weno_recon_lower ); - YAKL_SCOPE( idl , this->idl ); - YAKL_SCOPE( sigma , this->sigma ); - - YAKL_SCOPE( derivMatrix , this->derivMatrix ); - YAKL_SCOPE( dy , this->dy ); - YAKL_SCOPE( tracer_pos , this->tracer_pos ); - YAKL_SCOPE( num_tracers , this->num_tracers ); - YAKL_SCOPE( gamma , this->gamma ); - YAKL_SCOPE( C0 , this->C0 ); - - real6d stateLimits ("stateLimits" ,num_state ,2,nz,ny+1,nx,nens); - real6d tracerLimits("tracerLimits",num_tracers,2,nz,ny+1,nx,nens); - - // Loop through all cells, reconstruct in y-direction, compute centered tendencies, store cell-edge state estimates - parallel_for( "Spatial.h Y recon" , SimpleBounds<4>(nz,ny,nx,nens) , - YAKL_LAMBDA (int k, int j, int i, int iens) { - SArray r_DTs , rv_DTs; - - { // State - SArray ru_DTs , rw_DTs , rt_DTs, rvu_DTs , rvv_DTs , rvw_DTs , rvt_DTs , rt_gamma_DTs; - - { // Recon - SArray stencil; - SArray gll; - - // Density - for (int jj=0; jj < ord; jj++) { stencil(jj) = state(idR,hs+k,wrapy(j,jj,ny),hs+i,iens); } - reconstruct_gll_values( stencil , gll , c2g , s2g , s2c , weno_recon_lower , idl , sigma , weno_scalars ); - // Add hydrostasis back on - for (int jj=0; jj < ngll; jj++) { r_DTs(0,jj) = gll(jj); } - - // u - for (int jj=0; jj < ord; jj++) { stencil(jj) = state(idU,hs+k,wrapy(j,jj,ny),hs+i,iens); } - reconstruct_gll_values( stencil , gll , c2g , s2g , s2c , weno_recon_lower , idl , sigma , weno_winds ); - for (int jj=0; jj < ngll; jj++) { ru_DTs(0,jj) = gll(jj); } - - // v - for (int jj=0; jj < ord; jj++) { stencil(jj) = state(idV,hs+k,wrapy(j,jj,ny),hs+i,iens); } - reconstruct_gll_values( stencil , gll , c2g , s2g , s2c , weno_recon_lower , idl , sigma , weno_winds ); - for (int jj=0; jj < ngll; jj++) { rv_DTs(0,jj) = gll(jj); } - - // w - for (int jj=0; jj < ord; jj++) { stencil(jj) = state(idW,hs+k,wrapy(j,jj,ny),hs+i,iens); } - reconstruct_gll_values( stencil , gll , c2g , s2g , s2c , weno_recon_lower , idl , sigma , weno_winds ); - for (int jj=0; jj < ngll; jj++) { rw_DTs(0,jj) = gll(jj); } - - // theta - for (int jj=0; jj < ord; jj++) { stencil(jj) = state(idT,hs+k,wrapy(j,jj,ny),hs+i,iens); } - reconstruct_gll_values( stencil , gll , c2g , s2g , s2c , weno_recon_lower , idl , sigma , weno_scalars ); - // Add hydrostasis back on - for (int jj=0; jj < ngll; jj++) { rt_DTs(0,jj) = gll(jj); } - } - - for (int jj=0; jj < ngll; jj++) { - real r = r_DTs (0,jj); - real u = ru_DTs(0,jj) / r; - real v = rv_DTs(0,jj) / r; - real w = rw_DTs(0,jj) / r; - real t = rt_DTs(0,jj) / r; - rvu_DTs (0,jj) = r*v*u; - rvv_DTs (0,jj) = r*v*v; - rvw_DTs (0,jj) = r*v*w; - rvt_DTs (0,jj) = r*v*t; - rt_gamma_DTs(0,jj) = pow(r*t,gamma); - } - - if (nAder > 1) { - diffTransformEulerConsY( r_DTs , ru_DTs , rv_DTs , rw_DTs , rt_DTs , rvu_DTs , rvv_DTs , rvw_DTs , - rvt_DTs , rt_gamma_DTs , derivMatrix , C0 , gamma , dy ); - } - - SArray r_tavg, rv_tavg; - if (timeAvg) { - compute_timeAvg( r_DTs , r_tavg , dt ); - compute_timeAvg( ru_DTs , dt ); - compute_timeAvg( rv_DTs , rv_tavg , dt ); - compute_timeAvg( rw_DTs , dt ); - compute_timeAvg( rt_DTs , dt ); - } else { - for (int jj=0; jj < ngll; jj++) { - r_tavg (jj) = r_DTs (0,jj); - rv_tavg(jj) = rv_DTs(0,jj); - } - } - - // Left interface - stateLimits(idR,1,k,j ,i,iens) = r_tavg (0 ); - stateLimits(idU,1,k,j ,i,iens) = ru_DTs(0,0 ); - stateLimits(idV,1,k,j ,i,iens) = rv_tavg (0 ); - stateLimits(idW,1,k,j ,i,iens) = rw_DTs(0,0 ); - stateLimits(idT,1,k,j ,i,iens) = rt_DTs(0,0 ); - // Right interface - stateLimits(idR,0,k,j+1,i,iens) = r_tavg (ngll-1); - stateLimits(idU,0,k,j+1,i,iens) = ru_DTs(0,ngll-1); - stateLimits(idV,0,k,j+1,i,iens) = rv_tavg (ngll-1); - stateLimits(idW,0,k,j+1,i,iens) = rw_DTs(0,ngll-1); - stateLimits(idT,0,k,j+1,i,iens) = rt_DTs(0,ngll-1); - } - - { // Tracers - for (int tr=0; tr < num_tracers; tr++) { - SArray rt_DTs, rvt_DTs; - - { // Recon - SArray stencil; - SArray gll; - - for (int jj=0; jj < ord; jj++) { stencil(jj) = tracers(tr,hs+k,wrapy(j,jj,ny),hs+i,iens); } - reconstruct_gll_values( stencil , gll , c2g , s2g , s2c , weno_recon_lower , idl , sigma , weno_scalars ); - if (tracer_pos(tr)) { - for (int jj=0; jj < ngll; jj++) { gll(jj) = max( 0._fp , gll(jj) ); } - } - for (int jj=0; jj < ngll; jj++) { rt_DTs(0,jj) = gll(jj); } - } - - for (int jj=0; jj < ngll; jj++) { - rvt_DTs(0,jj) = rt_DTs(0,jj) * rv_DTs(0,jj) / r_DTs(0,jj); - } - - if (nAder > 1) { - diffTransformTracer( r_DTs , rv_DTs , rt_DTs , rvt_DTs , derivMatrix , dy ); - } - - if (timeAvg) { - compute_timeAvg( rt_DTs , dt ); - } - - if (tracer_pos(tr)) { - for (int jj=0; jj < ngll; jj++) { rt_DTs(0,jj) = max( 0._fp , rt_DTs(0,jj) ); } - } - - tracerLimits(tr,1,k,j ,i,iens) = rt_DTs (0,0 ); // Left interface - tracerLimits(tr,0,k,j+1,i,iens) = rt_DTs (0,ngll-1); // Right interface - } - } - }); - - real5d state_flux ("state_flux" ,num_state ,nz,ny+1,nx,nens); - real5d tracer_flux("tracer_flux",num_tracers,nz,ny+1,nx,nens); - - ////////////////////////////////////////////////////////// - // Compute the upwind fluxes - ////////////////////////////////////////////////////////// - parallel_for( "Spatial.h Y Riemann" , SimpleBounds<4>(nz,ny+1,nx,nens) , - YAKL_LAMBDA (int k, int j, int i, int iens) { - if (j == 0 ) { - for (int l=0; l < num_state ; l++) { stateLimits (l,0,k,0 ,i,iens) = stateLimits (l,0,k,ny,i,iens); } - for (int l=0; l < num_tracers; l++) { tracerLimits(l,0,k,0 ,i,iens) = tracerLimits(l,0,k,ny,i,iens); } - } - if (j == ny) { - for (int l=0; l < num_state ; l++) { stateLimits (l,1,k,ny,i,iens) = stateLimits (l,1,k,0 ,i,iens); } - for (int l=0; l < num_tracers; l++) { tracerLimits(l,1,k,ny,i,iens) = tracerLimits(l,1,k,0 ,i,iens); } - } - // Get left and right state - real r_L = stateLimits(idR,0,k,j,i,iens) ; real r_R = stateLimits(idR,1,k,j,i,iens) ; - real u_L = stateLimits(idU,0,k,j,i,iens)/r_L; real u_R = stateLimits(idU,1,k,j,i,iens)/r_R; - real v_L = stateLimits(idV,0,k,j,i,iens)/r_L; real v_R = stateLimits(idV,1,k,j,i,iens)/r_R; - real w_L = stateLimits(idW,0,k,j,i,iens)/r_L; real w_R = stateLimits(idW,1,k,j,i,iens)/r_R; - real t_L = stateLimits(idT,0,k,j,i,iens)/r_L; real t_R = stateLimits(idT,1,k,j,i,iens)/r_R; - // Compute average state - real r = 0.5_fp * (r_L + r_R); - real u = 0.5_fp * (u_L + u_R); - real v = 0.5_fp * (v_L + v_R); - real w = 0.5_fp * (w_L + w_R); - real t = 0.5_fp * (t_L + t_R); - real p = C0 * pow(r*t,gamma); - real cs2 = gamma*p/r; - real cs = sqrt(cs2); - - // COMPUTE UPWIND STATE FLUXES - // Get left and right fluxes - real q1_L = stateLimits(idR,0,k,j,i,iens); real q1_R = stateLimits(idR,1,k,j,i,iens); - real q2_L = stateLimits(idU,0,k,j,i,iens); real q2_R = stateLimits(idU,1,k,j,i,iens); - real q3_L = stateLimits(idV,0,k,j,i,iens); real q3_R = stateLimits(idV,1,k,j,i,iens); - real q4_L = stateLimits(idW,0,k,j,i,iens); real q4_R = stateLimits(idW,1,k,j,i,iens); - real q5_L = stateLimits(idT,0,k,j,i,iens); real q5_R = stateLimits(idT,1,k,j,i,iens); - // Compute upwind characteristics - // Waves 1-3, velocity: v - real w1, w2, w3; - if (v > 0) { - w1 = q1_L - q5_L/t; - w2 = q2_L - u*q5_L/t; - w3 = q4_L - w*q5_L/t; - } else { - w1 = q1_R - q5_R/t; - w2 = q2_R - u*q5_R/t; - w3 = q4_R - w*q5_R/t; - } - // Wave 5, velocity: v-cs - real w5 = v*q1_R/(2*cs) - q3_R/(2*cs) + q5_R/(2*t); - // Wave 6, velocity: v+cs - real w6 = -v*q1_L/(2*cs) + q3_L/(2*cs) + q5_L/(2*t); - // Use right eigenmatrix to compute upwind flux - real q1 = w1 + w5 + w6; - real q2 = w2 + u*w5 + u*w6; - real q3 = v*w1 + (v-cs)*w5 + (v+cs)*w6; - real q4 = w3 + w*w5 + w*w6; - real q5 = t*w5 + t*w6; - - state_flux(idR,k,j,i,iens) = q3; - state_flux(idU,k,j,i,iens) = q3*q2/q1; - state_flux(idV,k,j,i,iens) = q3*q3/q1 + C0*pow(q5,gamma); - state_flux(idW,k,j,i,iens) = q3*q4/q1; - state_flux(idT,k,j,i,iens) = q3*q5/q1; - - // COMPUTE UPWIND TRACER FLUXES - // Handle it one tracer at a time - for (int tr=0; tr < num_tracers; tr++) { - if (v > 0) { - tracer_flux(tr,k,j,i,iens) = q3 * tracerLimits(tr,0,k,j,i,iens) / r_L; - } else { - tracer_flux(tr,k,j,i,iens) = q3 * tracerLimits(tr,1,k,j,i,iens) / r_R; - } - } - }); - - stateLimits = real6d(); - tracerLimits = real6d(); - - ////////////////////////////////////////////////////////// - // Limit the tracer fluxes for positivity - ////////////////////////////////////////////////////////// - real5d fct_mult("fct_mult",num_tracers,nz,ny+1,nx,nens); - parallel_for( "Spatial.h Y FCT" , SimpleBounds<5>(num_tracers,nz,ny+1,nx,nens) , - YAKL_LAMBDA (int tr, int k, int j, int i, int iens) { - fct_mult(tr,k,j,i,iens) = 1.; - // Solid wall BCs mean u == 0 at boundaries, so we assume periodic if u != 0 - if (tracer_pos(tr)) { - // Compute and apply the flux reduction factor of the upwind cell - if (tracer_flux(tr,k,j,i,iens) > 0) { - // upwind is to the left of this interface - int ind_j = j-1; - if (ind_j == -1) ind_j = ny-1; - real f1 = min( tracer_flux(tr,k,ind_j ,i,iens) , 0._fp ); - real f2 = max( tracer_flux(tr,k,ind_j+1,i,iens) , 0._fp ); - real fluxOut = dt*(f2-f1)/dy; - real mass = tracers(tr,hs+k,hs+ind_j,hs+i,iens); - if (fluxOut > 0) { - fct_mult(tr,k,j,i,iens) = min( 1._fp , mass / fluxOut ); - } - } else if (tracer_flux(tr,k,j,i,iens) < 0) { - // upwind is to the right of this interface - int ind_j = j; - if (ind_j == ny) ind_j = 0; - real f1 = min( tracer_flux(tr,k,ind_j ,i,iens) , 0._fp ); - real f2 = max( tracer_flux(tr,k,ind_j+1,i,iens) , 0._fp ); - real fluxOut = dt*(f2-f1)/dy; - real mass = tracers(tr,hs+k,hs+ind_j,hs+i,iens); - if (fluxOut > 0) { - fct_mult(tr,k,j,i,iens) = min( 1._fp , mass / fluxOut ); - } - } - } - }); - - ////////////////////////////////////////////////////////// - // Compute the tendencies - ////////////////////////////////////////////////////////// - parallel_for( "Spatial.h Y tendendies" , SimpleBounds<4>(nz,ny,nx,nens) , - YAKL_LAMBDA(int k, int j, int i, int iens) { - for (int l=0; l < num_state; l++) { - stateTend(l,k,j,i,iens) = - ( state_flux(l,k,j+1,i,iens) - state_flux(l,k,j,i,iens) ) / dy; - } - for (int l=0; l < num_tracers; l++) { - // Compute the tracer tendency - tracerTend(l,k,j,i,iens) = - ( tracer_flux(l,k,j+1,i,iens)*fct_mult(l,k,j+1,i,iens) - - tracer_flux(l,k,j ,i,iens)*fct_mult(l,k,j ,i,iens) ) / dy; - } - }); - } - - - - YAKL_INLINE static int wrapz(int k, int kk, int nz) { - int ret = k+kk; - if (ret < hs+0 ) ret = hs+0; - if (ret > hs+nz-1) ret = hs+nz-1; - return ret; - } - - - - void computeTendenciesZ( real5d const &state , real5d const &stateTend , - real5d const &tracers , real5d const &tracerTend , - real &dt ) const { - using yakl::c::parallel_for; - using yakl::c::SimpleBounds; - - YAKL_SCOPE( nz , this->nz ); - YAKL_SCOPE( weno_scalars , this->weno_scalars ); - YAKL_SCOPE( weno_winds , this->weno_winds ); - YAKL_SCOPE( c2g , this->coefs_to_gll ); - YAKL_SCOPE( idl , this->idl ); - YAKL_SCOPE( sigma , this->sigma ); - YAKL_SCOPE( hyDensSten , this->hyDensSten ); - YAKL_SCOPE( hyDensThetaSten , this->hyDensThetaSten ); - YAKL_SCOPE( hyDensGLL , this->hyDensGLL ); - YAKL_SCOPE( hyDensThetaGLL , this->hyDensThetaGLL ); - YAKL_SCOPE( hyPressureGLL , this->hyPressureGLL ); - YAKL_SCOPE( sim2d , this->sim2d ); - YAKL_SCOPE( derivMatrix , this->derivMatrix ); - YAKL_SCOPE( dz , this->dz ); - YAKL_SCOPE( tracer_pos , this->tracer_pos ); - YAKL_SCOPE( num_tracers , this->num_tracers ); - YAKL_SCOPE( gllWts_ngll , this->gllWts_ngll ); - YAKL_SCOPE( gamma , this->gamma ); - YAKL_SCOPE( C0 , this->C0 ); - YAKL_SCOPE( vert_sten_to_gll , this->vert_sten_to_gll ); - YAKL_SCOPE( vert_sten_to_coefs , this->vert_sten_to_coefs ); - YAKL_SCOPE( vert_weno_recon_lower , this->vert_weno_recon_lower ); - YAKL_SCOPE( grav , this->grav ); - - // Pre-process the tracers by dividing by density inside the domain - // After this, we can reconstruct tracers only (not rho * tracer) - parallel_for( "Spatial.h Z tracer div dens" , SimpleBounds<5>(num_tracers,nz,ny,nx,nens) , - YAKL_LAMBDA (int tr, int k, int j, int i, int iens) { - tracers(tr,hs+k,hs+j,hs+i,iens) /= state(idR,hs+k,hs+j,hs+i,iens); - }); - - real6d stateLimits ("stateLimits" ,num_state ,2,nz+1,ny,nx,nens); - real6d tracerLimits("tracerLimits",num_tracers,2,nz+1,ny,nx,nens); - - // Loop through all cells, reconstruct in x-direction, compute centered tendencies, store cell-edge state estimates - parallel_for( "Spatial.h Z recon" , SimpleBounds<4>(nz,ny,nx,nens) , - YAKL_LAMBDA (int k, int j, int i, int iens) { - SArray s2g_loc; - SArray s2c_loc; - SArray weno_recon_lower_loc; - for (int jj=0; jj < ord; jj++) { - for (int ii=0; ii < ngll; ii++) { - s2g_loc(jj,ii) = vert_sten_to_gll(k,jj,ii,iens); - } - } - for (int jj=0; jj < ord; jj++) { - for (int ii=0; ii < ord; ii++) { - s2c_loc(jj,ii) = vert_sten_to_coefs(k,jj,ii,iens); - } - } - for (int kk=0; kk < hs+1; kk++) { - for (int jj=0; jj < hs+1; jj++) { - for (int ii=0; ii < hs+1; ii++) { - weno_recon_lower_loc(kk,jj,ii) = vert_weno_recon_lower(k,kk,jj,ii,iens); - } - } - } - - SArray r_DTs , rw_DTs; - - { // State - SArray ru_DTs , rv_DTs , rt_DTs, rwu_DTs , rwv_DTs , rww_DTs , rwt_DTs , rt_gamma_DTs; - { // Recon - SArray stencil; - SArray gll; - - // Density - for (int kk=0; kk < ord; kk++) { - if (k+kk < hs || k+kk > hs+nz-1) { - stencil(kk) = state(idR,hs+k,hs+j,hs+i,iens) - hyDensSten(k,hs,iens); - } else { - stencil(kk) = state(idR,k+kk,hs+j,hs+i,iens) - hyDensSten(k,kk,iens); - } - } - reconstruct_gll_values( stencil , gll , c2g , s2g_loc , s2c_loc , weno_recon_lower_loc , - idl , sigma , weno_scalars ); - // Add hydrostasis back on - for (int kk=0; kk < ngll; kk++) { r_DTs(0,kk) = gll(kk) + hyDensGLL(k,kk,iens); } - - // u values and derivatives - for (int kk=0; kk < ord; kk++) { stencil(kk) = state(idU,wrapz(k,kk,nz),hs+j,hs+i,iens); } - reconstruct_gll_values( stencil , gll , c2g , s2g_loc , s2c_loc , weno_recon_lower_loc , - idl , sigma , weno_winds ); - for (int kk=0; kk < ngll; kk++) { ru_DTs(0,kk) = gll(kk); } - - // v - for (int kk=0; kk < ord; kk++) { stencil(kk) = state(idV,wrapz(k,kk,nz),hs+j,hs+i,iens); } - reconstruct_gll_values( stencil , gll , c2g , s2g_loc , s2c_loc , weno_recon_lower_loc , - idl , sigma , weno_winds ); - for (int kk=0; kk < ngll; kk++) { rv_DTs(0,kk) = gll(kk); } - - // w - for (int kk=0; kk < ord; kk++) { - stencil(kk) = state(idW,wrapz(k,kk,nz),hs+j,hs+i,iens); - if (k+kk > hs+nz-1 || k+kk < hs) stencil(kk) = 0; - } - reconstruct_gll_values( stencil , gll , c2g , s2g_loc , s2c_loc , weno_recon_lower_loc , - idl , sigma , weno_winds ); - if (k == nz-1) gll(ngll-1) = 0; - if (k == 0 ) gll(0 ) = 0; - for (int kk=0; kk < ngll; kk++) { rw_DTs(0,kk) = gll(kk); } - - // rho*theta - for (int kk=0; kk < ord; kk++) { - if (k+kk < hs || k+kk > hs+nz-1) { - stencil(kk) = state(idT,hs+k,hs+j,hs+i,iens) - hyDensThetaSten(k,hs,iens); - } else { - stencil(kk) = state(idT,k+kk,hs+j,hs+i,iens) - hyDensThetaSten(k,kk,iens); - } - } - reconstruct_gll_values( stencil , gll , c2g , s2g_loc , s2c_loc , weno_recon_lower_loc , - idl , sigma , weno_scalars ); - // Add hydrostasis back on - for (int kk=0; kk < ngll; kk++) { rt_DTs(0,kk) = gll(kk) + hyDensThetaGLL(k,kk,iens); } - } - - for (int kk=0; kk < ngll; kk++) { - real r = r_DTs (0,kk); - real u = ru_DTs(0,kk) / r; - real v = rv_DTs(0,kk) / r; - real w = rw_DTs(0,kk) / r; - real t = rt_DTs(0,kk) / r; - rwu_DTs (0,kk) = r*w*u; - rwv_DTs (0,kk) = r*w*v; - rww_DTs (0,kk) = r*w*w; - rwt_DTs (0,kk) = r*w*t; - rt_gamma_DTs(0,kk) = pow(r*t,gamma); - } - - if (nAder > 1) { - diffTransformEulerConsZ( r_DTs , ru_DTs , rv_DTs , rw_DTs , rt_DTs , rwu_DTs , rwv_DTs , rww_DTs , - rwt_DTs , rt_gamma_DTs , derivMatrix , hyDensGLL , hyPressureGLL , C0 , gamma , - grav , k , dz(k,iens) , nz , iens ); - } - - SArray r_tavg, rw_tavg; - if (timeAvg) { - compute_timeAvg( r_DTs , r_tavg , dt ); - compute_timeAvg( ru_DTs , dt ); - compute_timeAvg( rv_DTs , dt ); - compute_timeAvg( rw_DTs , rw_tavg , dt ); - compute_timeAvg( rt_DTs , dt ); - } else { - for (int ii=0; ii < ngll; ii++) { - r_tavg (ii) = r_DTs (0,ii); - rw_tavg(ii) = rw_DTs(0,ii); - } - } - - // Left interface - stateLimits(idR,1,k ,j,i,iens) = r_tavg (0 ); - stateLimits(idU,1,k ,j,i,iens) = ru_DTs(0,0 ); - stateLimits(idV,1,k ,j,i,iens) = rv_DTs(0,0 ); - stateLimits(idW,1,k ,j,i,iens) = rw_tavg (0 ); - stateLimits(idT,1,k ,j,i,iens) = rt_DTs(0,0 ); - // Right interface - stateLimits(idR,0,k+1,j,i,iens) = r_tavg (ngll-1); - stateLimits(idU,0,k+1,j,i,iens) = ru_DTs(0,ngll-1); - stateLimits(idV,0,k+1,j,i,iens) = rv_DTs(0,ngll-1); - stateLimits(idW,0,k+1,j,i,iens) = rw_tavg (ngll-1); - stateLimits(idT,0,k+1,j,i,iens) = rt_DTs(0,ngll-1); - - real ravg = 0; - for (int kk=0; kk < ngll; kk++) { - ravg += r_tavg(kk) * gllWts_ngll(kk); - } - stateTend(idR,k,j,i,iens) = 0; - stateTend(idU,k,j,i,iens) = 0; - stateTend(idV,k,j,i,iens) = 0; - stateTend(idW,k,j,i,iens) = -grav * ravg; - stateTend(idT,k,j,i,iens) = 0; - } - - { // Tracers - for (int tr=0; tr < num_tracers; tr++) { - SArray rt_DTs; // Density * tracer - SArray rwt_DTs; // Density * wwind * tracer - { // Recon - SArray stencil; - SArray gll; - - for (int kk=0; kk < ord; kk++) { stencil(kk) = tracers(tr,wrapz(k,kk,nz),hs+j,hs+i,iens); } - reconstruct_gll_values( stencil , gll , c2g , s2g_loc , s2c_loc , weno_recon_lower_loc , - idl , sigma , weno_scalars ); - for (int kk=0; kk < ngll; kk++) { gll(kk) *= r_DTs(0,kk); } - if (tracer_pos(tr)) { - for (int kk=0; kk < ngll; kk++) { gll(kk) = max( 0._fp , gll(kk) ); } - } - for (int kk=0; kk < ngll; kk++) { rt_DTs(0,kk) = gll(kk); } - } - - for (int kk=0; kk < ngll; kk++) { - rwt_DTs(0,kk) = rt_DTs(0,kk) * rw_DTs(0,kk) / r_DTs(0,kk); - } - - if (nAder > 1) { - diffTransformTracer( r_DTs , rw_DTs , rt_DTs , rwt_DTs , derivMatrix , dz(k,iens) ); - } - - if (timeAvg) { - compute_timeAvg( rt_DTs , dt ); - } - - if (tracer_pos(tr)) { - for (int kk=0; kk < ngll; kk++) { rt_DTs(0,kk) = max( 0._fp , rt_DTs(0,kk) ); } - } - - if (k == nz-1) rwt_DTs(0,ngll-1) = 0; - if (k == 0 ) rwt_DTs(0,0 ) = 0; - - tracerLimits(tr,1,k ,j,i,iens) = rt_DTs (0,0 ); // Left interface - tracerLimits(tr,0,k+1,j,i,iens) = rt_DTs (0,ngll-1); // Right interface - } - } - }); - - real5d state_flux ("state_flux" ,num_state ,nz+1,ny,nx,nens); - real5d tracer_flux("tracer_flux",num_tracers,nz+1,ny,nx,nens); - - ////////////////////////////////////////////////////////// - // Compute the upwind fluxes - ////////////////////////////////////////////////////////// - parallel_for( "Spatial.h Z Riemann" , SimpleBounds<4>(nz+1,ny,nx,nens) , - YAKL_LAMBDA (int k, int j, int i, int iens) { - if (k == 0) { - for (int l = 0; l < num_state ; l++) { stateLimits (l,0,0 ,j,i,iens) = stateLimits (l,1,0 ,j,i,iens); } - for (int l = 0; l < num_tracers; l++) { tracerLimits(l,0,0 ,j,i,iens) = tracerLimits(l,1,0 ,j,i,iens); } - } - if (k == nz) { - for (int l = 0; l < num_state ; l++) { stateLimits (l,1,nz,j,i,iens) = stateLimits (l,0,nz,j,i,iens); } - for (int l = 0; l < num_tracers; l++) { tracerLimits(l,1,nz,j,i,iens) = tracerLimits(l,0,nz,j,i,iens); } - } - // Get left and right state - real r_L = stateLimits(idR,0,k,j,i,iens) ; real r_R = stateLimits(idR,1,k,j,i,iens) ; - real u_L = stateLimits(idU,0,k,j,i,iens)/r_L; real u_R = stateLimits(idU,1,k,j,i,iens)/r_R; - real v_L = stateLimits(idV,0,k,j,i,iens)/r_L; real v_R = stateLimits(idV,1,k,j,i,iens)/r_R; - real w_L = stateLimits(idW,0,k,j,i,iens)/r_L; real w_R = stateLimits(idW,1,k,j,i,iens)/r_R; - real t_L = stateLimits(idT,0,k,j,i,iens)/r_L; real t_R = stateLimits(idT,1,k,j,i,iens)/r_R; - // Compute average state - real r = 0.5_fp * (r_L + r_R); - real u = 0.5_fp * (u_L + u_R); - real v = 0.5_fp * (v_L + v_R); - real w = 0.5_fp * (w_L + w_R); - real t = 0.5_fp * (t_L + t_R); - real p = C0 * pow(r*t,gamma); - real cs2 = gamma*p/r; - real cs = sqrt(cs2); - // Get left and right fluxes - real q1_L = stateLimits(idR,0,k,j,i,iens); real q1_R = stateLimits(idR,1,k,j,i,iens); - real q2_L = stateLimits(idU,0,k,j,i,iens); real q2_R = stateLimits(idU,1,k,j,i,iens); - real q3_L = stateLimits(idV,0,k,j,i,iens); real q3_R = stateLimits(idV,1,k,j,i,iens); - real q4_L = stateLimits(idW,0,k,j,i,iens); real q4_R = stateLimits(idW,1,k,j,i,iens); - real q5_L = stateLimits(idT,0,k,j,i,iens); real q5_R = stateLimits(idT,1,k,j,i,iens); - // Compute upwind characteristics - // Waves 1-3, velocity: w - real w1, w2, w3; - if (w > 0) { - w1 = q1_L - q5_L/t; - w2 = q2_L - u*q5_L/t; - w3 = q3_L - v*q5_L/t; - } else { - w1 = q1_R - q5_R/t; - w2 = q2_R - u*q5_R/t; - w3 = q3_R - v*q5_R/t; - } - // Wave 5, velocity: w-cs - real w5 = w*q1_R/(2*cs) - q4_R/(2*cs) + q5_R/(2*t); - // Wave 6, velocity: w+cs - real w6 = -w*q1_L/(2*cs) + q4_L/(2*cs) + q5_L/(2*t); - // Use right eigenmatrix to compute upwind flux - real q1 = w1 + w5 + w6; - real q2 = w2 + u*w5 + u*w6; - real q3 = w3 + v*w5 + v*w6; - real q4 = w*w1 + (w-cs)*w5 + (w+cs)*w6; - real q5 = t*w5 + t*w6; - - state_flux(idR,k,j,i,iens) = q4; - state_flux(idU,k,j,i,iens) = q4*q2/q1; - state_flux(idV,k,j,i,iens) = q4*q3/q1; - state_flux(idW,k,j,i,iens) = q4*q4/q1 + C0*pow(q5,gamma); - state_flux(idT,k,j,i,iens) = q4*q5/q1; - - // COMPUTE UPWIND TRACER FLUXES - // Handle it one tracer at a time - for (int tr=0; tr < num_tracers; tr++) { - if (w > 0) { - tracer_flux(tr,k,j,i,iens) = q4 * tracerLimits(tr,0,k,j,i,iens) / r_L; - } else { - tracer_flux(tr,k,j,i,iens) = q4 * tracerLimits(tr,1,k,j,i,iens) / r_R; - } - } - }); - - stateLimits = real6d(); - tracerLimits = real6d(); - - ////////////////////////////////////////////////////////// - // Limit the tracer fluxes for positivity - ////////////////////////////////////////////////////////// - real5d fct_mult("fct_mult",num_tracers,nz+1,ny,nx,nens); - parallel_for( "Spatial.h Z FCT" , SimpleBounds<5>(num_tracers,nz+1,ny,nx,nens) , - YAKL_LAMBDA (int tr, int k, int j, int i, int iens) { - fct_mult(tr,k,j,i,iens) = 1.; - if (k == 0 || k == nz) tracer_flux(tr,k,j,i,iens) = 0; - // Solid wall BCs mean w == 0 at boundaries - if (tracer_pos(tr)) { - // Compute and apply the flux reduction factor of the upwind cell - if (tracer_flux(tr,k,j,i,iens) > 0) { - int ind_k = k-1; - // upwind is to the left of this interface - real f1 = min( tracer_flux(tr,ind_k ,j,i,iens) , 0._fp ); - real f2 = max( tracer_flux(tr,ind_k+1,j,i,iens) , 0._fp ); - real fluxOut = dt*(f2-f1)/dz(ind_k,iens); - real dens = state(idR,hs+ind_k,hs+j,hs+i,iens); - real mass = tracers(tr,hs+ind_k,hs+j,hs+i,iens) * dens; - if (fluxOut > 0) { - fct_mult(tr,k,j,i,iens) = min( 1._fp , mass / fluxOut ); - } - } else if (tracer_flux(tr,k,j,i,iens) < 0) { - int ind_k = k; - // upwind is to the right of this interface - real f1 = min( tracer_flux(tr,ind_k ,j,i,iens) , 0._fp ); - real f2 = max( tracer_flux(tr,ind_k+1,j,i,iens) , 0._fp ); - real fluxOut = dt*(f2-f1)/dz(ind_k,iens); - real dens = state(idR,hs+ind_k,hs+j,hs+i,iens); - real mass = tracers(tr,hs+ind_k,hs+j,hs+i,iens) * dens; - if (fluxOut > 0) { - fct_mult(tr,k,j,i,iens) = min( 1._fp , mass / fluxOut ); - } - } - } - }); - - ////////////////////////////////////////////////////////// - // Compute the tendencies - ////////////////////////////////////////////////////////// - parallel_for( "Spatial.h Z tendencies" , SimpleBounds<4>(nz,ny,nx,nens) , - YAKL_LAMBDA(int k, int j, int i, int iens) { - for (int l=0; l < num_state; l++) { - if (sim2d && l == idV) { - stateTend(l,k,j,i,iens) = 0; - } else { - stateTend(l,k,j,i,iens) += - ( state_flux(l,k+1,j,i,iens) - state_flux(l,k,j,i,iens) ) / dz(k,iens); - } - } - for (int l=0; l < num_tracers; l++) { - // Compute tracer tendency - tracerTend(l,k,j,i,iens) = - ( tracer_flux(l,k+1,j,i,iens)*fct_mult(l,k+1,j,i,iens) - - tracer_flux(l,k ,j,i,iens)*fct_mult(l,k ,j,i,iens) ) / dz(k,iens); - // Multiply density back onto the tracers - tracers(l,hs+k,hs+j,hs+i,iens) *= state(idR,hs+k,hs+j,hs+i,iens); - } - }); - } - - - - const char * getName() { return ""; } - - - - void finalize(real4d const &state , real4d const &tracers) {} - - - - // ord stencil values to ngll GLL values; store in DTs - YAKL_INLINE static void reconstruct_gll_values( SArray const stencil , - SArray &gll , - SArray const &coefs_to_gll , - SArray const &sten_to_gll , - SArray const &sten_to_coefs , - SArray const &weno_recon_lower , - SArray const &idl , - real sigma, bool doweno ) { - if (doweno) { - - // Reconstruct values - SArray wenoCoefs; - weno::compute_weno_coefs( weno_recon_lower , sten_to_coefs , stencil , wenoCoefs , idl , sigma ); - // Transform ord weno coefficients into ngll GLL points - for (int ii=0; ii &r , - SArray &ru , - SArray &rv , - SArray &rw , - SArray &rt , - SArray &ruu , - SArray &ruv , - SArray &ruw , - SArray &rut , - SArray &rt_gamma , - SArray const &deriv , - real C0, real gamma, real dx ) { - // zero out the non-linear DTs - for (int kt=1; kt < nAder; kt++) { - for (int ii=0; ii < ngll; ii++) { - ruu (kt,ii) = 0; - ruv (kt,ii) = 0; - ruw (kt,ii) = 0; - rut (kt,ii) = 0; - rt_gamma(kt,ii) = 0; - } - } - - // Loop over the time derivatives - for (int kt=0; kt &r , - SArray &ru , - SArray &rv , - SArray &rw , - SArray &rt , - SArray &rvu , - SArray &rvv , - SArray &rvw , - SArray &rvt , - SArray &rt_gamma , - SArray const &deriv , - real C0, real gamma, real dy ) { - // zero out the non-linear DTs - for (int kt=1; kt < nAder; kt++) { - for (int ii=0; ii < ngll; ii++) { - rvu (kt,ii) = 0; - rvv (kt,ii) = 0; - rvw (kt,ii) = 0; - rvt (kt,ii) = 0; - rt_gamma(kt,ii) = 0; - } - } - - // Loop over the time derivatives - for (int kt=0; kt &r , - SArray &ru , - SArray &rv , - SArray &rw , - SArray &rt , - SArray &rwu , - SArray &rwv , - SArray &rww , - SArray &rwt , - SArray &rt_gamma , - SArray const &deriv , - real3d const &hyDensGLL , - real3d const &hyPressureGLL , - real C0, real gamma , real grav , - int k , real dz , int nz , int iens ) { - // zero out the non-linear DTs - for (int kt=1; kt < nAder; kt++) { - for (int ii=0; ii < ngll; ii++) { - rwu (kt,ii) = 0; - rwv (kt,ii) = 0; - rww (kt,ii) = 0; - rwt (kt,ii) = 0; - rt_gamma(kt,ii) = 0; - } - } - - // Loop over the time derivatives - for (int kt=0; kt const &r , - SArray const &ru , - SArray &rt , - SArray &rut , - SArray const &deriv , - real dx ) { - // zero out the non-linear DT - for (int kt=1; kt < nAder; kt++) { - for (int ii=0; ii < ngll; ii++) { - rut(kt,ii) = 0; - } - } - // Loop over the time derivatives - for (int kt=0; kt &dts , real dt ) { - real dtmult = dt; - for (int kt=1; kt &dts , real dt ) { - real dtmult = dt; - for (int kt=1; kt const &dts , SArray &tavg , - real dt ) { - for (int ii=0; ii class Temporal_operator { -public: - static_assert(nTimeDerivs <= ngll , "ERROR: nTimeDerivs must be <= ngll."); - - real5d stateTend; - real5d tracerTend; - - Spatial space_op; - - real etime; - - void init(pam::PamCoupler &coupler) { - - space_op.init(coupler); - - stateTend = space_op.createStateTendArr (); - tracerTend = space_op.createTracerTendArr(); - etime = 0; - } - - - int add_tracer(std::string name , std::string desc , bool pos_def , bool adds_mass) { - return space_op.add_tracer(name , desc , pos_def , adds_mass); - } - - - void init_idealized_state_and_tracers( pam::PamCoupler &coupler ) { - space_op.init_idealized_state_and_tracers( coupler ); - } - - - real compute_time_step(pam::PamCoupler const &coupler, real cfl_in = -1) { - return space_op.compute_time_step(coupler, cfl_in); - } - - - std::vector compute_mass( pam::PamCoupler const &coupler , realConst5d state , realConst5d tracers ) const { - using yakl::c::parallel_for; - using yakl::c::SimpleBounds; - - int nz = coupler.get_nz(); - int ny = coupler.get_ny(); - int nx = coupler.get_nx(); - int nens = coupler.get_nens(); - - int idR = Spatial::idR; - int hs = Spatial::hs; - int num_tracers = space_op.num_tracers; - YAKL_SCOPE( dz , space_op.dz ); - - std::vector mass(num_tracers+1); - real4d tmp("tmp",nz,ny,nx,nens); - - parallel_for( "Temporal_ader.h state mass" , SimpleBounds<4>(nz,ny,nx,nens) , YAKL_LAMBDA (int k, int j, int i, int iens) { - tmp(k,j,i,iens) = state(idR,hs+k,hs+j,hs+i,iens) * dz(k,iens); - }); - mass[0] = yakl::intrinsics::sum(tmp); - - for (int l=0; l < num_tracers; l++) { - parallel_for( "Temporal_ader.h tracer mass" , SimpleBounds<4>(nz,ny,nx,nens) , YAKL_LAMBDA (int k, int j, int i, int iens) { - tmp(k,j,i,iens) = tracers(l,hs+k,hs+j,hs+i,iens) * dz(k,iens); - }); - mass[l+1] = yakl::intrinsics::sum(tmp); - } - return mass; - } - - - void timeStep( pam::PamCoupler &coupler ) { - using yakl::c::parallel_for; - using yakl::c::SimpleBounds; - - real dtphys = coupler.get_option("crm_dt"); - - YAKL_SCOPE( stateTend , this->stateTend ); - YAKL_SCOPE( tracerTend , this->tracerTend ); - - real dt = compute_time_step( coupler ); - - real5d state = space_op.createStateArr(); - real5d tracers = space_op.createTracerArr(); - - #ifdef PAM_DEBUG - memset(state , 0._fp); - memset(tracers , 0._fp); - #endif - - space_op.convert_coupler_state_to_dynamics( coupler , state , tracers ); - - int idR = space_op.idR; - int idU = space_op.idU; - int idV = space_op.idV; - int idW = space_op.idW; - int idT = space_op.idT; - int nx = space_op.nx; - int ny = space_op.ny; - int nz = space_op.nz; - int nens = space_op.nens; - int num_state = space_op.num_state; - int num_tracers = space_op.num_tracers; - int hs = space_op.hs; - - int n_iter = ceil( dtphys / dt ); - dt = dtphys / n_iter; - - for (int iter = 0; iter < n_iter; iter++) { - - #ifdef PAM_DEBUG - validate_array_positive(tracers); - validate_array_inf_nan(state); - validate_array_inf_nan(tracers); - std::vector mass_init = compute_mass( coupler , state , tracers ); - #endif - - ScalarLiveOut neg_too_large(false); - - // Loop over different items in the spatial splitting - for (int spl = 0 ; spl < space_op.numSplit() ; spl++) { - real dtloc = dt; - - // Compute the tendencies for state and tracers - space_op.computeTendencies( state , stateTend , tracers , tracerTend , dtloc , spl ); - - parallel_for( "Temporal_ader.h apply tendencies" , SimpleBounds<4>(nz,ny,nx,nens) , - YAKL_LAMBDA (int k, int j, int i, int iens) { - for (int l=0; l < num_state; l++) { - state(l,hs+k,hs+j,hs+i,iens) += dtloc * stateTend(l,k,j,i,iens); - } - for (int l=0; l < num_tracers; l++) { - tracers(l,hs+k,hs+j,hs+i,iens) += dtloc * tracerTend(l,k,j,i,iens); - #ifdef PAM_DEBUG - if (tracers(l,hs+k,hs+j,hs+i,iens) < -1.e-10) { - neg_too_large = true; - } - #endif - tracers(l,hs+k,hs+j,hs+i,iens) = max( 0._fp , tracers(l,hs+k,hs+j,hs+i,iens) ); - } - }); - } - - #ifdef PAM_DEBUG - if (neg_too_large.hostRead()) { - std::cerr << "WARNING: Correcting a non-machine-precision negative tracer value" << std::endl; - // endrun(); - } - std::vector mass_final = compute_mass( coupler , state , tracers ); - for (int l=0; l < mass_final.size(); l++) { - real mass_diff; - if (mass_init[l] > 0) { - mass_diff = abs(mass_final[l] - mass_init[l]) / abs(mass_init[l]); - } else { - mass_diff = mass_final[l]; - } - real tol = 1.e-12; - if (std::is_same::value) {tol = 1.e-5;} - if (mass_diff > tol) { - std::cout << "Dycore mass change is too large. Abs Diff: " << abs(mass_final[l] - mass_init[l]) - << "; Rel Diff: " << mass_diff - << "; Initial Mass: " << mass_init[l] << std::endl; - // endrun("ERROR: mass not conserved by dycore"); - } - } - validate_array_positive(tracers); - validate_array_inf_nan(state); - validate_array_inf_nan(tracers); - #endif - - space_op.switch_directions(); - - } - - space_op.convert_dynamics_to_coupler_state( coupler , state , tracers ); - - etime += dtphys; - } - - - void finalize(pam::PamCoupler &coupler) { } - - - const char * dycore_name() const { return "AWFL"; } - -}; diff --git a/dynamics/awfl/Temporal_ssprk3_needs_update.h b/dynamics/awfl/Temporal_ssprk3_needs_update.h deleted file mode 100644 index 87367412..00000000 --- a/dynamics/awfl/Temporal_ssprk3_needs_update.h +++ /dev/null @@ -1,243 +0,0 @@ - -#pragma once - -#include "awfl_const.h" -#include "DataManager.h" - -int constexpr nTimeDerivs = 1; -bool constexpr timeAvg = false; -int constexpr nAder = 1; - -template class Temporal_operator { -public: - static_assert(nTimeDerivs <= ngll , "ERROR: nTimeDerivs must be <= ngll."); - - real5d stateTmp; - real5d tracersTmp; - - real5d stateTend; - real5d tracerTend; - - real5d stateTendAccum; - real5d tracerTendAccum; - - Spatial space_op; - - void init(std::string inFile, int ny, int nx, int nens, real xlen, real ylen, int num_tracers, DataManager &dm) { - space_op.init(inFile, ny, nx, nens, xlen, ylen, num_tracers, dm); - stateTmp = space_op.createStateArr (); - tracersTmp = space_op.createTracerArr (); - - stateTend = space_op.createStateTendArr (); - tracerTend = space_op.createTracerTendArr(); - - stateTendAccum = space_op.createStateTendArr (); - tracerTendAccum = space_op.createTracerTendArr(); - } - - - template - void convert_dynamics_to_coupler_state( DataManager &dm , MICRO µ ) { - space_op.convert_dynamics_to_coupler_state( dm , micro ); - } - - - template - void convert_coupler_state_to_dynamics( DataManager &dm , MICRO µ ) { - space_op.convert_coupler_state_to_dynamics( dm , micro ); - } - - - int add_tracer(DataManager &dm , std::string name , std::string desc , bool pos_def , bool adds_mass) { - return space_op.add_tracer(dm , name , desc , pos_def , adds_mass); - } - - - template - void init_state_and_tracers( DataManager &dm , MICRO const µ ) { - space_op.init_state_and_tracers( dm , micro ); - } - - - template - void init_tracer_by_location(std::string name , F const &init_mass , DataManager &dm, MICRO const µ) const { - space_op.init_tracer_by_location(name , init_mass , dm, micro); - } - - - template - void output(DataManager &dm, MICRO const µ, real etime) const { - space_op.output(dm , micro , etime); - } - - - template - real compute_time_step(DataManager &dm, MICRO const µ, real cfl=0.8) { - return space_op.compute_time_step(dm, micro, cfl); - } - - - inline void zero_accum_arrays( real5d &stateTendAccum , real5d &tracerTendAccum ) { - int num_state = stateTendAccum .dimension[0]; - int nz = stateTendAccum .dimension[1]; - int ny = stateTendAccum .dimension[2]; - int nx = stateTendAccum .dimension[3]; - int nens = stateTendAccum .dimension[4]; - - int num_tracers = tracerTendAccum.dimension[0]; - - parallel_for( SimpleBounds<4>(nz,ny,nx,nens) , YAKL_LAMBDA (int k, int j, int i, int iens) { - for (int l=0; l < num_state; l++) { - stateTendAccum(l,k,j,i,iens) = 0; - } - for (int l=0; l < num_tracers; l++) { - tracerTendAccum(l,k,j,i,iens) = 0; - } - }); - } - - - inline void tendency_accum( real5d &stateTendAccum , real5d const &stateTend , - real5d &tracerTendAccum , real5d const &tracerTend ) { - int num_state = stateTendAccum .dimension[0]; - int nz = stateTendAccum .dimension[1]; - int ny = stateTendAccum .dimension[2]; - int nx = stateTendAccum .dimension[3]; - int nens = stateTendAccum .dimension[4]; - - int num_tracers = tracerTendAccum.dimension[0]; - - parallel_for( SimpleBounds<4>(nz,ny,nx,nens) , YAKL_LAMBDA (int k, int j, int i, int iens) { - for (int l=0; l < num_state; l++) { - stateTendAccum(l,k,j,i,iens) += stateTend(l,k,j,i,iens); - } - for (int l=0; l < num_tracers; l++) { - tracerTendAccum(l,k,j,i,iens) += tracerTend(l,k,j,i,iens); - } - }); - } - - - template - void timeStep( DataManager &dm , MICRO const µ , real dtphys ) { - YAKL_SCOPE( stateTend , this->stateTend ); - YAKL_SCOPE( stateTmp , this->stateTmp ); - YAKL_SCOPE( stateTendAccum , this->stateTendAccum ); - YAKL_SCOPE( tracerTend , this->tracerTend ); - YAKL_SCOPE( tracersTmp , this->tracersTmp ); - YAKL_SCOPE( tracerTendAccum , this->tracerTendAccum ); - - space_op.convert_coupler_state_to_dynamics( dm , micro ); - - real5d state = dm.get("dynamics_state"); - real5d tracers = dm.get("dynamics_tracers"); - - int nx = space_op.nx; - int ny = space_op.ny; - int nz = space_op.nz; - int nens = space_op.nens; - int num_state = space_op.num_state; - int num_tracers = space_op.num_tracers; - int hs = space_op.hs; - - real dt = compute_time_step( dm , micro ); - - real loctime = 0.; - while (loctime < dtphys) { - if (loctime + dt > dtphys) { dt = dtphys - loctime; } - - real dtloc = dt; - - ScalarLiveOut neg_too_large(false); - - ///////////////////////////////////// - // Stage 1 - ///////////////////////////////////// - zero_accum_arrays( stateTendAccum , tracerTendAccum ); - for (int spl = 0 ; spl < space_op.numSplit() ; spl++) { - space_op.computeTendencies( state , stateTend , tracers , tracerTend , micro , dtloc , spl ); - tendency_accum( stateTendAccum , stateTend , tracerTendAccum , tracerTend ); - } - parallel_for( SimpleBounds<4>(nz,ny,nx,nens) , YAKL_LAMBDA (int k, int j, int i, int iens) { - for (int l=0; l < num_state; l++) { - stateTmp (l,hs+k,hs+j,hs+i,iens) = state (l,hs+k,hs+j,hs+i,iens) + dtloc * stateTendAccum (l,k,j,i,iens); - } - for (int l=0; l < num_tracers; l++) { - tracersTmp(l,hs+k,hs+j,hs+i,iens) = tracers(l,hs+k,hs+j,hs+i,iens) + dtloc * tracerTendAccum(l,k,j,i,iens); - if (tracersTmp(l,hs+k,hs+j,hs+i,iens) < -1.e-10) { - neg_too_large = true; - } - tracersTmp(l,hs+k,hs+j,hs+i,iens) = max( 0._fp , tracersTmp(l,hs+k,hs+j,hs+i,iens) ); - } - }); - - ///////////////////////////////////// - // Stage 2 - ///////////////////////////////////// - zero_accum_arrays( stateTendAccum , tracerTendAccum ); - for (int spl = 0 ; spl < space_op.numSplit() ; spl++) { - space_op.computeTendencies( stateTmp , stateTend , tracersTmp , tracerTend , micro , dtloc , spl ); - tendency_accum( stateTendAccum , stateTend , tracerTendAccum , tracerTend ); - } - parallel_for( SimpleBounds<4>(nz,ny,nx,nens) , YAKL_LAMBDA (int k, int j, int i, int iens) { - for (int l=0; l < num_state; l++) { - stateTmp(l,hs+k,hs+j,hs+i,iens) = 0.75_fp * state (l,hs+k,hs+j,hs+i,iens) + - 0.25_fp * stateTmp(l,hs+k,hs+j,hs+i,iens) + - 0.25_fp * dtloc * stateTendAccum(l,k,j,i,iens); - } - for (int l=0; l < num_tracers; l++) { - tracersTmp(l,hs+k,hs+j,hs+i,iens) = 0.75_fp * tracers (l,hs+k,hs+j,hs+i,iens) + - 0.25_fp * tracersTmp(l,hs+k,hs+j,hs+i,iens) + - 0.25_fp * dtloc * tracerTendAccum(l,k,j,i,iens); - if (tracersTmp(l,hs+k,hs+j,hs+i,iens) < -1.e-10) { - neg_too_large = true; - } - tracersTmp(l,hs+k,hs+j,hs+i,iens) = max( 0._fp , tracersTmp(l,hs+k,hs+j,hs+i,iens) ); - } - }); - - ///////////////////////////////////// - // Stage 3 - ///////////////////////////////////// - zero_accum_arrays( stateTendAccum , tracerTendAccum ); - for (int spl = 0 ; spl < space_op.numSplit() ; spl++) { - space_op.computeTendencies( stateTmp , stateTend , tracersTmp , tracerTend , micro , dtloc , spl ); - tendency_accum( stateTendAccum , stateTend , tracerTendAccum , tracerTend ); - } - parallel_for( SimpleBounds<4>(nz,ny,nx,nens) , YAKL_LAMBDA (int k, int j, int i, int iens) { - for (int l=0; l < num_state; l++) { - state(l,hs+k,hs+j,hs+i,iens) = (1._fp/3._fp) * state (l,hs+k,hs+j,hs+i,iens) + - (2._fp/3._fp) * stateTmp(l,hs+k,hs+j,hs+i,iens) + - (2._fp/3._fp) * dtloc * stateTendAccum(l,k,j,i,iens); - } - for (int l=0; l < num_tracers; l++) { - tracers(l,hs+k,hs+j,hs+i,iens) = (1._fp/3._fp) * tracers (l,hs+k,hs+j,hs+i,iens) + - (2._fp/3._fp) * tracersTmp(l,hs+k,hs+j,hs+i,iens) + - (2._fp/3._fp) * dtloc * tracerTendAccum(l,k,j,i,iens); - if (tracers(l,hs+k,hs+j,hs+i,iens) < -1.e-10) { - neg_too_large = true; - } - tracers(l,hs+k,hs+j,hs+i,iens) = max( 0._fp , tracers(l,hs+k,hs+j,hs+i,iens) ); - } - }); - - #ifdef PAM_DEBUG - if (neg_too_large.hostRead()) { - std::cerr << "WARNING: Correcting a non-machine-precision negative tracer value" << std::endl; - // endrun(); - } - #endif - - loctime += dt; - } - - space_op.convert_dynamics_to_coupler_state( dm , micro ); - } - - - void finalize(DataManager &dm) { } - - - const char * dycore_name() const { return "AWFL"; } - -}; diff --git a/pam_core/DataManager.h b/pam_core/DataManager.h index d2f56706..7df15c5a 100644 --- a/pam_core/DataManager.h +++ b/pam_core/DataManager.h @@ -73,6 +73,7 @@ namespace pam { int dimid = find_dimension( name ); if (dimid > 0) { if ( dimensions[dimid].len != len ) { + std::cerr << "ERROR: For add_dimension(name,len): (" << name << "," << len << "):" << std::endl; endrun("ERROR: Adding a dimension that already exists with a different size than the existing dimension"); } return; // Avoid adding a duplicate entry @@ -94,15 +95,27 @@ namespace pam { std::vector dim_names = std::vector() , bool positive = false ) { if (name == "") { + std::cerr << "ERROR: For register_and_allocate(name,desc,dims,dim_names,positive):" << std::endl + << "name: " << name << std::endl + << "desc: " << desc << std::endl + << "positive: " << positive << std::endl; endrun("ERROR: You cannot register_and_allocate with an empty string"); } // Make sure we don't have a duplicate entry if ( find_entry(name) != -1) { + std::cerr << "ERROR: For register_and_allocate(name,desc,dims,dim_names,positive):" << std::endl + << "name: " << name << std::endl + << "desc: " << desc << std::endl + << "positive: " << positive << std::endl; endrun("ERROR: Duplicate entry name"); } if (dim_names.size() > 0) { if (dims.size() != dim_names.size()) { + std::cerr << "ERROR: For register_and_allocate(name,desc,dims,dim_names,positive):" << std::endl + << "name: " << name << std::endl + << "desc: " << desc << std::endl + << "positive: " << positive << std::endl; endrun("ERROR: Must have the same number of dims and dim_names"); } // Make sure the dimensions are the same size as existing ones of the same name @@ -115,6 +128,10 @@ namespace pam { dimensions.push_back(loc); } else { if (dimensions[dimid].len != dims[i]) { + std::cerr << "ERROR: For register_and_allocate(name,desc,dims,dim_names,positive):" << std::endl + << "name: " << name << std::endl + << "desc: " << desc << std::endl + << "positive: " << positive << std::endl; endrun("ERROR: Dimension already exists but has a different length"); } } @@ -162,15 +179,30 @@ namespace pam { std::vector dim_names = std::vector() , bool positive = false ) { if (name == "") { + std::cerr << "ERROR: For register_existing(name,desc,dims,ptr,dim_names,positive):" << std::endl + << "name: " << name << std::endl + << "desc: " << desc << std::endl + << "ptr: " << ptr << std::endl + << "positive: " << positive << std::endl; endrun("ERROR: You cannot register_and_allocate with an empty string"); } // Make sure we don't have a duplicate entry if ( find_entry(name) != -1) { + std::cerr << "ERROR: For register_existing(name,desc,dims,ptr,dim_names,positive):" << std::endl + << "name: " << name << std::endl + << "desc: " << desc << std::endl + << "ptr: " << ptr << std::endl + << "positive: " << positive << std::endl; endrun("ERROR: Duplicate entry name"); } if (dim_names.size() > 0) { if (dims.size() != dim_names.size()) { + std::cerr << "ERROR: For register_existing(name,desc,dims,ptr,dim_names,positive):" << std::endl + << "name: " << name << std::endl + << "desc: " << desc << std::endl + << "ptr: " << ptr << std::endl + << "positive: " << positive << std::endl; endrun("ERROR: Must have the same number of dims and dim_names"); } // Make sure the dimensions are the same size as existing ones of the same name @@ -183,6 +215,11 @@ namespace pam { dimensions.push_back(loc); } else { if (dimensions[dimid].len != dims[i]) { + std::cerr << "ERROR: For register_existing(name,desc,dims,ptr,dim_names,positive):" << std::endl + << "name: " << name << std::endl + << "desc: " << desc << std::endl + << "ptr: " << ptr << std::endl + << "positive: " << positive << std::endl; endrun("ERROR: Dimension already exists but has a different length"); } } @@ -302,7 +339,10 @@ namespace pam { Array get( std::string name ) { // Make sure we have this name as an entry int id = find_entry_or_error( name ); - if (entries[id].read_only) endrun("ERROR: Trying to get() a read-only arary without a const type"); + if (entries[id].read_only) { + std::cerr << "ERROR: For get(name) with name: " << name << std::endl; + endrun("ERROR: Trying to get() a read-only arary without a const type"); + } entries[id].dirty = true; // Make sure it's the right type and dimensionality validate_type(id); @@ -346,7 +386,10 @@ namespace pam { // Make sure we have this name as an entry int id = find_entry_or_error( name ); entries[id].dirty = true; - if (entries[id].read_only) endrun("ERROR: Trying to get_lev_col() a read-only arary without a const type"); + if (entries[id].read_only) { + std::cerr << "ERROR: For get_lev_col(name) with name: " << name << std::endl; + endrun("ERROR: Trying to get_lev_col() a read-only arary without a const type"); + } // Make sure it's the right type validate_type(id); validate_dims_lev_col(id); @@ -390,7 +433,10 @@ namespace pam { // Make sure we have this name as an entry int id = find_entry_or_error( name ); entries[id].dirty = true; - if (entries[id].read_only) endrun("ERROR: Trying to get_collapsed() a read-only arary without a const type"); + if (entries[id].read_only) { + std::cerr << "ERROR: For get_collapsed(name) with name: " << name << std::endl; + endrun("ERROR: Trying to get_collapsed() a read-only arary without a const type"); + } // Make sure it's the right type validate_type(id); int ncells = entries[id].dims[0]; @@ -472,6 +518,7 @@ namespace pam { auto arr = get_collapsed(name).createHostCopy(); for (int i=0; i < arr.get_elem_count(); i++) { if ( std::isnan( arr(i) ) ) { + std::cerr << "ERROR: For validate_single_nan(name) with name: " << name << std::endl; std::cerr << "WARNING: NaN discovered in: " << name << " at global index: " << i << "\n"; if (die_on_failed_check) endrun(""); } @@ -485,6 +532,7 @@ namespace pam { auto arr = get_collapsed(name).createHostCopy(); for (int i=0; i < arr.get_elem_count(); i++) { if ( std::isinf( arr(i) ) ) { + std::cerr << "ERROR: For validate_single_inf(name) with name: " << name << std::endl; std::cerr << "WARNING: inf discovered in: " << name << " at global index: " << i << "\n"; if (die_on_failed_check) endrun(""); } @@ -500,6 +548,7 @@ namespace pam { auto arr = get_collapsed(name).createHostCopy(); for (int i=0; i < arr.get_elem_count(); i++) { if ( arr(i) < 0. ) { + std::cerr << "ERROR: For validate_single_pos(name) with name: " << name << std::endl; std::cerr << "WARNING: negative value discovered in positive-definite entry: " << name << " at global index: " << i << "\n"; if (die_on_failed_check) endrun(""); @@ -531,6 +580,7 @@ namespace pam { int find_entry_or_error( std::string name ) const { int id = find_entry( name ); if (id >= 0) return id; + std::cerr << "ERROR: For find_entry_or_error(name) with name: " << name << std::endl; endrun("ERROR: Could not find entry in coupler data"); return -1; } @@ -547,7 +597,10 @@ namespace pam { // INTERNAL USE: Return the size of the named dimension or kill the run if it isn't found int get_dimension_size( std::string name ) const { int id = find_dimension( name ); - if (id == -1) { endrun("ERROR: Could not find dimension."); } + if (id == -1) { + std::cerr << "ERROR: For get_dimension_size(name) with name: " << name << std::endl; + endrun("ERROR: Could not find dimension."); + } return dimensions[id].len; } @@ -568,6 +621,7 @@ namespace pam { template void validate_type(int id) const { if ( entries[id].type_hash != get_type_hash() ) { + std::cerr << "ERROR: For validate_type(id) with id: " << id << std::endl; endrun("ERROR: Requested Array type does not match entry type"); } } @@ -578,6 +632,7 @@ namespace pam { template void validate_dims(int id) const { if ( N != entries[id].dims.size() ) { + std::cerr << "ERROR: For validate_dims(id) with id: " << id << std::endl; endrun("ERROR: Requested dimensions is different from the entry dimensions"); } } @@ -586,6 +641,7 @@ namespace pam { // INTERNAL USE: End the run if the entry id's of dimensions < 2 void validate_dims_lev_col(int id) const { if ( entries[id].dims.size() < 2 ) { + std::cerr << "ERROR: For validate_dims_lev_col(id) with id: " << id << std::endl; endrun("ERROR: Requested data is only one-dimensional"); } } diff --git a/pam_core/Options.h b/pam_core/Options.h index 26ef8c80..e4039e04 100644 --- a/pam_core/Options.h +++ b/pam_core/Options.h @@ -69,7 +69,10 @@ namespace pam { T * ptr = new T(value); options.push_back({ key , (void *) ptr , get_type_hash() , false }); } else { - if (options[id].readonly) endrun("ERROR: Trying to overwrite a readonly coupler option"); + if (options[id].readonly) { + std::cout << "ERROR: For add_option(key,value): (" << key << "," << value << "):" << std::endl; + endrun("ERROR: Trying to overwrite a readonly coupler option"); + } *((T *) options[id].data) = value; } } @@ -90,7 +93,10 @@ namespace pam { T get_option( std::string key ) const { validate_type(); int id = find_option_or_die( key ); - if (get_type_hash() != options[id].type_hash) endrun("ERROR: Requesting option using the wrong type"); + if (get_type_hash() != options[id].type_hash) { + std::cout << "ERROR: For get_option(key): (" << key << "):" << std::endl; + endrun("ERROR: Requesting option using the wrong type"); + } return *( (T *) options[id].data); } @@ -106,6 +112,7 @@ namespace pam { int find_option_or_die( std::string key ) const { int id = find_option(key); if (id >= 0) return id; + std::cout << "ERROR: For find_option_or_die(key): (" << key << "):" << std::endl; endrun("ERROR: option not found"); return -1; } diff --git a/pam_core/pam_coupler.h b/pam_core/pam_coupler.h index 175f77f7..c538d220 100644 --- a/pam_core/pam_coupler.h +++ b/pam_core/pam_coupler.h @@ -11,36 +11,6 @@ namespace pam { - - YAKL_INLINE real hydrostatic_pressure( realConst3d hy_params , real z_in , real z0 , real dz , - int k, int iens ) { - real z = ( z_in - z0 ) / dz; - real a0 = hy_params(k,0,iens); - real a1 = hy_params(k,1,iens); - real a2 = hy_params(k,2,iens); - real a3 = hy_params(k,3,iens); - real a4 = hy_params(k,4,iens); - real lnp = a0 + ( a1 + ( a2 + ( a3 + ( a4)*z)*z)*z)*z; - return exp(lnp); - } - - - - YAKL_INLINE real hydrostatic_density( realConst3d hy_params , real z_in , real z0 , real dz , - int k, int iens , real grav ) { - real z = ( z_in - z0 ) / dz; - real a1 = hy_params(k,1,iens); - real a2 = hy_params(k,2,iens); - real a3 = hy_params(k,3,iens); - real a4 = hy_params(k,4,iens); - real p = hydrostatic_pressure( hy_params , z_in , z0 , dz , k , iens ); - real mult = a1 + (2*a2 + (3*a3 + (4*a4)*z)*z)*z; - real dpdz = mult*p/dz; - return -dpdz/grav; - } - - - YAKL_INLINE real compute_pressure( real rho_d, real rho_v, real T, real R_d, real R_v ) { return rho_d*R_d*T + rho_v*R_v*T; } @@ -294,13 +264,14 @@ namespace pam { dm.register_and_allocate("vertical_interface_height","vertical interface height" ,{nz+1 ,nens},{"zp1" ,"nens"}); dm.register_and_allocate("vertical_cell_dz" ,"vertical grid spacing" ,{nz ,nens},{"z" ,"nens"}); dm.register_and_allocate("vertical_midpoint_height" ,"vertical midpoint height" ,{nz ,nens},{"z" ,"nens"}); - dm.register_and_allocate("hydrostasis_parameters" ,"hydrostasis parameters" ,{nz,5 ,nens},{"z","nhy" ,"nens"}); dm.register_and_allocate("gcm_density_dry" ,"GCM column dry density" ,{nz ,nens},{"z" ,"nens"}); dm.register_and_allocate("gcm_uvel" ,"GCM column u-velocity" ,{nz ,nens},{"z" ,"nens"}); dm.register_and_allocate("gcm_vvel" ,"GCM column v-velocity" ,{nz ,nens},{"z" ,"nens"}); dm.register_and_allocate("gcm_wvel" ,"GCM column w-velocity" ,{nz ,nens},{"z" ,"nens"}); dm.register_and_allocate("gcm_temp" ,"GCM column temperature" ,{nz ,nens},{"z" ,"nens"}); dm.register_and_allocate("gcm_water_vapor" ,"GCM column water vapor mass",{nz ,nens},{"z" ,"nens"}); + dm.register_and_allocate("gcm_cloud_water" ,"GCM column cloud water mass",{nz ,nens},{"z" ,"nens"}); + dm.register_and_allocate("gcm_cloud_ice" ,"GCM column cloud ice mass" ,{nz ,nens},{"z" ,"nens"}); auto density_dry = dm.get_collapsed("density_dry" ); auto uvel = dm.get_collapsed("uvel" ); @@ -310,13 +281,14 @@ namespace pam { auto zint = dm.get_collapsed("vertical_interface_height"); auto dz = dm.get_collapsed("vertical_cell_dz" ); auto zmid = dm.get_collapsed("vertical_midpoint_height" ); - auto hy_params = dm.get_collapsed("hydrostasis_parameters" ); auto gcm_rho_d = dm.get_collapsed("gcm_density_dry" ); auto gcm_uvel = dm.get_collapsed("gcm_uvel" ); auto gcm_vvel = dm.get_collapsed("gcm_vvel" ); auto gcm_wvel = dm.get_collapsed("gcm_wvel" ); auto gcm_temp = dm.get_collapsed("gcm_temp" ); auto gcm_rho_v = dm.get_collapsed("gcm_water_vapor" ); + auto gcm_rho_c = dm.get_collapsed("gcm_cloud_water" ); + auto gcm_rho_i = dm.get_collapsed("gcm_cloud_ice" ); parallel_for( YAKL_AUTO_LABEL() , SimpleBounds<1>(nz*ny*nx*nens) , YAKL_LAMBDA (int i) { if (i < density_dry.size()) density_dry(i) = 0; @@ -333,78 +305,9 @@ namespace pam { if (i < gcm_wvel .size()) gcm_wvel (i) = 0; if (i < gcm_temp .size()) gcm_temp (i) = 0; if (i < gcm_rho_v .size()) gcm_rho_v (i) = 0; - if (i < hy_params .size()) hy_params (i) = 0; - }); - } - - - - void update_hydrostasis() { - using yakl::c::parallel_for; - using yakl::c::SimpleBounds; - using yakl::intrinsics::matmul_cr; - using yakl::intrinsics::matinv_ge; - using yakl::atomicAdd; - - auto zint = dm.get("vertical_interface_height"); - auto zmid = dm.get("vertical_midpoint_height" ); - auto dz = dm.get("vertical_cell_dz"); - auto hy_params = dm.get("hydrostasis_parameters" ); - - auto dens_dry = dm.get("density_dry"); - auto dens_wv = dm.get("water_vapor"); - auto temp = dm.get("temp"); - - int nz = get_nz(); - int ny = get_ny(); - int nx = get_nx(); - int nens = get_nens(); - - auto R_d = get_option("R_d"); - auto R_v = get_option("R_v"); - - // Compute average column of pressure for each ensemble - real2d pressure_col("pressure_col",nz,nens); - memset( pressure_col , 0._fp ); - real r_nx_ny = 1._fp / (nx*ny); - parallel_for( SimpleBounds<4>(nz,ny,nx,nens) , YAKL_LAMBDA (int k, int j, int i, int iens) { - real rho_d = dens_dry(k,j,i,iens); - real rho_v = dens_wv (k,j,i,iens); - real T = temp (k,j,i,iens); - atomicAdd( pressure_col(k,iens) , compute_pressure( rho_d , rho_v , T , R_d , R_v )*r_nx_ny ); - }); - - parallel_for( SimpleBounds<2>(nz,nens) , YAKL_LAMBDA (int k, int iens) { - int constexpr npts = 5; - int kmid = k; - int kbot = k-2; - int ktop = k+2; - while (kbot < 0 ) { kbot++; ktop++; } - while (ktop > nz-1) { kbot--; ktop--; } - - SArray z; - real z0 = zmid(kmid,iens); - for (int i=0; i < npts; i++) { z(i) = ( zmid(kbot+i,iens) - z0 ) / dz(k,iens); } - - SArray vand; - for (int j=0; j < npts; j++) { - for (int i=0; i < npts; i++) { - vand(j,i) = pow( z(i) , (double) j ); - } - } - - auto vand_inv = matinv_ge( vand ); - - SArray logp; - for (int i=0; i < npts; i++) { - logp(i) = log(pressure_col(kbot+i,iens)); - } - - auto params = matmul_cr( vand_inv , logp ); - - for (int i=0; i < npts; i++) { hy_params(k,i,iens) = params(i); } + if (i < gcm_rho_c .size()) gcm_rho_c (i) = 0; + if (i < gcm_rho_i .size()) gcm_rho_i (i) = 0; }); - } diff --git a/standalone/mmf_simplified/build/cmakescript_pama.sh b/standalone/mmf_simplified/build/cmakescript_pama.sh index 0de0597d..e98c7dec 100755 --- a/standalone/mmf_simplified/build/cmakescript_pama.sh +++ b/standalone/mmf_simplified/build/cmakescript_pama.sh @@ -17,7 +17,6 @@ cmake \ -DPAM_MICRO="p3" \ -DPAM_RAD="none" \ -DPAM_SGS="shoc" \ - -DPAM_RAD="none" \ -DPAM_NLEV=${PAM_NLEV} \ -DSCREAM_CXX_LIBS_DIR=${SCREAM_CXX_LIBS_DIR} \ -DPAM_SCREAM_USE_CXX=${PAM_SCREAM_USE_CXX} \ diff --git a/standalone/mmf_simplified/driver.cpp b/standalone/mmf_simplified/driver.cpp index c35d1074..068254af 100644 --- a/standalone/mmf_simplified/driver.cpp +++ b/standalone/mmf_simplified/driver.cpp @@ -106,6 +106,8 @@ int main(int argc, char** argv) { auto gcm_wvel = dm.get("gcm_wvel" ); auto gcm_temp = dm.get("gcm_temp" ); auto gcm_rho_v = dm.get("gcm_water_vapor"); + auto gcm_rho_c = dm.get("gcm_cloud_water"); + auto gcm_rho_i = dm.get("gcm_cloud_ice" ); using yakl::c::parallel_for; using yakl::c::SimpleBounds; @@ -116,6 +118,8 @@ int main(int argc, char** argv) { gcm_wvel (k,iens) = wvel_col (k); gcm_temp (k,iens) = temp_col (k); gcm_rho_v(k,iens) = rho_v_col(k); + gcm_rho_c(k,iens) = 0; + gcm_rho_i(k,iens) = 0; }); if (mainproc) { @@ -139,8 +143,9 @@ int main(int argc, char** argv) { // Initialize the CRM internal state from the initial GCM column and random temperature perturbations modules::broadcast_initial_gcm_column( coupler ); - // Now that we have an initial state, define hydrostasis for each ensemble member - coupler.update_hydrostasis(); +#ifdef PAMA_DYCORE + dycore.declare_current_profile_as_hydrostatic( coupler , true ); +#endif int1d seeds("seeds",nens); seeds = 0; diff --git a/standalone/mmf_simplified/inputs/input_pama.yaml b/standalone/mmf_simplified/inputs/input_pama.yaml index 54136ea7..4eccd1b8 100644 --- a/standalone/mmf_simplified/inputs/input_pama.yaml +++ b/standalone/mmf_simplified/inputs/input_pama.yaml @@ -4,10 +4,10 @@ # simTime : 86400 # 1 day # simTime : 432000 # 5 days # simTime : 3600 # 30 min -simTime : 7200 +simTime : 86400 # Number of cells to use in the CRMs -crm_nx : 65 +crm_nx : 64 crm_ny : 1 # Number of CRMs From 82d220e7df7c52b1ffdc2f9417592ee8b6ed0113 Mon Sep 17 00:00:00 2001 From: Matt Norman Date: Wed, 20 Sep 2023 15:49:53 -0400 Subject: [PATCH 2/2] Removing coupler hydrostasis update from idealized --- standalone/idealized/driver.cpp | 3 --- 1 file changed, 3 deletions(-) diff --git a/standalone/idealized/driver.cpp b/standalone/idealized/driver.cpp index e508d61d..3835ec0a 100644 --- a/standalone/idealized/driver.cpp +++ b/standalone/idealized/driver.cpp @@ -127,9 +127,6 @@ int main(int argc, char** argv) { dycore.pre_time_loop(coupler); yakl::timer_stop("dycore"); - // Now that we have an initial state, define hydrostasis for each ensemble member - if (use_coupler_hydrostasis) coupler.update_hydrostasis( ); - real etime = 0; // There are two ways of time control- setting total simulation time (simTime) or setting number of physics time steps (simSteps) if (simTime == 0.0) { simTime = simSteps * dtphys_in; }