From 19f7c0dfe74643b6cedaf36dc55812494fd5998d Mon Sep 17 00:00:00 2001 From: Luca Bertagna Date: Wed, 13 Nov 2024 21:58:03 -0700 Subject: [PATCH 1/4] EAMxx: rework vertical remapper * Allow to build from existing grids * Handle masked as well as P0 extrapolation * Handle top/bot extrapolations separately --- .../share/grid/remap/vertical_remapper.cpp | 526 +++++++++----- .../share/grid/remap/vertical_remapper.hpp | 59 +- .../eamxx/src/share/io/scorpio_output.cpp | 10 +- .../share/tests/vertical_remapper_tests.cpp | 657 ++++++++++++------ 4 files changed, 841 insertions(+), 411 deletions(-) diff --git a/components/eamxx/src/share/grid/remap/vertical_remapper.cpp b/components/eamxx/src/share/grid/remap/vertical_remapper.cpp index 9b524cec5e49..8e13fd4e6df7 100644 --- a/components/eamxx/src/share/grid/remap/vertical_remapper.cpp +++ b/components/eamxx/src/share/grid/remap/vertical_remapper.cpp @@ -17,61 +17,56 @@ namespace scream { +std::shared_ptr VerticalRemapper:: -VerticalRemapper (const grid_ptr_type& src_grid, - const std::string& map_file, - const Field& pmid_src, - const Field& pint_src) - : VerticalRemapper(src_grid,map_file,pmid_src,pint_src,constants::DefaultFillValue::value) +create_tgt_grid (const grid_ptr_type& src_grid, + const std::string& map_file) { - // Nothing to do here -} - -VerticalRemapper:: -VerticalRemapper (const grid_ptr_type& src_grid, - const std::string& map_file, - const Field& pmid_src, - const Field& pint_src, - const Real mask_val) - : AbstractRemapper() - , m_comm (src_grid->get_comm()) - , m_mask_val(mask_val) -{ - using namespace ShortFieldTagsNames; - - // Sanity checks - EKAT_REQUIRE_MSG (src_grid->type()==GridType::Point, - "Error! VerticalRemapper only works on PointGrid grids.\n" - " - src grid name: " + src_grid->name() + "\n" - " - src_grid_type: " + e2str(src_grid->type()) + "\n"); - EKAT_REQUIRE_MSG (src_grid->is_unique(), - "Error! VerticalRemapper requires a unique source grid.\n"); - - // This is a vertical remapper. We only go in one direction - m_bwd_allowed = false; - - // Create tgt_grid that is a clone of the src grid but with - // the correct number of levels. Note that when vertically - // remapping the target field will be defined on the same DOFs - // as the source field, but will have a different number of - // vertical levels. + // Create tgt_grid as a clone of src_grid with different nlevs scorpio::register_file(map_file,scorpio::FileMode::Read); auto nlevs_tgt = scorpio::get_dimlen(map_file,"lev"); auto tgt_grid = src_grid->clone("vertical_remap_tgt_grid",true); tgt_grid->reset_num_vertical_lev(nlevs_tgt); - this->set_grids(src_grid,tgt_grid); - - // Set the LEV and ILEV vertical profiles for interpolation from - set_source_pressure_fields(pmid_src,pint_src); // Gather the pressure level data for vertical remapping - set_pressure_levels(map_file); + auto layout = tgt_grid->get_vertical_layout(true); + Field p_tgt(FieldIdentifier("p_levs",layout,ekat::units::Pa,tgt_grid->name())); + p_tgt.get_header().get_alloc_properties().request_allocation(SCREAM_PACK_SIZE); + p_tgt.allocate_view(); + scorpio::read_var(map_file,"p_levs",p_tgt.get_view().data()); + p_tgt.sync_to_dev(); // Add tgt pressure levels to the tgt grid - tgt_grid->set_geometry_data(m_tgt_pressure); + tgt_grid->set_geometry_data(p_tgt); scorpio::release_file(map_file); + + return tgt_grid; +} + +VerticalRemapper:: +VerticalRemapper (const grid_ptr_type& src_grid, + const std::string& map_file) + : VerticalRemapper(src_grid,create_tgt_grid(src_grid,map_file)) +{ + // NOTE: we prescribe a uniform tgt pressure levels, so pmid_tgt = pint_tgt (1d field) + // Cannot call set_target_pressure(p_tgt,p_tgt), since in there we do check the + // number of levels (i.e., pint/pmid cannot have the same nlevs). Since we remap + // every field (mid or int) to the same pressure coords, we just hard-code them. + m_tgt_pmid = m_tgt_pint = m_tgt_grid->get_geometry_data("p_levs"); +} + +VerticalRemapper:: +VerticalRemapper (const grid_ptr_type& src_grid, + const grid_ptr_type& tgt_grid) +{ + m_bwd_allowed = false; + + EKAT_REQUIRE_MSG (src_grid->get_2d_scalar_layout().congruent(tgt_grid->get_2d_scalar_layout()), + "Error! Source and target grid can only differ for their number of level.\n"); + + this->set_grids (src_grid,tgt_grid); } FieldLayout VerticalRemapper:: @@ -136,32 +131,30 @@ create_layout (const FieldLayout& fl_in, } void VerticalRemapper:: -set_pressure_levels(const std::string& map_file) +set_extrapolation_type (const ExtrapType etype, const TopBot where) { - // Ensure each map file gets a different decomp name - static std::map file2idx; - if (file2idx.find(map_file)==file2idx.end()) { - file2idx[map_file] = file2idx.size(); + if (where & Top) { + m_etype_top = etype; } + if (where & Bot) { + m_etype_bot = etype; + } +} - using namespace ShortFieldTagsNames; - auto layout = m_tgt_grid->get_vertical_layout(true); - FieldIdentifier fid("p_levs",layout,ekat::units::Pa,m_tgt_grid->name()); - m_tgt_pressure = Field(fid); - // Just in case input fields are packed - m_tgt_pressure.get_header().get_alloc_properties().request_allocation(SCREAM_PACK_SIZE); - m_tgt_pressure.allocate_view(); - - auto remap_pres_data = m_tgt_pressure.get_view().data(); - scorpio::read_var(map_file,"p_levs",remap_pres_data); +void VerticalRemapper:: +set_mask_value (const Real mask_val) +{ + EKAT_REQUIRE_MSG (not ekat::is_invalid(mask_val), + "[VerticalRemapper::set_mask_value] Error! Input mask value must be a valid number.\n"); - m_tgt_pressure.sync_to_dev(); + m_mask_val = mask_val; } void VerticalRemapper:: -set_source_pressure_fields(const Field& pmid, const Field& pint) +set_source_pressure (const Field& pmid, const Field& pint) { using namespace ShortFieldTagsNames; + using PackT = ekat::Pack; EKAT_REQUIRE_MSG(pmid.is_allocated(), "Error! Source midpoint pressure field is not yet allocated.\n" @@ -171,23 +164,60 @@ set_source_pressure_fields(const Field& pmid, const Field& pint) "Error! Source interface pressure field is not yet allocated.\n" " - field name: " + pint.name() + "\n"); + EKAT_REQUIRE_MSG(pmid.get_header().get_alloc_properties().is_compatible(), + "Error! Source midpoints pressure field not compatible with default pack size.\n" + " - pack size: " + std::to_string(SCREAM_PACK_SIZE) + "\n"); + EKAT_REQUIRE_MSG(pint.get_header().get_alloc_properties().is_compatible(), + "Error! Source interfaces pressure field not compatible with default pack size.\n" + " - pack size: " + std::to_string(SCREAM_PACK_SIZE) + "\n"); + const auto& pmid_layout = pmid.get_header().get_identifier().get_layout(); const auto& pint_layout = pint.get_header().get_identifier().get_layout(); - EKAT_REQUIRE_MSG(pmid_layout.congruent(m_src_grid->get_3d_scalar_layout(true)), + EKAT_REQUIRE_MSG(pmid_layout.dim(LEV)==m_src_grid->get_num_vertical_levels(), "Error! Source midpoint pressure field has the wrong layout.\n" " - field name: " + pmid.name() + "\n" " - field layout: " + pmid_layout.to_string() + "\n" - " - expected layout: " + m_src_grid->get_3d_scalar_layout(true).to_string() + "\n"); - EKAT_REQUIRE_MSG(pint_layout.congruent(m_src_grid->get_3d_scalar_layout(false)), + " - expected num levels: " + std::to_string(m_src_grid->get_num_vertical_levels()) + "\n"); + EKAT_REQUIRE_MSG(pint_layout.dim(ILEV)==m_src_grid->get_num_vertical_levels()+1, "Error! Source interface pressure field has the wrong layout.\n" " - field name: " + pint.name() + "\n" " - field layout: " + pint_layout.to_string() + "\n" - " - expected layout: " + m_src_grid->get_3d_scalar_layout(false).to_string() + "\n"); + " - expected num levels: " + std::to_string(m_src_grid->get_num_vertical_levels()+1) + "\n"); m_src_pmid = pmid; m_src_pint = pint; } +void VerticalRemapper:: +set_target_pressure (const Field& pmid, const Field& pint) +{ + using namespace ShortFieldTagsNames; + + EKAT_REQUIRE_MSG(pmid.is_allocated(), + "Error! Target midpoint pressure field is not yet allocated.\n" + " - field name: " + pmid.name() + "\n"); + + EKAT_REQUIRE_MSG(pint.is_allocated(), + "Error! Target interface pressure field is not yet allocated.\n" + " - field name: " + pint.name() + "\n"); + + const auto& pmid_layout = pmid.get_header().get_identifier().get_layout(); + const auto& pint_layout = pint.get_header().get_identifier().get_layout(); + EKAT_REQUIRE_MSG(pmid_layout.dim(LEV)==m_tgt_grid->get_num_vertical_levels(), + "Error! Target midpoint pressure field has the wrong layout.\n" + " - field name: " + pmid.name() + "\n" + " - field layout: " + pmid_layout.to_string() + "\n" + " - expected num levels: " + std::to_string(m_tgt_grid->get_num_vertical_levels()) + "\n"); + EKAT_REQUIRE_MSG(pint_layout.dim(ILEV)==m_tgt_grid->get_num_vertical_levels()+1, + "Error! Target interface pressure field has the wrong layout.\n" + " - field name: " + pint.name() + "\n" + " - field layout: " + pint_layout.to_string() + "\n" + " - expected num levels: " + std::to_string(m_tgt_grid->get_num_vertical_levels()+1) + "\n"); + + m_tgt_pmid = pmid; + m_tgt_pint = pint; +} + void VerticalRemapper:: do_register_field (const identifier_type& src, const identifier_type& tgt) { @@ -198,7 +228,7 @@ do_register_field (const identifier_type& src, const identifier_type& tgt) // could have src with ILEV and tgt with LEV) auto src_layout = src.get_layout().clone(); auto tgt_layout = tgt.get_layout().clone(); - EKAT_REQUIRE_MSG(src_layout.strip_dims({ILEV,LEV}).congruent(tgt_layout.strip_dims({LEV})), + EKAT_REQUIRE_MSG(src_layout.strip_dims({ILEV,LEV}).congruent(tgt_layout.strip_dims({LEV,ILEV})), "[VerticalRemapper] Error! Once vertical level tag is stripped, src/tgt layouts are incompatible.\n" " - src field name: " + src.name() + "\n" " - tgt field name: " + tgt.name() + "\n" @@ -231,54 +261,56 @@ do_bind_field (const int ifield, const field_type& src, const field_type& tgt) ft.packed = src.get_header().get_alloc_properties().is_compatible() and tgt.get_header().get_alloc_properties().is_compatible(); - // NOTE: for now we assume that masking is determined only by the COL,LEV location in space - // and that fields with multiple components will have the same masking for each component - // at a specific COL,LEV - src_layout.strip_dims({CMP}); + if (m_etype_top==Mask or m_etype_bot==Mask) { + // NOTE: for now we assume that masking is determined only by the COL,LEV location in space + // and that fields with multiple components will have the same masking for each component + // at a specific COL,LEV + src_layout.strip_dims({CMP}); - // I this mask has already been created, retrieve it, otherwise create it - const auto mask_name = m_tgt_grid->name() + "_" + ekat::join(src_layout.names(),"_") + "_mask"; - Field tgt_mask; - if (m_field2type.count(mask_name)==0) { - auto nondim = ekat::units::Units::nondimensional(); - // Create this src/tgt mask fields, and assign them to these src/tgt fields extra data + // I this mask has already been created, retrieve it, otherwise create it + const auto mask_name = m_tgt_grid->name() + "_" + ekat::join(src_layout.names(),"_") + "_mask"; + Field tgt_mask; + if (m_field2type.count(mask_name)==0) { + auto nondim = ekat::units::Units::nondimensional(); + // Create this src/tgt mask fields, and assign them to these src/tgt fields extra data - FieldIdentifier src_mask_fid (mask_name, src_layout, nondim, m_src_grid->name() ); - FieldIdentifier tgt_mask_fid = create_tgt_fid(src_mask_fid); + FieldIdentifier src_mask_fid (mask_name, src_layout, nondim, m_src_grid->name() ); + FieldIdentifier tgt_mask_fid = create_tgt_fid(src_mask_fid); - Field src_mask (src_mask_fid); - src_mask.allocate_view(); + Field src_mask (src_mask_fid); + src_mask.allocate_view(); - tgt_mask = Field (tgt_mask_fid); - tgt_mask.allocate_view(); + tgt_mask = Field (tgt_mask_fid); + tgt_mask.allocate_view(); - // Initialize the src mask values to 1.0 - src_mask.deep_copy(1.0); + // Initialize the src mask values to 1.0 + src_mask.deep_copy(1.0); - m_src_masks.push_back(src_mask); - m_tgt_masks.push_back(tgt_mask); + m_src_masks.push_back(src_mask); + m_tgt_masks.push_back(tgt_mask); - auto& mt = m_field2type[src_mask_fid.name()]; - mt.packed = false; - mt.midpoints = src_layout.has_tag(LEV); - } else { - for (size_t i=0; i void VerticalRemapper:: setup_lin_interp (const ekat::LinInterp& lin_interp, - const Field& p_src) const + const Field& p_src, const Field& p_tgt) const { using LI_t = ekat::LinInterp; using ESU = ekat::ExeSpaceUtils; using PackT = ekat::Pack; - auto p_src_v = p_src.get_view(); - auto p_tgt_v = m_tgt_pressure.get_view(); + using view2d = typename KokkosTypes::view; + using view1d = typename KokkosTypes::view; + + auto src1d = p_src.rank()==1; + auto tgt1d = p_tgt.rank()==1; + + view2d p_src2d_v, p_tgt2d_v; + view1d p_src1d_v, p_tgt1d_v; + if (src1d) { + p_src1d_v = p_src.get_view(); + } else { + p_src2d_v = p_src.get_view(); + } + if (tgt1d) { + p_tgt1d_v = p_tgt.get_view(); + } else { + p_tgt2d_v = p_tgt.get_view(); + } auto lambda = KOKKOS_LAMBDA(typename LI_t::MemberType const& team) { const int icol = team.league_rank(); - lin_interp.setup(team,ekat::subview(p_src_v,icol), - p_tgt_v); + // Extract subviews if src/tgt were not 1d to start with + auto x_src = p_src1d_v; + if (not src1d) + x_src = ekat::subview(p_src2d_v,icol); + auto x_tgt = p_tgt1d_v; + if (not tgt1d) + x_tgt = ekat::subview(p_tgt2d_v,icol); + + lin_interp.setup(team,x_src,x_tgt); }; - const int ncols = m_src_grid->get_num_local_dofs(); const int nlevs_tgt = m_tgt_grid->get_num_vertical_levels(); const int npacks_tgt = ekat::PackInfo::num_packs(nlevs_tgt); @@ -465,8 +523,7 @@ template void VerticalRemapper:: apply_vertical_interpolation(const ekat::LinInterp& lin_interp, const Field& f_src, const Field& f_tgt, - const Field& p_src, - const Real mask_val) const + const Field& p_src, const Field& p_tgt) const { // Note: if Packsize==1, we grab packs of size 1, which are for sure // compatible with the allocation @@ -474,44 +531,51 @@ apply_vertical_interpolation(const ekat::LinInterp& lin_interp, using PackT = ekat::Pack; using ESU = ekat::ExeSpaceUtils; - auto p_src_v = p_src.get_view(); - auto x_tgt = m_tgt_pressure.get_view(); - const auto& f_src_l = f_src.get_header().get_identifier().get_layout(); + using view2d = typename KokkosTypes::view; + using view1d = typename KokkosTypes::view; + + auto src1d = p_src.rank()==1; + auto tgt1d = p_tgt.rank()==1; + + view2d p_src2d_v, p_tgt2d_v; + view1d p_src1d_v, p_tgt1d_v; + if (src1d) { + p_src1d_v = p_src.get_view(); + } else { + p_src2d_v = p_src.get_view(); + } + if (tgt1d) { + p_tgt1d_v = p_tgt.get_view(); + } else { + p_tgt2d_v = p_tgt.get_view(); + } + + const auto& f_tgt_l = f_tgt.get_header().get_identifier().get_layout(); const int ncols = m_src_grid->get_num_local_dofs(); - const int nlevs_tgt = m_tgt_grid->get_num_vertical_levels(); - const int nlevs_src = f_src_l.dims().back(); + const int nlevs_tgt = f_tgt_l.dims().back(); const int npacks_tgt = ekat::PackInfo::num_packs(nlevs_tgt); - const int last_src_pack_idx = ekat::PackInfo::last_pack_idx(nlevs_src); - const int last_src_pack_end = ekat::PackInfo::last_vec_end(nlevs_src); - switch(f_src.rank()) { case 2: { auto f_src_v = f_src.get_view(); auto f_tgt_v = f_tgt.get_view< PackT**>(); auto policy = ESU::get_default_team_policy(ncols,npacks_tgt); - auto lambda = KOKKOS_LAMBDA(typename LI_t::MemberType const& team) { - - // Interpolate + auto lambda = KOKKOS_LAMBDA(typename LI_t::MemberType const& team) + { const int icol = team.league_rank(); - auto x_src = ekat::subview(p_src_v,icol); + + // Extract subviews if src/tgt pressures were not 1d to start with + auto x_src = p_src1d_v; + auto x_tgt = p_tgt1d_v; + if (not src1d) + x_src = ekat::subview(p_src2d_v,icol); + if (not tgt1d) + x_tgt = ekat::subview(p_tgt2d_v,icol); + auto y_src = ekat::subview(f_src_v,icol); auto y_tgt = ekat::subview(f_tgt_v,icol); lin_interp.lin_interp(team,x_src,x_tgt,y_src,y_tgt,icol); - team.team_barrier(); - - // If x_tgt is extrapolated, set to mask_val - auto x_min = x_src[0][0]; - auto x_max = x_src[last_src_pack_idx][last_src_pack_end-1]; - auto set_mask = [&](const int ipack) { - auto in_range = ekat::range(ipack*Packsize) < nlevs_tgt; - auto oob = (x_tgt[ipack]x_max) and in_range; - if (oob.any()) { - y_tgt[ipack].set(oob,mask_val); - } - }; - Kokkos::parallel_for (Kokkos::TeamThreadRange(team,npacks_tgt), set_mask); }; Kokkos::parallel_for("VerticalRemapper::apply_vertical_interpolation",policy,lambda); break; @@ -520,31 +584,25 @@ apply_vertical_interpolation(const ekat::LinInterp& lin_interp, { auto f_src_v = f_src.get_view(); auto f_tgt_v = f_tgt.get_view< PackT***>(); - const auto& layout = f_src.get_header().get_identifier().get_layout(); - const int ncomps = layout.get_vector_dim(); + const int ncomps = f_tgt_l.get_vector_dim(); auto policy = ESU::get_default_team_policy(ncols*ncomps,npacks_tgt); auto lambda = KOKKOS_LAMBDA(typename LI_t::MemberType const& team) { - // Interpolate const int icol = team.league_rank() / ncomps; const int icmp = team.league_rank() % ncomps; - auto x_src = ekat::subview(p_src_v,icol); + + // Extract subviews if src/tgt pressures were not 1d to start with + auto x_src = p_src1d_v; + auto x_tgt = p_tgt1d_v; + if (not src1d) + x_src = ekat::subview(p_src2d_v,icol); + if (not tgt1d) + x_tgt = ekat::subview(p_tgt2d_v,icol); + auto y_src = ekat::subview(f_src_v,icol,icmp); auto y_tgt = ekat::subview(f_tgt_v,icol,icmp); lin_interp.lin_interp(team,x_src,x_tgt,y_src,y_tgt,icol); - team.team_barrier(); - - // If x_tgt is extrapolated, set to mask_val - auto x_min = x_src[0][0]; - auto x_max = x_src[last_src_pack_idx][last_src_pack_end-1]; - auto set_mask = [&](const int ipack) { - auto oob = x_tgt[ipack]x_max; - if (oob.any()) { - y_tgt[ipack].set(oob,mask_val); - } - }; - Kokkos::parallel_for (Kokkos::TeamThreadRange(team,npacks_tgt), set_mask); }; Kokkos::parallel_for("VerticalRemapper::apply_vertical_interpolation",policy,lambda); break; @@ -557,4 +615,148 @@ apply_vertical_interpolation(const ekat::LinInterp& lin_interp, } } +void VerticalRemapper:: +extrapolate (const Field& f_src, + const Field& f_tgt, + const Field& p_src, + const Field& p_tgt, + const Real mask_val) const +{ + using ESU = ekat::ExeSpaceUtils; + + using view2d = typename KokkosTypes::view; + using view1d = typename KokkosTypes::view; + + auto src1d = p_src.rank()==1; + auto tgt1d = p_tgt.rank()==1; + + view2d p_src2d_v, p_tgt2d_v; + view1d p_src1d_v, p_tgt1d_v; + if (src1d) { + p_src1d_v = p_src.get_view(); + } else { + p_src2d_v = p_src.get_view(); + } + if (tgt1d) { + p_tgt1d_v = p_tgt.get_view(); + } else { + p_tgt2d_v = p_tgt.get_view(); + } + + const auto& f_tgt_l = f_tgt.get_header().get_identifier().get_layout(); + const auto& f_src_l = f_src.get_header().get_identifier().get_layout(); + const int ncols = m_src_grid->get_num_local_dofs(); + const int nlevs_tgt = f_tgt_l.dims().back(); + const int nlevs_src = f_src_l.dims().back(); + + auto etop = m_etype_top; + auto ebot = m_etype_bot; + auto mid = nlevs_tgt / 2; + switch(f_src.rank()) { + case 2: + { + auto f_src_v = f_src.get_view(); + auto f_tgt_v = f_tgt.get_view< Real**>(); + auto policy = ESU::get_default_team_policy(ncols,nlevs_tgt); + auto lambda = KOKKOS_LAMBDA(const auto& team) + { + const int icol = team.league_rank(); + + // Extract subviews if src/tgt pressures were not 1d to start with + auto x_src = p_src1d_v; + auto x_tgt = p_tgt1d_v; + if (not src1d) + x_src = ekat::subview(p_src2d_v,icol); + if (not tgt1d) + x_tgt = ekat::subview(p_tgt2d_v,icol); + + auto y_src = ekat::subview(f_src_v,icol); + auto y_tgt = ekat::subview(f_tgt_v,icol); + + auto x_min = x_src[0]; + auto x_max = x_src[nlevs_src-1]; + auto extrapolate = [&](const int ilev) { + if (ilev>=mid) { + // Near surface + if (x_tgt[ilev]>x_max) { + if (ebot==P0) { + y_tgt[ilev] = y_src[nlevs_src-1]; + } else { + y_tgt[ilev] = mask_val; + } + } + } else { + // Near top + if (x_tgt[ilev](); + auto f_tgt_v = f_tgt.get_view< Real***>(); + const int ncomps = f_tgt_l.get_vector_dim(); + auto policy = ESU::get_default_team_policy(ncols*ncomps,nlevs_tgt); + + auto lambda = KOKKOS_LAMBDA(const auto& team) + { + const int icol = team.league_rank() / ncomps; + const int icmp = team.league_rank() % ncomps; + + // Extract subviews if src/tgt pressures were not 1d to start with + auto x_src = p_src1d_v; + auto x_tgt = p_tgt1d_v; + if (not src1d) + x_src = ekat::subview(p_src2d_v,icol); + if (not tgt1d) + x_tgt = ekat::subview(p_tgt2d_v,icol); + + auto y_src = ekat::subview(f_src_v,icol,icmp); + auto y_tgt = ekat::subview(f_tgt_v,icol,icmp); + auto x_min = x_src[0]; + auto x_max = x_src[nlevs_src-1]; + auto extrapolate = [&](const int ilev) { + if (ilev>=mid) { + // Near surface + if (x_tgt[ilev]>x_max) { + if (ebot==P0) { + y_tgt[ilev] = y_src[nlevs_src-1]; + } else { + y_tgt[ilev] = mask_val; + } + } + } else { + // Near top + if (x_tgt[ilev] + create_tgt_grid (const grid_ptr_type& src_grid, + const std::string& map_file); + protected: FieldLayout create_layout (const FieldLayout& fl_in, @@ -89,25 +105,23 @@ class VerticalRemapper : public AbstractRemapper EKAT_ERROR_MSG ("VerticalRemapper only supports fwd remapping.\n"); } - void set_pressure_levels (const std::string& map_file); - void do_print(); - #ifdef KOKKOS_ENABLE_CUDA public: #endif template void apply_vertical_interpolation (const ekat::LinInterp& lin_interp, const Field& f_src, const Field& f_tgt, - const Field& p_src, - const Real mask_value) const; + const Field& p_src, const Field& p_tgt) const; + void extrapolate (const Field& f_src, const Field& f_tgt, + const Field& p_src, const Field& p_tgt, + const Real mask_val) const; template void setup_lin_interp (const ekat::LinInterp& lin_interp, - const Field& p_src) const; + const Field& p_src, const Field& p_tgt) const; protected: - void set_source_pressure_fields(const Field& pmid, const Field& pint); void create_lin_interp (); using KT = KokkosTypes; @@ -122,14 +136,19 @@ class VerticalRemapper : public AbstractRemapper // Source and target fields std::vector m_src_fields; std::vector m_tgt_fields; - std::vector m_tgt_masks; std::vector m_src_masks; + std::vector m_tgt_masks; // Vertical profile fields, both for source and target - Real m_mask_val; - Field m_tgt_pressure; - Field m_src_pmid; // Src vertical profile for LEV layouts - Field m_src_pint; // Src vertical profile for ILEV layouts + Field m_src_pmid; + Field m_src_pint; + Field m_tgt_pmid; + Field m_tgt_pint; + + // Extrapolation settings at top/bottom. Default to P0 extrapolation + ExtrapType m_etype_top = P0; + ExtrapType m_etype_bot = P0; + Real m_mask_val = std::numeric_limits::quiet_NaN(); // We need to remap mid/int fields separately, and we want to use packs if possible, // so we need to divide input fields into 4 separate categories diff --git a/components/eamxx/src/share/io/scorpio_output.cpp b/components/eamxx/src/share/io/scorpio_output.cpp index e8f322b85f32..47a5616aec55 100644 --- a/components/eamxx/src/share/io/scorpio_output.cpp +++ b/components/eamxx/src/share/io/scorpio_output.cpp @@ -227,9 +227,13 @@ AtmosphereOutput (const ekat::Comm& comm, const ekat::ParameterList& params, if (use_vertical_remap_from_file) { // We build a remapper, to remap fields from the fm grid to the io grid auto vert_remap_file = params.get("vertical_remap_file"); - auto f_lev = get_field("p_mid","sim"); - auto f_ilev = get_field("p_int","sim"); - m_vert_remapper = std::make_shared(io_grid,vert_remap_file,f_lev,f_ilev,m_fill_value); + auto p_mid = get_field("p_mid","sim"); + auto p_int = get_field("p_int","sim"); + auto vert_remapper = std::make_shared(io_grid,vert_remap_file); + vert_remapper->set_source_pressure (p_mid,p_int); + vert_remapper->set_mask_value(m_fill_value); + vert_remapper->set_extrapolation_type(VerticalRemapper::Mask); // both Top AND Bot + m_vert_remapper = vert_remapper; io_grid = m_vert_remapper->get_tgt_grid(); set_grid(io_grid); diff --git a/components/eamxx/src/share/tests/vertical_remapper_tests.cpp b/components/eamxx/src/share/tests/vertical_remapper_tests.cpp index 34351c65db8f..b4f1627a16ec 100644 --- a/components/eamxx/src/share/tests/vertical_remapper_tests.cpp +++ b/components/eamxx/src/share/tests/vertical_remapper_tests.cpp @@ -4,37 +4,46 @@ #include "share/grid/remap/coarsening_remapper.hpp" #include "share/grid/point_grid.hpp" #include "share/io/scream_scorpio_interface.hpp" +#include "share/util/scream_timing.hpp" +#include "share/field/field_utils.hpp" namespace scream { -template -typename ViewT::HostMirror -cmvc (const ViewT& v) { - auto vh = Kokkos::create_mirror_view(v); - Kokkos::deep_copy(vh,v); - return vh; -} +constexpr int vec_dim = 3; +constexpr auto P0 = VerticalRemapper::P0; +constexpr auto Mask = VerticalRemapper::Mask; +constexpr auto Top = VerticalRemapper::Top; +constexpr auto Bot = VerticalRemapper::Bot; +constexpr auto TopBot = VerticalRemapper::TopAndBot; +constexpr Real mask_val = -99999.0; -void print (const std::string& msg, const ekat::Comm& comm) { +template +void print (const std::string& fmt, const ekat::Comm& comm, Args&&... args) { + if (comm.am_i_root()) { + printf(fmt.c_str(),std::forward(args)...); + } +} +// Overload for when there are no additional arguments +void print(const std::string& fmt, const ekat::Comm& comm) { if (comm.am_i_root()) { - printf("%s",msg.c_str()); + printf(fmt.c_str()); } } // Helper function to create a grid given the number of dof's and a comm group. std::shared_ptr -build_src_grid(const ekat::Comm& comm, const int nldofs_src, const int nlevs_src) +build_grid(const ekat::Comm& comm, const int nldofs, const int nlevs) { using gid_type = AbstractGrid::gid_type; - auto src_grid = std::make_shared("src",nldofs_src,nlevs_src,comm); + auto grid = std::make_shared("src",nldofs,nlevs,comm); - auto src_dofs = src_grid->get_dofs_gids(); - auto src_dofs_h = src_dofs.get_view(); - std::iota(src_dofs_h.data(),src_dofs_h.data()+nldofs_src,nldofs_src*comm.rank()); - src_dofs.sync_to_dev(); + auto dofs = grid->get_dofs_gids(); + auto dofs_h = dofs.get_view(); + std::iota(dofs_h.data(),dofs_h.data()+nldofs,nldofs*comm.rank()); + dofs.sync_to_dev(); - return src_grid; + return grid; } // Helper function to create fields @@ -64,7 +73,133 @@ Real data_func(const int col, const int vec, const Real pres) { // - pres, the current pressure // Should ensure that the interpolated values match exactly, since vertical interp is also a linear interpolator. // Note, we don't use the level, because here the vertical interpolation is over pressure, so it represents the level. - return col*pres + vec*100.0; + return (col+1)*pres + vec*100.0; +} + +void compute_field (const Field& f, const Field& p) +{ + Field::view_host_t p1d; + Field::view_host_t p2d; + bool rank1 = p.rank()==1; + const auto& l = f.get_header().get_identifier().get_layout(); + const int ncols = l.dims().front(); + const int nlevs = l.dims().back(); + if (rank1) { + p1d = p.get_view(); + } else { + p2d = p.get_view(); + } + + // Grab correct pressure (1d or 2d) + auto pval = [&](int i, int k) { + if (rank1) return p1d(k); + else return p2d(i,k); + }; + + switch (l.type()) { + case LayoutType::Scalar2D: + { + const auto v = f.get_view(); + for (int i=0; i(); + for (int i=0; i(); + for (int i=0; i(); + for (int i=0; i p1d_src,p1d_tgt; + Field::view_host_t p2d_src,p2d_tgt; + if (p_src.rank()==1) { + p1d_src = p_src.get_view(); + } else { + p2d_src = p_src.get_view(); + } + if (p_tgt.rank()==1) { + p1d_tgt = p_tgt.get_view(); + } else { + p2d_tgt = p_tgt.get_view(); + } + + auto pval = [&](auto p1d, auto p2d, int i, int k, int rank) { + if (rank==1) return p1d(k); + else return p2d(i,k); + }; + + const auto& l = f.get_header().get_identifier().get_layout(); + const int ncols = l.dims().front(); + const int nlevs = l.dims().back(); + const int nlevs_src = p_src.get_header().get_identifier().get_layout().dims().back(); + // print_field_hyperslab(p_src); + switch (l.type()) { + case LayoutType::Scalar2D: break; + case LayoutType::Vector2D: break; + case LayoutType::Scalar3D: + { + const auto v = f.get_view(); + for (int i=0; ipmax) { + v(i,j) = etype_bot==Mask ? mask_val : data_func(i,0,pmax); + } else if (p(); + for (int i=0; ipmax) { + v(i,j,k) = etype_bot==Mask ? mask_val : data_func(i,j,pmax); + } else if (p dofs_p(nlevs_tgt); std::iota(dofs_p.begin(),dofs_p.end(),0); std::vector p_tgt; - for (int ii=0; ii creating map file ... done!\n",comm); // -------------------------------------- // - // Build src grid and remapper // + // Build src grid // // -------------------------------------- // - print (" -> creating grid and remapper ...\n",comm); + print (" -> creating src grid ...\n",comm); + auto src_grid = build_grid(comm, nldofs, nlevs_src); + print (" -> creating src grid ... done!\n",comm); - const Real mask_val = -99999.0; + // -------------------------------------- // + // Retrieve tgt grid // + // -------------------------------------- // - auto src_grid = build_src_grid(comm, nldofs_src, nlevs_src); - - // We need the source pressure level fields for both p_mid and p_int - auto pmid_src = create_field("p_mid", src_grid, false, false, true, SCREAM_PACK_SIZE); - auto pint_src = create_field("p_int", src_grid, false, false, false, SCREAM_PACK_SIZE); - // Set the source pressures - { - // By adding 1 to the pbot_tgt and subtrating 1 from ptop_tgt we ensure some masking, which - // we also want to check. - const Real ptop_src = ptop_tgt-1; - const Real pbot_src = pbot_tgt+1; - const Real dp_src = (pbot_src-ptop_src)/(nlevs_src-1); - auto pmid_v = pmid_src.get_view(); - auto pint_v = pint_src.get_view(); - for (int ii=0; ii(src_grid,filename,pmid_src,pint_src,mask_val); - print (" -> creating grid and remapper ... done!\n",comm); + print (" -> retreiving tgt grid ...\n",comm); + auto tgt_grid = VerticalRemapper::create_tgt_grid (src_grid,filename); + print (" -> retreiving tgt grid ... done!\n",comm); // -------------------------------------- // - // Create src/tgt grid fields // + // Check tgt grid // // -------------------------------------- // - print (" -> creating fields ...\n",comm); - constexpr int vec_dim = 3; + print (" -> checking tgt grid ...\n",comm); + REQUIRE (tgt_grid->get_num_local_dofs()==src_grid->get_num_local_dofs()); + REQUIRE (tgt_grid->get_num_vertical_levels()==nlevs_tgt); + REQUIRE (tgt_grid->has_geometry_data("p_levs")); + auto p_levs = tgt_grid->get_geometry_data("p_levs"); + auto p_levs_v = p_levs.get_view(); + for (int k=0; kget_tgt_grid(); - // Check that the target grid made by the remapper has the same number of columns as the source grid. - // Also check that the number of levels matches the expectation. - REQUIRE(tgt_grid->get_num_vertical_levels()==nlevs_tgt); - REQUIRE(tgt_grid->get_num_global_dofs()==src_grid->get_num_global_dofs()); - - auto src_s2d = create_field("s2d", src_grid,true,false); - auto src_v2d = create_field("v2d", src_grid,true,true); - // For now we can only support PS = SCREAM_PACK_SIZE because the source and target pressure levels will assume that packsize. - // If we use a smaller packsize we throw an error in the property check step of the vertical interpolation scheme. - // TODO: Fix that. - auto src_s3d_m = create_field("s3d_m",src_grid,false,false,true, 1); - auto src_s3d_i = create_field("s3d_i",src_grid,false,false,false,SCREAM_PACK_SIZE); - auto src_v3d_m = create_field("v3d_m",src_grid,false,true ,true, 1); - auto src_v3d_i = create_field("v3d_i",src_grid,false,true ,false,SCREAM_PACK_SIZE); - - auto tgt_s2d = create_field("s2d", tgt_grid,true,false); - auto tgt_v2d = create_field("v2d", tgt_grid,true,true); - auto tgt_s3d_m = create_field("s3d_m",tgt_grid,false,false,true, 1); - auto tgt_s3d_i = create_field("s3d_i",tgt_grid,false,false,true, SCREAM_PACK_SIZE); - auto tgt_v3d_m = create_field("v3d_m",tgt_grid,false,true ,true, 1); - auto tgt_v3d_i = create_field("v3d_i",tgt_grid,false,true ,true, SCREAM_PACK_SIZE); - - std::vector src_f = {src_s2d,src_v2d,src_s3d_m,src_s3d_i,src_v3d_m,src_v3d_i}; - std::vector tgt_f = {tgt_s2d,tgt_v2d,tgt_s3d_m,tgt_s3d_i,tgt_v3d_m,tgt_v3d_i}; - - print (" -> creating fields ... done!\n",comm); + print (" -> checking tgt grid ... done!\n",comm); - // -------------------------------------- // - // Register fields in the remapper // - // -------------------------------------- // + // Clean up scorpio stuff + scorpio::finalize_subsystem(); - print (" -> registering fields ...\n",comm); - remap->registration_begins(); - remap->register_field(src_s2d, tgt_s2d); - remap->register_field(src_v2d, tgt_v2d); - remap->register_field(src_s3d_m,tgt_s3d_m); - remap->register_field(src_s3d_i,tgt_s3d_i); - remap->register_field(src_v3d_m,tgt_v3d_m); - remap->register_field(src_v3d_i,tgt_v3d_i); - remap->registration_ends(); - print (" -> registering fields ... done!\n",comm); + print ("Testing retrieval of tgt grid from map file ...\n",comm); +} +TEST_CASE ("vertical_remapper") { // -------------------------------------- // - // Check remapper internals // + // Init MPI and PIO // // -------------------------------------- // - print (" -> Checking remapper internal state ...\n",comm); + ekat::Comm comm(MPI_COMM_WORLD); + print ("Testing vertical remapper ...\n",comm); - // Check the pressure levels read from map file + scorpio::init_subsystem(comm); // -------------------------------------- // - // Generate data for src fields // + // Set grid/map sizes // // -------------------------------------- // - print (" -> generate src fields data ...\n",comm); - using namespace ShortFieldTagsNames; - // Generate data in a deterministic way, so that when we check results, - // we know a priori what the input data that generated the tgt field's - // values was, even if that data was off rank. - auto pmid_v = pmid_src.get_view(); - auto pint_v = pint_src.get_view(); - auto src_gids = remap->get_src_grid()->get_dofs_gids().get_view(); - for (const auto& f : src_f) { - const auto& l = f.get_header().get_identifier().get_layout(); - switch (l.type()) { - case LayoutType::Scalar2D: - { - const auto v_src = f.get_view(); - for (int i=0; i creating src grid ...\n",comm); + auto src_grid = build_grid(comm, nldofs, nlevs_src); + print (" nlevs src: %d\n",comm,nlevs_src); + print (" -> creating src grid ...done!\n",comm); + + // Tgt grid must have same 2d layout as src grid + REQUIRE_THROWS (std::make_shared(src_grid,build_grid(comm,nldofs+1,nlevs_src))); + + // Helper lambda, to create p_int profile. If it is a 3d field, make same profile on each col + auto create_pint = [&](const auto& grid, const bool one_d, const Real ptop, const Real pbot) { + auto layout = one_d ? grid->get_vertical_layout(false) + : grid->get_3d_scalar_layout(false); + FieldIdentifier fid("p_int",layout,ekat::units::Pa,grid->name()); + Field pint (fid); + pint.get_header().get_alloc_properties().request_allocation(SCREAM_PACK_SIZE); + pint.allocate_view(); + + int nlevs = grid->get_num_vertical_levels(); + const Real dp = (pbot-ptop)/nlevs; + + if (one_d) { + auto pv = pint.get_view(); + pv(nlevs) = pbot; + for (int k=nlevs; k>0; --k) { + pv(k-1) = pv(k) - dp; + } + } else { + auto pv = pint.get_view(); + for (int i=0; i0; --k) { + pv(i,k-1) = pv(i,k) - dp; } - } break; - case LayoutType::Vector2D: - { - const auto v_src = f.get_view(); - for (int i=0; i(); - for (int i=0; i(); - for (int i=0; i generate src fields data ... done!\n",comm); - - // No bwd remap - REQUIRE_THROWS(remap->remap(false)); - - for (int irun=0; irun<5; ++irun) { - print (" -> run remap ...\n",comm); - remap->remap(true); - print (" -> run remap ... done!\n",comm); - - // -------------------------------------- // - // Check remapped fields // - // -------------------------------------- // - - print (" -> check tgt fields ...\n",comm); - const auto tgt_gids = tgt_grid->get_dofs_gids().get_view(); - const int ntgt_gids = tgt_gids.size(); - for (size_t ifield=0; ifield Checking field with source layout " + ls +" " + dots + "\n",comm); - - f.sync_to_host(); - - switch (lsrc.type()) { - case LayoutType::Scalar2D: - { - // This is a flat array w/ no LEV tag so the interpolated value for source and target should match. - const auto v_src = fsrc.get_view(); - const auto v_tgt = f.get_view(); - for (int i=0; i(); - const auto v_src = fsrc.get_view(); - for (int i=0; i(); - for (int i=0; ip_v(i,nlevs_p-1) || p_tgt[j](); - for (int i=0; ip_v(i,nlevs_p-1) || p_tgt[k](); + auto pmid_v = pmid.get_view< Real*,Host>(); + for (int k=0; k(); + auto pmid_v = pmid.get_view< Real**,Host>(); + for (int i=0; i Checking field with source layout " + ls + " " + dots + " OK!\n",comm); + // Test tgt grid with 2x and 0.5x as many levels as src grid + for (int nlevs_tgt : {nlevs_src/2, 2*nlevs_src}) { + for (bool src_1d : {true, false}) { + for (bool tgt_1d : {true, false}) { + for (auto etype_top : {P0, Mask}) { + for (auto etype_bot : {P0, Mask}) { + print ("************************************************\n",comm); + print (" nlevs tgt: %d\n",comm,nlevs_tgt); + print (" src pressure is 1d: %s\n",comm,src_1d ? "true" : "false"); + print (" tgt pressure is 1d: %s\n",comm,tgt_1d ? "true" : "false"); + print (" extrap type at top: %s\n",comm,etype_top==P0 ? "p0" : "masked"); + print (" extrap type at bot: %s\n",comm,etype_bot==P0 ? "p0" : "masked"); + print ("************************************************\n",comm); + + print (" -> creating tgt grid ...\n",comm); + auto tgt_grid = src_grid->clone("tgt",true); + tgt_grid->reset_num_vertical_lev(nlevs_tgt); + print (" -> creating tgt grid ...done!\n",comm); + + print (" -> creating src/tgt pressure fields ...\n",comm); + auto pint_src = create_pint(src_grid, src_1d, ptop_src, pbot_src); + auto pmid_src = create_pmid(pint_src); + + // Make ptop_tgtpsurf_src, so we do have extrapolation + const Real ptop_tgt = 10; + const Real pbot_tgt = 1020; + auto pint_tgt = create_pint(tgt_grid, tgt_1d, ptop_tgt, pbot_tgt); + auto pmid_tgt = create_pmid(pint_tgt); + print (" -> creating src/tgt pressure fields ... done!\n",comm); + + print (" -> creating fields ... done!\n",comm); + auto src_s2d = create_field("s2d", src_grid,true,false); + auto src_v2d = create_field("v2d", src_grid,true,true); + auto src_s3d_m = create_field("s3d_m",src_grid,false,false,true, 1); + auto src_s3d_i = create_field("s3d_i",src_grid,false,false,false,SCREAM_PACK_SIZE); + auto src_v3d_m = create_field("v3d_m",src_grid,false,true ,true, 1); + auto src_v3d_i = create_field("v3d_i",src_grid,false,true ,false,SCREAM_PACK_SIZE); + + auto tgt_s2d = create_field("s2d", tgt_grid,true,false); + auto tgt_v2d = create_field("v2d", tgt_grid,true,true); + auto tgt_s3d_m = create_field("s3d_m",tgt_grid,false,false,true, 1); + auto tgt_s3d_i = create_field("s3d_i",tgt_grid,false,false,true, SCREAM_PACK_SIZE); + auto tgt_v3d_m = create_field("v3d_m",tgt_grid,false,true ,true, 1); + auto tgt_v3d_i = create_field("v3d_i",tgt_grid,false,true ,true, SCREAM_PACK_SIZE); + + auto expected_s2d = tgt_s2d.clone(); + auto expected_v2d = tgt_v2d.clone(); + auto expected_s3d_m = tgt_s3d_m.clone(); + auto expected_s3d_i = tgt_s3d_i.clone(); + auto expected_v3d_m = tgt_v3d_m.clone(); + auto expected_v3d_i = tgt_v3d_i.clone(); + print (" -> creating fields ... done!\n",comm); + + // -------------------------------------- // + // Register fields in the remapper // + // -------------------------------------- // + + print (" -> creating and initializing remapper ...\n",comm); + auto remap = std::make_shared(src_grid,tgt_grid); + remap->set_source_pressure (pmid_src, pint_src); + remap->set_target_pressure (pmid_tgt, pint_tgt); + remap->set_extrapolation_type(etype_top,Top); + remap->set_extrapolation_type(etype_bot,Bot); + REQUIRE_THROWS (remap->set_mask_value(std::numeric_limits::quiet_NaN())); + remap->set_mask_value(mask_val); // Only needed if top and/or bot use etype=Mask + + remap->registration_begins(); + remap->register_field(src_s2d, tgt_s2d); + remap->register_field(src_v2d, tgt_v2d); + remap->register_field(src_s3d_m,tgt_s3d_m); + remap->register_field(src_s3d_i,tgt_s3d_i); + remap->register_field(src_v3d_m,tgt_v3d_m); + remap->register_field(src_v3d_i,tgt_v3d_i); + remap->registration_ends(); + print (" -> creating and initializing remapper ... done!\n",comm); + + // -------------------------------------- // + // Generate data for src fields // + // -------------------------------------- // + + print (" -> generate fields data ...\n",comm); + compute_field(src_s2d, pmid_src); + compute_field(src_v2d, pmid_src); + compute_field(src_s3d_m,pmid_src); + compute_field(src_s3d_i,pint_src); + compute_field(src_v3d_m,pmid_src); + compute_field(src_v3d_i,pint_src); + + // Pre-compute what we expect the tgt fields to be + compute_field(expected_s2d, pmid_tgt); + compute_field(expected_v2d, pmid_tgt); + compute_field(expected_s3d_m,pmid_tgt); + compute_field(expected_s3d_i,pint_tgt); + compute_field(expected_v3d_m,pmid_tgt); + compute_field(expected_v3d_i,pint_tgt); + + extrapolate(pmid_src,pmid_tgt,expected_s2d, etype_top,etype_bot); + extrapolate(pmid_src,pmid_tgt,expected_v2d, etype_top,etype_bot); + extrapolate(pmid_src,pmid_tgt,expected_s3d_m,etype_top,etype_bot); + extrapolate(pint_src,pint_tgt,expected_s3d_i,etype_top,etype_bot); + extrapolate(pmid_src,pmid_tgt,expected_v3d_m,etype_top,etype_bot); + extrapolate(pint_src,pint_tgt,expected_v3d_i,etype_top,etype_bot); + print (" -> generate fields data ... done!\n",comm); + + // -------------------------------------- // + // Perform remap // + // -------------------------------------- // + + // No bwd remap + REQUIRE_THROWS(remap->remap(false)); + + print (" -> run remap ...\n",comm); + remap->remap(true); + print (" -> run remap ... done!\n",comm); + + // -------------------------------------- // + // Check remapped fields // + // -------------------------------------- // + + using namespace Catch::Matchers; + Real tol = 10*std::numeric_limits::epsilon(); + + print (" -> check tgt fields ...\n",comm); + { + auto diff = tgt_s2d.clone("diff"); + auto ex_norm = frobenius_norm(expected_s2d); + diff.update(expected_s2d,1/ex_norm,-1/ex_norm); + REQUIRE (frobenius_norm(diff)(expected_v2d); + diff.update(expected_v2d,1/ex_norm,-1/ex_norm); + REQUIRE (frobenius_norm(diff)(expected_s3d_m); + diff.update(expected_s3d_m,1/ex_norm,-1/ex_norm); + REQUIRE (frobenius_norm(diff)(expected_s3d_i); + diff.update(expected_s3d_i,1/ex_norm,-1/ex_norm); + REQUIRE (frobenius_norm(diff)(expected_v3d_m); + diff.update(expected_v3d_m,1 / ex_norm,-1 / ex_norm); + REQUIRE (frobenius_norm(diff)(expected_v3d_i); + diff.update(expected_v3d_i,1 / ex_norm,-1 / ex_norm); + REQUIRE (frobenius_norm(diff) check tgt fields ... done!\n",comm); + } + } + } } - print ("check tgt fields ... done!\n",comm); } // Clean up scorpio stuff scorpio::finalize_subsystem(); + + print ("Testing vertical remapper ... done!\n",comm); } } // namespace scream From aebf3cb33a073148ac5efca15ef5bf008eecb54d Mon Sep 17 00:00:00 2001 From: Luca Bertagna Date: Wed, 13 Nov 2024 22:16:27 -0700 Subject: [PATCH 2/4] EAMxx: fix creation of layouts in VerticalRemapper --- .../share/grid/remap/vertical_remapper.cpp | 47 +++++++++---------- .../share/grid/remap/vertical_remapper.hpp | 7 +-- 2 files changed, 25 insertions(+), 29 deletions(-) diff --git a/components/eamxx/src/share/grid/remap/vertical_remapper.cpp b/components/eamxx/src/share/grid/remap/vertical_remapper.cpp index 8e13fd4e6df7..39056a50d466 100644 --- a/components/eamxx/src/share/grid/remap/vertical_remapper.cpp +++ b/components/eamxx/src/share/grid/remap/vertical_remapper.cpp @@ -55,6 +55,8 @@ VerticalRemapper (const grid_ptr_type& src_grid, // number of levels (i.e., pint/pmid cannot have the same nlevs). Since we remap // every field (mid or int) to the same pressure coords, we just hard-code them. m_tgt_pmid = m_tgt_pint = m_tgt_grid->get_geometry_data("p_levs"); + + m_tgt_int_same_as_mid = true; } VerticalRemapper:: @@ -72,13 +74,10 @@ VerticalRemapper (const grid_ptr_type& src_grid, FieldLayout VerticalRemapper:: create_src_layout (const FieldLayout& tgt_layout) const { - using namespace ShortFieldTagsNames; - - EKAT_REQUIRE_MSG (is_valid_tgt_layout(tgt_layout), - "[VerticalRemapper] Error! Input target layout is not valid for this remapper.\n" - " - input layout: " + tgt_layout.to_string()); - - return create_layout(tgt_layout,m_src_grid); + // Since we don't know if the tgt layout is "LEV for everything", + // we cannot infer what the corresponding src layout was. + // This function should never be used for this remapper. + EKAT_ERROR_MSG ("Error! VerticalRemapper does not support creating a src layout from a tgt layout.\n"); } FieldLayout VerticalRemapper:: @@ -90,44 +89,40 @@ create_tgt_layout (const FieldLayout& src_layout) const "[VerticalRemapper] Error! Input source layout is not valid for this remapper.\n" " - input layout: " + src_layout.to_string()); - return create_layout(src_layout,m_tgt_grid); -} - -FieldLayout VerticalRemapper:: -create_layout (const FieldLayout& fl_in, - const grid_ptr_type& grid_out) const -{ - // NOTE: for the vert remapper, it doesn't really make sense to distinguish - // between midpoints and interfaces: we're simply asking for a quantity - // at a given set of pressure levels. So we choose to have fl_out - // to *always* have LEV as vertical tag. - auto fl_out = FieldLayout::invalid(); - switch (fl_in.type()) { + // If we remap to a fixed set of pressure levels during I/O, + // it doesn't really make sense to distinguish between midpoints + // and interfaces, so choose fl_out to have LEV as vertical tag. + auto tgt_layout = FieldLayout::invalid(); + bool midpoints; + switch (src_layout.type()) { case LayoutType::Scalar0D: [[ fallthrough ]]; case LayoutType::Vector0D: [[ fallthrough ]]; case LayoutType::Scalar2D: [[ fallthrough ]]; case LayoutType::Vector2D: [[ fallthrough ]]; case LayoutType::Tensor2D: // These layouts do not have vertical dim tags, so no change - fl_out = fl_in; + tgt_layout = src_layout; break; case LayoutType::Scalar1D: - fl_out = grid_out->get_vertical_layout(true); + midpoints = m_tgt_int_same_as_mid || src_layout.tags().back()==LEV; + tgt_layout = m_tgt_grid->get_vertical_layout(midpoints); break; case LayoutType::Scalar3D: - fl_out = grid_out->get_3d_scalar_layout(true); + midpoints = m_tgt_int_same_as_mid || src_layout.tags().back()==LEV; + tgt_layout = m_tgt_grid->get_3d_scalar_layout(midpoints); break; case LayoutType::Vector3D: - fl_out = grid_out->get_3d_vector_layout(true,fl_in.get_vector_dim()); + midpoints = m_tgt_int_same_as_mid || src_layout.tags().back()==LEV; + tgt_layout = m_tgt_grid->get_3d_vector_layout(midpoints,src_layout.get_vector_dim()); break; default: // NOTE: this also include Tensor3D. We don't really have any atm proc // that needs to handle a tensor3d quantity, so no need to add it EKAT_ERROR_MSG ( "[VerticalRemapper] Error! Layout not supported by VerticalRemapper.\n" - " - input layout: " + fl_in.to_string() + "\n"); + " - input layout: " + src_layout.to_string() + "\n"); } - return fl_out; + return tgt_layout; } void VerticalRemapper:: diff --git a/components/eamxx/src/share/grid/remap/vertical_remapper.hpp b/components/eamxx/src/share/grid/remap/vertical_remapper.hpp index d763cddceab7..a22a9bb56df7 100644 --- a/components/eamxx/src/share/grid/remap/vertical_remapper.hpp +++ b/components/eamxx/src/share/grid/remap/vertical_remapper.hpp @@ -75,9 +75,6 @@ class VerticalRemapper : public AbstractRemapper protected: - FieldLayout create_layout (const FieldLayout& fl_in, - const grid_ptr_type& grid_out) const; - const identifier_type& do_get_src_field_id (const int ifield) const override { return m_src_fields[ifield].get_header().get_identifier(); } @@ -145,6 +142,10 @@ class VerticalRemapper : public AbstractRemapper Field m_tgt_pmid; Field m_tgt_pint; + // If we remap to a fixed set of pressure levels during I/O, + // our tgt pint would be the same as tgt pmid. + bool m_tgt_int_same_as_mid = false; + // Extrapolation settings at top/bottom. Default to P0 extrapolation ExtrapType m_etype_top = P0; ExtrapType m_etype_bot = P0; From 14df8d281ea57ecee39b2474c0d86d7842e618f4 Mon Sep 17 00:00:00 2001 From: Luca Bertagna Date: Thu, 14 Nov 2024 15:16:10 -0700 Subject: [PATCH 3/4] EAMxx: add comment regarding no bwd remap --- components/eamxx/src/share/grid/remap/vertical_remapper.cpp | 3 +++ 1 file changed, 3 insertions(+) diff --git a/components/eamxx/src/share/grid/remap/vertical_remapper.cpp b/components/eamxx/src/share/grid/remap/vertical_remapper.cpp index 39056a50d466..94f873951063 100644 --- a/components/eamxx/src/share/grid/remap/vertical_remapper.cpp +++ b/components/eamxx/src/share/grid/remap/vertical_remapper.cpp @@ -63,6 +63,9 @@ VerticalRemapper:: VerticalRemapper (const grid_ptr_type& src_grid, const grid_ptr_type& tgt_grid) { + // We only go in one direction for simplicity, since we need to setup some + // infrsatructures, and we don't want to setup 2x as many "just in case". + // If you need to remap bwd, just create another remapper with src/tgt grids swapped. m_bwd_allowed = false; EKAT_REQUIRE_MSG (src_grid->get_2d_scalar_layout().congruent(tgt_grid->get_2d_scalar_layout()), From 84d6d09bdb600d505c941af689a37543d5e6a7fa Mon Sep 17 00:00:00 2001 From: Luca Bertagna Date: Fri, 22 Nov 2024 18:59:40 -0700 Subject: [PATCH 4/4] EAMxx: remove generic host-device lambdas in VerticalRemapper --- .../eamxx/src/share/grid/remap/vertical_remapper.cpp | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/components/eamxx/src/share/grid/remap/vertical_remapper.cpp b/components/eamxx/src/share/grid/remap/vertical_remapper.cpp index 94f873951063..67c72e045a78 100644 --- a/components/eamxx/src/share/grid/remap/vertical_remapper.cpp +++ b/components/eamxx/src/share/grid/remap/vertical_remapper.cpp @@ -656,7 +656,9 @@ extrapolate (const Field& f_src, auto f_src_v = f_src.get_view(); auto f_tgt_v = f_tgt.get_view< Real**>(); auto policy = ESU::get_default_team_policy(ncols,nlevs_tgt); - auto lambda = KOKKOS_LAMBDA(const auto& team) + + using MemberType = typename decltype(policy)::member_type; + auto lambda = KOKKOS_LAMBDA(const MemberType& team) { const int icol = team.league_rank(); @@ -706,7 +708,8 @@ extrapolate (const Field& f_src, const int ncomps = f_tgt_l.get_vector_dim(); auto policy = ESU::get_default_team_policy(ncols*ncomps,nlevs_tgt); - auto lambda = KOKKOS_LAMBDA(const auto& team) + using MemberType = typename decltype(policy)::member_type; + auto lambda = KOKKOS_LAMBDA(const MemberType& team) { const int icol = team.league_rank() / ncomps; const int icmp = team.league_rank() % ncomps;