diff --git a/components/eamxx/src/diagnostics/field_at_height.cpp b/components/eamxx/src/diagnostics/field_at_height.cpp index 3982ed9a3ac..f61cd3a76c1 100644 --- a/components/eamxx/src/diagnostics/field_at_height.cpp +++ b/components/eamxx/src/diagnostics/field_at_height.cpp @@ -44,7 +44,7 @@ FieldAtHeight (const ekat::Comm& comm, const ekat::ParameterList& params) " - field name: " + m_field_name + "\n" " - surface reference: " + surf_ref + "\n" " - valid options: sealevel, surface\n"); - m_z_name = (surf_ref == "sealevel") ? "z" : "geopotential"; + m_z_name = (surf_ref == "sealevel") ? "z" : "height"; const auto& location = m_params.get("vertical_location"); auto chars_start = location.find_first_not_of("0123456789."); EKAT_REQUIRE_MSG (chars_start!=0 && chars_start!=std::string::npos, diff --git a/components/eamxx/src/diagnostics/register_diagnostics.hpp b/components/eamxx/src/diagnostics/register_diagnostics.hpp index efb55980a2f..21d82b045da 100644 --- a/components/eamxx/src/diagnostics/register_diagnostics.hpp +++ b/components/eamxx/src/diagnostics/register_diagnostics.hpp @@ -36,9 +36,11 @@ inline void register_diagnostics () { diag_factory.register_product("Exner",&create_atmosphere_diagnostic); diag_factory.register_product("VirtualTemperature",&create_atmosphere_diagnostic); diag_factory.register_product("z_int",&create_atmosphere_diagnostic); - diag_factory.register_product("geopotential_int",&create_atmosphere_diagnostic); diag_factory.register_product("z_mid",&create_atmosphere_diagnostic); + diag_factory.register_product("geopotential_int",&create_atmosphere_diagnostic); diag_factory.register_product("geopotential_mid",&create_atmosphere_diagnostic); + diag_factory.register_product("height_int",&create_atmosphere_diagnostic); + diag_factory.register_product("height_mid",&create_atmosphere_diagnostic); diag_factory.register_product("dz",&create_atmosphere_diagnostic); diag_factory.register_product("DryStaticEnergy",&create_atmosphere_diagnostic); diag_factory.register_product("SeaLevelPressure",&create_atmosphere_diagnostic); diff --git a/components/eamxx/src/diagnostics/tests/field_at_height_tests.cpp b/components/eamxx/src/diagnostics/tests/field_at_height_tests.cpp index 57a05711610..e26d34dd1a5 100644 --- a/components/eamxx/src/diagnostics/tests/field_at_height_tests.cpp +++ b/components/eamxx/src/diagnostics/tests/field_at_height_tests.cpp @@ -46,8 +46,8 @@ TEST_CASE("field_at_height") FieldIdentifier z_surf_fid ("z_surf", FieldLayout({COL },{ncols }),m,grid->name()); FieldIdentifier z_mid_fid ("z_mid", FieldLayout({COL, LEV},{ncols, nlevs }),m,grid->name()); FieldIdentifier z_int_fid ("z_int", FieldLayout({COL, ILEV},{ncols, nlevs+1}),m,grid->name()); - FieldIdentifier geo_mid_fid ("geopotential_mid",FieldLayout({COL, LEV},{ncols, nlevs }),m,grid->name()); - FieldIdentifier geo_int_fid ("geopotential_int",FieldLayout({COL, ILEV},{ncols, nlevs+1}),m,grid->name()); + FieldIdentifier h_mid_fid ("height_mid",FieldLayout({COL, LEV},{ncols, nlevs }),m,grid->name()); + FieldIdentifier h_int_fid ("height_int",FieldLayout({COL, ILEV},{ncols, nlevs+1}),m,grid->name()); // Keep track of reference fields for comparison FieldIdentifier s_tgt_fid ("scalar_target",FieldLayout({COL },{ncols }),m,grid->name()); FieldIdentifier v_tgt_fid ("vector_target",FieldLayout({COL,CMP},{ncols,ndims}),m,grid->name()); @@ -59,8 +59,8 @@ TEST_CASE("field_at_height") Field z_surf (z_surf_fid); Field z_mid (z_mid_fid); Field z_int (z_int_fid); - Field geo_mid (geo_mid_fid); - Field geo_int (geo_int_fid); + Field h_mid (h_mid_fid); + Field h_int (h_int_fid); Field s_tgt (s_tgt_fid); Field v_tgt (v_tgt_fid); @@ -71,8 +71,8 @@ TEST_CASE("field_at_height") z_surf.allocate_view(); z_mid.allocate_view(); z_int.allocate_view(); - geo_mid.allocate_view(); - geo_int.allocate_view(); + h_mid.allocate_view(); + h_int.allocate_view(); s_tgt.allocate_view(); v_tgt.allocate_view(); @@ -83,8 +83,8 @@ TEST_CASE("field_at_height") z_surf.get_header().get_tracking().update_time_stamp(t0); z_mid.get_header().get_tracking().update_time_stamp(t0); z_int.get_header().get_tracking().update_time_stamp(t0); - geo_mid.get_header().get_tracking().update_time_stamp(t0); - geo_int.get_header().get_tracking().update_time_stamp(t0); + h_mid.get_header().get_tracking().update_time_stamp(t0); + h_int.get_header().get_tracking().update_time_stamp(t0); s_tgt.get_header().get_tracking().update_time_stamp(t0); v_tgt.get_header().get_tracking().update_time_stamp(t0); @@ -124,9 +124,9 @@ TEST_CASE("field_at_height") // Set up vertical structure for the tests. Note, // z_mid/int represents the height in m above sealevel - // geo_mid/int represente the hegith in m above the surface + // h_mid/int represente the hegith in m above the surface // So we first construct z_mid/int using z_surf as reference, and - // then can build geo_mid/int from z_mid/int + // then can build h_mid/int from z_mid/int // Furthermore, z_mid is just the midpoint between two adjacent z_int // points, so we back z_mid out of z_int. // @@ -137,8 +137,8 @@ TEST_CASE("field_at_height") const auto& zint_v = z_int.get_view(); const auto& zmid_v = z_mid.get_view(); const auto& zsurf_v = z_surf.get_view(); - const auto& geoint_v = geo_int.get_view(); - const auto& geomid_v = geo_mid.get_view(); + const auto& geoint_v = h_int.get_view(); + const auto& geomid_v = h_mid.get_view(); int min_col_thickness = z_top; int max_surf = 0; for (int ii=0; ii Testing throws error with unsupported reference height...\n"); { - REQUIRE_THROWS(run_diag (s_mid,geo_mid,"1m","foobar")); + REQUIRE_THROWS(run_diag (s_mid,h_mid,"1m","foobar")); } print(" -> Testing throws error with unsupported reference height... OK\n"); @@ -182,8 +182,8 @@ TEST_CASE("field_at_height") std::string loc; for (std::string surf_ref : {"sealevel","surface"}) { printf(" -> Testing for a reference height above %s...\n",surf_ref.c_str()); - const auto mid_src = surf_ref == "sealevel" ? z_mid : geo_mid; - const auto int_src = surf_ref == "sealevel" ? z_int : geo_int; + const auto mid_src = surf_ref == "sealevel" ? z_mid : h_mid; + const auto int_src = surf_ref == "sealevel" ? z_int : h_int; const int max_surf_4test = surf_ref == "sealevel" ? max_surf : 0; for (int irun=0; irun Forced extrapolation ...............\n"); + print(" -> Forced extrapolation at top...............\n"); auto slope = pdf_m(engine); auto inter = pdf_y0(engine); f_z_src(inter, slope, int_src, s_int); - print(" -> at top...............\n"); z_tgt = 2*z_top; std::string loc = std::to_string(z_tgt) + "m"; auto dtop = run_diag(s_int,int_src,loc,surf_ref); f_z_tgt(inter,slope,z_tgt,int_src,s_tgt); REQUIRE (views_are_approx_equal(dtop,s_tgt,tol)); - print(" -> at bot...............\n"); + print(" -> Forced extrapolation at top............... OK!\n"); + print(" -> Forced extrapolation at bot...............\n"); z_tgt = 0; loc = std::to_string(z_tgt) + "m"; auto dbot = run_diag(s_int,int_src,loc,surf_ref); f_z_tgt(inter,slope,z_tgt,int_src,s_tgt); REQUIRE (views_are_approx_equal(dbot,s_tgt,tol)); - print(" -> Forced extrapolation............... OK!\n"); + print(" -> Forced extrapolation at bot............... OK!\n"); } printf(" -> Testing for a reference height above %s... OK!\n",surf_ref.c_str()); } diff --git a/components/eamxx/src/diagnostics/tests/vertical_layer_tests.cpp b/components/eamxx/src/diagnostics/tests/vertical_layer_tests.cpp index 32402955f8d..fe75114611d 100644 --- a/components/eamxx/src/diagnostics/tests/vertical_layer_tests.cpp +++ b/components/eamxx/src/diagnostics/tests/vertical_layer_tests.cpp @@ -1,20 +1,8 @@ #include "catch2/catch.hpp" -#include "share/grid/mesh_free_grids_manager.hpp" -#include "diagnostics/vertical_layer.hpp" #include "diagnostics/register_diagnostics.hpp" - #include "physics/share/physics_constants.hpp" - -#include "share/util/scream_setup_random_test.hpp" -#include "share/util/scream_common_physics_functions.hpp" -#include "share/field/field_utils.hpp" - -#include "ekat/ekat_pack.hpp" -#include "ekat/kokkos/ekat_kokkos_utils.hpp" -#include "ekat/util/ekat_test_utils.hpp" - -#include +#include "share/grid/mesh_free_grids_manager.hpp" namespace scream { @@ -39,23 +27,13 @@ create_gm (const ekat::Comm& comm, const int ncols, const int nlevs) { } //-----------------------------------------------------------------------------------------------// -template -void run(std::mt19937_64& engine, std::string diag_type, const bool from_sea_level = false) +template +void run (const std::string& diag_name, const std::string& location) { - using PF = scream::PhysicsFunctions; - using PC = scream::physics::Constants; - using Pack = ekat::Pack; - using KT = ekat::KokkosTypes; - using ExecSpace = typename KT::ExeSpace; - using MemberType = typename KT::MemberType; - using view_1d = typename KT::template view_1d; - using rview_1d = typename KT::template view_1d; - using view_2d = typename KT::template view_2d; - - const int packsize = SCREAM_PACK_SIZE; + using PC = scream::physics::Constants; + + const int packsize = N; constexpr int num_levs = packsize*2 + 1; // Number of levels to use for tests, make sure the last pack can also have some empty slots (packsize>1). - const int num_mid_packs = ekat::npack(num_levs); - const int num_mid_packs_p1 = ekat::npack(num_levs+1); // A world comm ekat::Comm comm(MPI_COMM_WORLD); @@ -64,51 +42,25 @@ void run(std::mt19937_64& engine, std::string diag_type, const bool from_sea_lev const int ncols = 1; auto gm = create_gm(comm,ncols,num_levs); - // Kokkos Policy - auto policy = ekat::ExeSpaceUtils::get_default_team_policy(ncols, num_mid_packs); - - // Input (randomized) views - view_1d temperature("temperature",num_mid_packs), - pseudodensity("pseudodensity",num_mid_packs), - pressure("pressure",num_mid_packs), - watervapor("watervapor",num_mid_packs); - - auto dview_as_real = [&] (const view_1d& v) -> rview_1d { - return rview_1d(reinterpret_cast(v.data()),v.size()*packsize); - }; - - // Construct random input data - using RPDF = std::uniform_real_distribution; - RPDF pdf_qv(1e-6,1e-3), - pdf_pseudodens(1.0,100.0), - pdf_pres(0.0,PC::P0), - pdf_temp(200.0,400.0), - pdf_phis(0.0,10000.0); - // A time stamp util::TimeStamp t0 ({2022,1,1},{0,0,0}); + register_diagnostics(); + auto& diag_factory = AtmosphereDiagnosticFactory::instance(); + // Construct the Diagnostic ekat::ParameterList params; - std::string diag_name; - if (diag_type == "thickness") { - diag_name = "dz"; + std::string name = diag_name; + if (location=="midpoints") { + name += "_mid"; + } else if (location=="interfaces") { + name += "_int"; } - else if (diag_type == "interface") { - diag_name = from_sea_level ? "z_int" : "geopotential_int"; - } else if (diag_type == "midpoint") { - diag_name = from_sea_level ? "z_mid" : "geopotential_mid"; - } - params.set("diag_name", diag_name); - register_diagnostics(); - auto& diag_factory = AtmosphereDiagnosticFactory::instance(); - auto diag = diag_factory.create(diag_name,comm,params); + params.set("diag_name", name); + auto diag = diag_factory.create(name,comm,params); diag->set_grids(gm); - // Helpful bools - const bool only_compute_dz = (diag_type == "thickness"); - const bool is_interface_layout = (diag_type == "interface"); - const bool generate_phis_data = (not only_compute_dz and not from_sea_level); + const bool needs_phis = diag_name=="z" or diag_name=="geopotential"; // Set the required fields for the diagnostic. std::map input_fields; @@ -120,89 +72,80 @@ void run(std::mt19937_64& engine, std::string diag_type, const bool from_sea_lev const auto name = f.name(); f.get_header().get_tracking().update_time_stamp(t0); diag->set_required_field(f.get_const()); - REQUIRE_THROWS(diag->set_computed_field(f)); input_fields.emplace(name,f); } - // Initialize the diagnostic + // Can't set computed fields in the diag + REQUIRE_THROWS(diag->set_computed_field(input_fields.begin()->second)); + + // Note: we are not testing the calculate_dz utility. We are testing + // the diag class, so use some inputs that make checking results easier + // With these inputs, T_virt=T, and dz=8*rd/g + const Real g = PC::gravit; + const Real rho_val = 4; + const Real qv_val = 0; + const Real p_val = 2; + const Real T_val = 4; + const Real c1 = -PC::ONE + PC::ONE / PC::ep_2; + const Real Tvirt_val = T_val*(PC::ONE + c1*qv_val); + const Real dz_val = (PC::RD/g) * rho_val*Tvirt_val / p_val; + const Real phis_val = 3; + + input_fields["T_mid"].deep_copy(T_val); + input_fields["p_mid"].deep_copy(p_val); + input_fields["pseudo_density"].deep_copy(rho_val); + input_fields["qv"].deep_copy(qv_val); + if (needs_phis) { + input_fields["phis"].deep_copy(phis_val); + } + + // Initialize and run the diagnostic diag->initialize(t0,RunType::Initial); + diag->compute_diagnostic(); + const auto& diag_out = diag->get_diagnostic(); + diag_out.sync_to_host(); + auto d_h = diag_out.get_view(); + + // Compare against expected value + const auto last_int = num_levs; + const auto last_mid = last_int-1; + + // Precompute surface value and increment depending on the diag type + Real delta, surf_val; + if (diag_name=="height") { + surf_val = 0; + delta = dz_val; + } else if (diag_name=="z") { + surf_val = phis_val/g; + delta = dz_val; + } else { + surf_val = phis_val; + delta = dz_val*g; + } - // Run tests - { - // Construct random data to use for test - // Get views of input data and set to random values - const auto& T_mid_f = input_fields["T_mid"]; - const auto& T_mid_v = T_mid_f.get_view(); - const auto& pseudo_dens_f = input_fields["pseudo_density"]; - const auto& pseudo_dens_v = pseudo_dens_f.get_view(); - const auto& p_mid_f = input_fields["p_mid"]; - const auto& p_mid_v = p_mid_f.get_view(); - const auto& qv_mid_f = input_fields["qv"]; - const auto& qv_mid_v = qv_mid_f.get_view(); - Field phis_f; - rview_1d phis_v; - if (generate_phis_data) { - phis_f = input_fields["phis"]; - phis_v = phis_f.get_view(); - } + for (int icol=0; icolcompute_diagnostic(); - const auto& diag_out = diag->get_diagnostic(); - - // Need to generate temporary values for calculation - const auto& dz_v = view_2d("",ncols, num_mid_packs); - const auto& zmid_v = view_2d("",ncols, num_mid_packs); - const auto& zint_v = view_2d("",ncols, num_mid_packs_p1); - - Kokkos::parallel_for("", policy, KOKKOS_LAMBDA(const MemberType& team) { - const int icol = team.league_rank(); - - const auto& dz_s = ekat::subview(dz_v,icol); - Kokkos::parallel_for(Kokkos::TeamVectorRange(team,num_mid_packs), [&] (const Int& jpack) { - dz_s(jpack) = PF::calculate_dz(pseudo_dens_v(icol,jpack),p_mid_v(icol,jpack),T_mid_v(icol,jpack),qv_mid_v(icol,jpack)); - }); - team.team_barrier(); - if (not only_compute_dz) { - const auto& zint_s = ekat::subview(zint_v,icol); - const Real surf_geopotential = from_sea_level ? 0.0 : phis_v(icol); - PF::calculate_z_int(team,num_levs,dz_s,surf_geopotential,zint_s); - - if (not is_interface_layout) { - const auto& zmid_s = ekat::subview(zmid_v,icol); - PF::calculate_z_mid(team,num_levs,zint_s,zmid_s); + for (int ilev=last_mid; ilev>=0; --ilev) { + if (diag_name=="dz") { + REQUIRE (d_h(icol,ilev)==dz_val); + } else { + // If interface, check value, otherwise perform int->mid averaging and check value + auto int_val = prev_int_val + delta; + if (location=="interfaces") { + REQUIRE_THAT(d_h(icol,ilev), Catch::Matchers::WithinRel(int_val,1e-5)); + } else { + auto mid_val = (int_val + prev_int_val) / 2; + REQUIRE_THAT(d_h(icol,ilev), Catch::Matchers::WithinRel(mid_val,1e-5)); } + prev_int_val = int_val; } - }); - Kokkos::fence(); - - Field diag_calc = diag_out.clone(); - auto field_v = diag_calc.get_view(); - - if (diag_type == "thickness") Kokkos::deep_copy(field_v, dz_v); - else if (diag_type == "interface") Kokkos::deep_copy(field_v, zint_v); - else if (diag_type == "midpoint") Kokkos::deep_copy(field_v, zmid_v); - diag_calc.sync_to_host(); - REQUIRE(views_are_equal(diag_out,diag_calc)); + } } // Finalize the diagnostic @@ -215,29 +158,36 @@ TEST_CASE("vertical_layer_test", "vertical_layer_test]"){ using scream::Real; using Device = scream::DefaultDevice; - constexpr int num_runs = 5; - - auto engine = scream::setup_random_test(); - - printf(" -> Number of randomized runs: %d, Pack scalar type\n\n", num_runs, SCREAM_PACK_SIZE); + ekat::Comm comm(MPI_COMM_WORLD); + auto root_print = [&](const std::string& msg) { + if (comm.am_i_root()) { + printf("%s",msg.c_str()); + } + }; + auto do_run = [&] (auto int_const) { + constexpr int N = decltype(int_const)::value; + root_print("\n"); + root_print(" -> Testing diagnostic for pack_size=" + std::to_string(N) + "\n"); + for (std::string loc : {"midpoints","interfaces"}) { + for (std::string diag : {"geopotential","height","z"}) { + std::string msg = " -> Testing diag=" + diag + " at " + loc + " "; + std::string dots (50-msg.size(),'.'); + root_print (msg + dots + "\n"); + run(diag, loc); + root_print (msg + dots + " PASS!\n"); + } + } + std::string msg = " -> Testing diag=dz "; + std::string dots (50-msg.size(),'.'); + root_print (msg + dots + "\n"); + run("dz", "UNUSED"); + root_print (msg + dots + " PASS!\n"); + }; - printf(" -> Testing dz..."); - for (int irun=0; irun(engine, "thickness"); - } - printf("ok!\n"); - printf(" -> Testing z_int/geopotential_int..."); - for (int irun=0; irun(engine, "interface", irun%2==0); // alternate from_sea_level=true/false + if (SCREAM_PACK_SIZE!=1) { + do_run(std::integral_constant()); } - printf("ok!\n"); - printf(" -> Testing z_mid/geopotential_mid..."); - for (int irun=0; irun(engine, "midpoint", irun%2==0); // alternate from_sea_level=true/false - } - printf("ok!\n"); - - printf("\n"); + do_run(std::integral_constant()); } // TEST_CASE diff --git a/components/eamxx/src/diagnostics/vertical_layer.cpp b/components/eamxx/src/diagnostics/vertical_layer.cpp index 8766649cbaf..32d870da03c 100644 --- a/components/eamxx/src/diagnostics/vertical_layer.cpp +++ b/components/eamxx/src/diagnostics/vertical_layer.cpp @@ -1,5 +1,9 @@ #include "diagnostics/vertical_layer.hpp" +#include "physics/share/physics_constants.hpp" +#include "share/util/scream_common_physics_functions.hpp" +#include "share/util/scream_column_ops.hpp" + namespace scream { @@ -9,18 +13,25 @@ VerticalLayerDiagnostic (const ekat::Comm& comm, const ekat::ParameterList& para : AtmosphereDiagnostic(comm,params) { m_diag_name = params.get("diag_name"); - EKAT_REQUIRE_MSG(m_diag_name == "z_int" or m_diag_name == "z_mid" or - m_diag_name == "geopotential_int" or m_diag_name == "geopotential_mid" or - m_diag_name == "dz", - "Error! VerticalLayerDiagnostic has been given an unknown name: "+m_diag_name+".\n"); + std::vector supported = { + "z_int", + "z_mid", + "geopotential_int", + "geopotential_mid", + "height_int", + "height_mid", + "dz" + }; + + EKAT_REQUIRE_MSG(ekat::contains(supported,m_diag_name), + "[VerticalLayerDiagnostic] Error! Invalid diag_name.\n" + " - diag_name : " + m_diag_name + "\n" + " - valid names: " + ekat::join(supported,", ") + "\n"); - m_only_compute_dz = (m_diag_name == "dz"); m_is_interface_layout = m_diag_name.find("_int") != std::string::npos; - // Whether or not diagnostic is computed from sea level depends on the name. - // "z_" -> from sea level, "geopotential_" -> from topography data. - // This boolean is irrelevant for vertical layer thickness (dz). - m_from_sea_level = m_diag_name.find("z_") != std::string::npos; + m_geopotential = m_diag_name.substr(0,12)=="geopotential"; + m_from_sea_level = m_diag_name[0]=='z' or m_geopotential; } // ======================================================================================== void VerticalLayerDiagnostic:: @@ -37,101 +48,182 @@ set_grids(const std::shared_ptr grids_manager) m_num_cols = grid->get_num_local_dofs(); // Number of columns on this rank m_num_levs = grid->get_num_vertical_levels(); // Number of levels per column - FieldLayout scalar2d_layout { {COL }, {m_num_cols } }; - FieldLayout scalar3d_layout_mid { {COL,LEV }, {m_num_cols,m_num_levs } }; - FieldLayout scalar3d_layout_int { {COL,ILEV}, {m_num_cols,m_num_levs+1} }; - constexpr int ps = Pack::n; + const auto scalar2d = grid->get_2d_scalar_layout(); + const auto scalar3d_mid = grid->get_3d_scalar_layout(true); + const auto scalar3d_int = grid->get_3d_scalar_layout(false); // The fields required for this diagnostic to be computed - add_field("T_mid", scalar3d_layout_mid, K, grid_name, ps); - add_field("pseudo_density", scalar3d_layout_mid, Pa, grid_name, ps); - add_field("p_mid", scalar3d_layout_mid, Pa, grid_name, ps); - add_field("qv", scalar3d_layout_mid, kg/kg, grid_name, ps); + add_field("T_mid", scalar3d_mid, K, grid_name); + add_field("pseudo_density", scalar3d_mid, Pa, grid_name); + add_field("p_mid", scalar3d_mid, Pa, grid_name); + add_field("qv", scalar3d_mid, kg/kg, grid_name); // Only need phis if computing geopotential_* - if (not m_only_compute_dz and not m_from_sea_level) { - add_field("phis", scalar2d_layout, m2/s2, grid_name); + if (m_from_sea_level) { + add_field("phis", scalar2d, m2/s2, grid_name); } +} + +void VerticalLayerDiagnostic:: +initialize_impl (const RunType /*run_type*/) +{ + using namespace ShortFieldTagsNames; + using namespace ekat::units; + + auto m2 = pow(m,2); + auto s2 = pow(s,2); + + const auto& T = get_field_in("T_mid"); + const auto& rho = get_field_in("pseudo_density"); + const auto& p = get_field_in("p_mid"); + const auto& qv = get_field_in("qv"); + const auto& phis = m_from_sea_level ? get_field_in("phis") : T; // unused if m_from_sea_level=false + + // Construct and allocate the diagnostic field. + // Notes: + // - consider diag name to set long name + // - check input fields alloc props to set alloc props for output + + const auto& grid_name = T.get_header().get_identifier().get_grid_name(); + const auto VLEV = m_is_interface_layout ? ILEV : LEV; + const auto nlevs = m_is_interface_layout ? m_num_levs+1 : m_num_levs; + FieldLayout diag_layout ({COL,VLEV},{m_num_cols,nlevs}); + FieldIdentifier fid (name(), diag_layout, m_geopotential ? m2/s2 : m, grid_name); - // Construct and allocate the diagnostic field based on the diagnostic name. - const auto diag_layout = m_is_interface_layout ? scalar3d_layout_int : scalar3d_layout_mid; - FieldIdentifier fid (name(), diag_layout, m, grid_name); m_diagnostic_output = Field(fid); - auto& C_ap = m_diagnostic_output.get_header().get_alloc_properties(); - C_ap.request_allocation(ps); + auto& diag_fap = m_diagnostic_output.get_header().get_alloc_properties(); + int ps = SCREAM_PACK_SIZE; + for (const auto& f : {T,rho,p,qv,phis}) { + const auto& fap = f.get_header().get_alloc_properties(); + const auto& f_ps = fap.get_largest_pack_size(); + + // We must use a pack size that works with all inputs, so pick the smallest + ps = std::min(ps,f_ps); + } + diag_fap.request_allocation(ps); m_diagnostic_output.allocate_view(); - // Initialize temporary views based on need. - if (not m_only_compute_dz) { - if (m_is_interface_layout) { - const auto npacks = ekat::npack(m_num_levs); - m_tmp_midpoint_view = view_2d("tmp_mid",m_num_cols,npacks); - } else { - const auto npacks_p1 = ekat::npack(m_num_levs+1); - m_tmp_interface_view = view_2d("tmp_int",m_num_cols,npacks_p1); - } + using stratts_t = std::map; + auto& io_atts = m_diagnostic_output.get_header().get_extra_data("io: string attributes"); + auto& long_name = io_atts["long_name"]; + if (m_diag_name=="dz") { + long_name = "level thickness"; + } else if (m_diag_name=="z_mid") { + long_name = "elevation above sealevel at level midpoints"; + } else if (m_diag_name=="z_int") { + long_name = "elevation above sealevel at level interfaces"; + } else if (m_diag_name=="height_mid") { + long_name= "elevation above surface at level midpoints"; + } else if (m_diag_name=="height_int") { + long_name = "elevation above surface at level interfaces"; + } else if (m_diag_name=="geopotential_mid") { + long_name = "geopotential height relative to sealevel at level midpoints"; + } else { + long_name = "geopotential height relative to sealevel at level interfaces"; + } + + // Initialize temporary views based on need. Can alias the diag if a temp is not needed + auto create_temp = [&](const std::string& name, int levs) { + auto u = Units::nondimensional(); + FieldLayout fl({COL,LEV},{m_num_cols,levs}); + FieldIdentifier fid (name,fl,u,grid_name); + Field f = Field(fid); + f.get_header().get_alloc_properties().request_allocation(ps); + f.allocate_view(); + return f; + }; + if (m_diag_name == "dz") { + m_tmp_midpoint = m_diagnostic_output; + m_tmp_interface = m_diagnostic_output; // Not really used + } else if (m_is_interface_layout) { + m_tmp_midpoint = create_temp("tmp_mid",m_num_levs); + m_tmp_interface = m_diagnostic_output; + } else { + m_tmp_interface = create_temp("tmp_int",m_num_levs+1); + m_tmp_midpoint = m_diagnostic_output; } } // ========================================================================================= void VerticalLayerDiagnostic::compute_diagnostic_impl() { - const auto npacks = ekat::npack(m_num_levs); - const auto default_policy = ekat::ExeSpaceUtils::get_thread_range_parallel_scan_team_policy(m_num_cols, npacks); + const auto& fap = m_diagnostic_output.get_header().get_alloc_properties(); + if (fap.get_largest_pack_size()==SCREAM_PACK_SIZE) { + do_compute_diagnostic_impl(); + } else { + do_compute_diagnostic_impl<1>(); + } +} - const auto& T_mid = get_field_in("T_mid").get_view(); - const auto& p_mid = get_field_in("p_mid").get_view(); - const auto& qv_mid = get_field_in("qv").get_view(); - const auto& pseudo_density_mid = get_field_in("pseudo_density").get_view(); +template +void VerticalLayerDiagnostic::do_compute_diagnostic_impl() +{ + using column_ops = ColumnOps; + using PackT = ekat::Pack; + using KT = KokkosTypes; + using MemberType = typename KT::MemberType; + using PF = PhysicsFunctions; - view_1d_const phis; - if (not m_only_compute_dz and not m_from_sea_level) { - phis = get_field_in("phis").get_view(); - } + // To use in column_ops, since we integrate from surface + constexpr bool FromTop = false; + + const auto npacks = ekat::npack(m_num_levs); + const auto policy = ekat::ExeSpaceUtils::get_thread_range_parallel_scan_team_policy(m_num_cols, npacks); + + const auto& T = get_field_in("T_mid").get_view(); + const auto& p = get_field_in("p_mid").get_view(); + const auto& qv = get_field_in("qv").get_view(); + const auto& rho = get_field_in("pseudo_density").get_view(); + const auto phis = m_from_sea_level ? get_field_in("phis").get_view() : typename KT::view_1d(); - const bool only_compute_dz = m_only_compute_dz; + const bool only_compute_dz = m_diag_name=="dz"; const bool is_interface_layout = m_is_interface_layout; const bool from_sea_level = m_from_sea_level; + const bool geopotential = m_geopotential; const int num_levs = m_num_levs; + constexpr auto g = scream::physics::Constants::gravit; // Alias correct view for diagnostic output and for tmp class views - view_2d interface_view; - view_2d midpoint_view; - if (only_compute_dz) { - midpoint_view = m_diagnostic_output.get_view(); - } else if (is_interface_layout) { - interface_view = m_diagnostic_output.get_view(); - midpoint_view = m_tmp_midpoint_view; - } else { - midpoint_view = m_diagnostic_output.get_view(); - interface_view = m_tmp_interface_view; - } + auto tmp_mid = m_tmp_midpoint.get_view(); + auto tmp_int = m_tmp_interface.get_view(); - Kokkos::parallel_for("VerticalLayerDiagnostic", - default_policy, - KOKKOS_LAMBDA(const MemberType& team) { + // Define the lambda, then dispatch the ||for + auto lambda = KOKKOS_LAMBDA(const MemberType& team) { const int icol = team.league_rank(); - // Calculate dz. Use the memory in tmp_mid_view for dz and z_mid, - // since we don't set z_mid until after dz is no longer needed. - const auto& dz_s = ekat::subview(midpoint_view, icol); - Kokkos::parallel_for(Kokkos::TeamVectorRange(team, npacks), [&] (const Int& jpack) { - dz_s(jpack) = PF::calculate_dz(pseudo_density_mid(icol,jpack), p_mid(icol,jpack), T_mid(icol,jpack), qv_mid(icol,jpack)); - }); + // Whatever the output needs, the first thing to compute is dz. + const auto& dz = ekat::subview(tmp_mid, icol); + PF::calculate_dz(team,ekat::subview(rho,icol), + ekat::subview(p,icol), + ekat::subview(T,icol), + ekat::subview(qv,icol), + dz); team.team_barrier(); - if (not only_compute_dz) { - // Calculate z_int if this diagnostic is not dz - const auto& z_int_s = ekat::subview(interface_view, icol); - const Real surf_geopotential = from_sea_level ? 0.0 : phis(icol); - PF::calculate_z_int(team,num_levs,dz_s,surf_geopotential,z_int_s); - - if (not is_interface_layout) { - // Calculate z_mid if this diagnostic is not dz or an interface value - const auto& z_mid_s = ekat::subview(midpoint_view, icol); - PF::calculate_z_mid(team,num_levs,z_int_s,z_mid_s); - } + // If dz is all we need, we're done + if (only_compute_dz) { return; } + + // Now integrate to compute quantity at interfaces + const auto& v_int = ekat::subview(tmp_int, icol); + + // phi and z are related by phi=z*g, so dphi=dz*g, and z_surf = phis/g + if (geopotential) { + auto dphi = [&](const int ilev) { + return dz(ilev) * g; + }; + column_ops::template column_scan(team,num_levs,dphi,v_int,phis(icol)); + } else { + const Real surf_val = from_sea_level ? phis(icol)/g : 0; + column_ops::template column_scan(team,num_levs,dz,v_int,surf_val); } - }); + + // If we need quantity at midpoints, simply do int->mid averaging + if (not is_interface_layout) { + team.team_barrier(); + const auto& v_mid = ekat::subview(tmp_mid, icol); + column_ops::compute_midpoint_values(team,num_levs,v_int,v_mid); + } + }; + Kokkos::parallel_for(m_diag_name, policy, lambda); } -// ========================================================================================= + } //namespace scream diff --git a/components/eamxx/src/diagnostics/vertical_layer.hpp b/components/eamxx/src/diagnostics/vertical_layer.hpp index a7a6bbf68f5..805fd70028f 100644 --- a/components/eamxx/src/diagnostics/vertical_layer.hpp +++ b/components/eamxx/src/diagnostics/vertical_layer.hpp @@ -2,8 +2,6 @@ #define EAMXX_VERTICAL_LAY_MID_DIAGNOSTIC_HPP #include "share/atm_process/atmosphere_diagnostic.hpp" -#include "share/util/scream_common_physics_functions.hpp" -#include "ekat/kokkos/ekat_subview_utils.hpp" namespace scream { @@ -22,13 +20,6 @@ namespace scream class VerticalLayerDiagnostic : public AtmosphereDiagnostic { public: - using Pack = ekat::Pack; - using PF = scream::PhysicsFunctions; - using KT = KokkosTypes; - using MemberType = typename KT::MemberType; - using view_1d_const = typename KT::template view_1d; - using view_2d = typename KT::template view_2d; - // Constructors VerticalLayerDiagnostic (const ekat::Comm& comm, const ekat::ParameterList& params); @@ -39,35 +30,35 @@ class VerticalLayerDiagnostic : public AtmosphereDiagnostic void set_grids (const std::shared_ptr grids_manager); protected: + void compute_diagnostic_impl (); + void initialize_impl (const RunType /* run_type */); #ifdef KOKKOS_ENABLE_CUDA public: #endif - void compute_diagnostic_impl (); + template + void do_compute_diagnostic_impl (); protected: // Keep track of field dimensions Int m_num_cols; Int m_num_levs; - // Temporary view to set dz, z_mid, and z_int - view_2d m_tmp_interface_view; - view_2d m_tmp_midpoint_view; + // Temporaries to use for calculation of dz, z_int, and z_mid + Field m_tmp_interface; + Field m_tmp_midpoint; - // The diagnostic name. This will dictate which - // field in the computation is output (dz, z_int, or z_mid). + // The diagnostic name. std::string m_diag_name; - // Store if we only need to compute dz to save computation/memory requirements. - bool m_only_compute_dz; - // Store if the diagnostic output field exists on interface values bool m_is_interface_layout; - // If z_int or z_mid is computed, determine whether the BC - // is from sea level or not (from topography data). + // True z_mid/int and geopotential_mid/int, false for height_mid/int. Unused for dz bool m_from_sea_level; -}; // class VerticalLayerDiagnostic + // If true, output is a geopotential (units m2/s2), otherwise an elevation + bool m_geopotential; +}; } //namespace scream diff --git a/components/eamxx/src/share/io/scorpio_output.cpp b/components/eamxx/src/share/io/scorpio_output.cpp index 1852785eeac..e79d75e97b1 100644 --- a/components/eamxx/src/share/io/scorpio_output.cpp +++ b/components/eamxx/src/share/io/scorpio_output.cpp @@ -955,9 +955,6 @@ register_variables(const std::string& filename, auto vec_of_dims = set_vec_of_dims(layout); std::string units = fid.get_units().to_string(); - // Gather longname - auto longname = m_longnames.get_longname(name); - // TODO Need to change dtype to allow for other variables. // Currently the field_manager only stores Real variables so it is not an issue, // but in the future if non-Real variables are added we will want to accomodate that. @@ -991,8 +988,6 @@ register_variables(const std::string& filename, scorpio::define_var (filename, name, units, vec_of_dims, "real",fp_precision, m_add_time_dim); - scorpio::set_attribute(filename, name, "long_name", longname); - // Add FillValue as an attribute of each variable // FillValue is a protected metadata, do not add it if it already existed if (fp_precision=="double" or @@ -1033,6 +1028,12 @@ register_variables(const std::string& filename, for (const auto& [att_name,att_val] : str_atts) { scorpio::set_attribute(filename,name,att_name,att_val); } + + // Gather longname (if not already in the io: string attributes) + if (str_atts.count("long_name")==0) { + auto longname = m_longnames.get_longname(name); + scorpio::set_attribute(filename, name, "long_name", longname); + } } } // Now register the average count variables @@ -1315,25 +1316,25 @@ AtmosphereOutput::create_diagnostic (const std::string& diag_field_name) { if (units.find("_above_") != std::string::npos) { // The field is at a height above a specific reference. // Currently we only support FieldAtHeight above "sealevel" or "surface" - auto subtokens = ekat::split(units,"_above_"); - params.set("surface_reference",subtokens[1]); - units = subtokens[0]; - // Need to reset the vertical location to strip the "_above_" part of the string. - params.set("vertical_location", tokens[1].substr(0,units_start)+subtokens[0]); - // If the slice is "above_sealevel" then we need to track the avg cnt uniquely. - // Note, "above_surface" is expected to never have masking and can thus use - // the typical 2d layout avg cnt. - if (subtokens[1]=="sealevel") { - diag_avg_cnt_name = "_" + tokens[1]; // Set avg_cnt tracking for this specific slice - // If we have 2D slices we need to be tracking the average count, - // if m_avg_type is not Instant - m_track_avg_cnt = m_track_avg_cnt || m_avg_type!=OutputAvgType::Instant; - } + auto subtokens = ekat::split(units,"_above_"); + params.set("surface_reference",subtokens[1]); + units = subtokens[0]; + // Need to reset the vertical location to strip the "_above_" part of the string. + params.set("vertical_location", tokens[1].substr(0,units_start)+subtokens[0]); + // If the slice is "above_sealevel" then we need to track the avg cnt uniquely. + // Note, "above_surface" is expected to never have masking and can thus use + // the typical 2d layout avg cnt. + if (subtokens[1]=="sealevel") { + diag_avg_cnt_name = "_" + tokens[1]; // Set avg_cnt tracking for this specific slice + // If we have 2D slices we need to be tracking the average count, + // if m_avg_type is not Instant + m_track_avg_cnt = m_track_avg_cnt || m_avg_type!=OutputAvgType::Instant; + } } if (units=="m") { diag_name = "FieldAtHeight"; - EKAT_REQUIRE_MSG(params.isParameter("surface_reference"),"Error! Output field request for " + diag_field_name + " is missing a surface reference." - " Please add either '_above_sealevel' or '_above_surface' to the field name"); + EKAT_REQUIRE_MSG(params.isParameter("surface_reference"),"Error! Output field request for " + diag_field_name + " is missing a surface reference." + " Please add either '_above_sealevel' or '_above_surface' to the field name"); } else if (units=="mb" or units=="Pa" or units=="hPa") { diag_name = "FieldAtPressureLevel"; diag_avg_cnt_name = "_" + tokens[1]; // Set avg_cnt tracking for this specific slice @@ -1390,9 +1391,10 @@ AtmosphereOutput::create_diagnostic (const std::string& diag_field_name) { } // These fields are special case of VerticalLayer diagnostic. - // The diagnostics requires the names be given as param value. - if (diag_name == "z_int" or diag_name == "geopotential_int" or - diag_name == "z_mid" or diag_name == "geopotential_mid" or + // The diagnostics requires the name to be given as param value. + if (diag_name == "z_int" or diag_name == "z_mid" or + diag_name == "geopotential_int" or diag_name == "geopotential_mid" or + diag_name == "height_int" or diag_name == "height_mid" or diag_name == "dz") { params.set("diag_name", diag_name); } @@ -1401,7 +1403,7 @@ AtmosphereOutput::create_diagnostic (const std::string& diag_field_name) { auto diag = diag_factory.create(diag_name,m_comm,params); diag->set_grids(m_grids_manager); - // Add empty entry for this map, so .at(..) always works + // Ensure there's an entry in the map for this diag, so .at(diag_name) always works auto& deps = m_diag_depends_on_diags[diag->name()]; // Initialize the diagnostic diff --git a/components/eamxx/src/share/io/scream_io_utils.hpp b/components/eamxx/src/share/io/scream_io_utils.hpp index 781b376858a..efb2a4fd65b 100644 --- a/components/eamxx/src/share/io/scream_io_utils.hpp +++ b/components/eamxx/src/share/io/scream_io_utils.hpp @@ -80,9 +80,9 @@ struct LongNames { std::map name_2_longname = { {"lev","hybrid level at midpoints (1000*(A+B))"}, {"hyai","hybrid A coefficient at layer interfaces"}, - {"hybi","hybrid B coefficient at layer interfaces"}, - {"hyam","hybrid A coefficient at layer midpoints"}, - {"hybm","hybrid B coefficient at layer midpoints"} + {"hybi","hybrid B coefficient at layer interfaces"}, + {"hyam","hybrid A coefficient at layer midpoints"}, + {"hybm","hybrid B coefficient at layer midpoints"} }; }; diff --git a/components/eamxx/tests/single-process/p3/output.yaml b/components/eamxx/tests/single-process/p3/output.yaml index e52e02c4a19..251adec1f3c 100644 --- a/components/eamxx/tests/single-process/p3/output.yaml +++ b/components/eamxx/tests/single-process/p3/output.yaml @@ -4,7 +4,7 @@ filename_prefix: p3_standalone_output Averaging Type: Instant Field Names: - T_mid - - T_mid_at_2m_above_sealevel + - T_mid_at_2m_above_surface - T_prev_micro_step - qv - qc