From abba54751773bbe55e640f6d7a4ce32fb06330a6 Mon Sep 17 00:00:00 2001 From: Chris Vogl Date: Wed, 19 Feb 2025 18:12:09 -0800 Subject: [PATCH 1/8] introduce zonal average diagnostic with unit test for EAMxx --- .../eamxx/src/diagnostics/CMakeLists.txt | 1 + .../src/diagnostics/register_diagnostics.hpp | 5 + .../src/diagnostics/tests/CMakeLists.txt | 4 +- .../src/diagnostics/tests/zonal_avg_test.cpp | 193 ++++++++++++++++++ .../eamxx/src/diagnostics/zonal_avg.cpp | 193 ++++++++++++++++++ .../eamxx/src/diagnostics/zonal_avg.hpp | 53 +++++ .../eamxx/src/share/io/scorpio_output.cpp | 7 +- 7 files changed, 452 insertions(+), 4 deletions(-) create mode 100644 components/eamxx/src/diagnostics/tests/zonal_avg_test.cpp create mode 100644 components/eamxx/src/diagnostics/zonal_avg.cpp create mode 100644 components/eamxx/src/diagnostics/zonal_avg.hpp diff --git a/components/eamxx/src/diagnostics/CMakeLists.txt b/components/eamxx/src/diagnostics/CMakeLists.txt index 812157390c96..8d722ffd2f0b 100644 --- a/components/eamxx/src/diagnostics/CMakeLists.txt +++ b/components/eamxx/src/diagnostics/CMakeLists.txt @@ -23,6 +23,7 @@ set(DIAGNOSTIC_SRCS water_path.cpp wind_speed.cpp vert_contract.cpp + zonal_avg.cpp ) add_library(diagnostics ${DIAGNOSTIC_SRCS}) diff --git a/components/eamxx/src/diagnostics/register_diagnostics.hpp b/components/eamxx/src/diagnostics/register_diagnostics.hpp index 1a512fb15ee0..f16c04c4479c 100644 --- a/components/eamxx/src/diagnostics/register_diagnostics.hpp +++ b/components/eamxx/src/diagnostics/register_diagnostics.hpp @@ -26,6 +26,7 @@ #include "diagnostics/atm_backtend.hpp" #include "diagnostics/horiz_avg.hpp" #include "diagnostics/vert_contract.hpp" +#include "diagnostics/zonal_avg.hpp" namespace scream { @@ -54,7 +55,11 @@ inline void register_diagnostics () { diag_factory.register_product("AeroComCld",&create_atmosphere_diagnostic); diag_factory.register_product("AtmBackTendDiag",&create_atmosphere_diagnostic); diag_factory.register_product("HorizAvgDiag",&create_atmosphere_diagnostic); +<<<<<<< HEAD diag_factory.register_product("VertContractDiag",&create_atmosphere_diagnostic); +======= + diag_factory.register_product("ZonalAvgDiag",&create_atmosphere_diagnostic); +>>>>>>> be844e83c7 (initial work on zonal diagnostic for EAMxx) } } // namespace scream diff --git a/components/eamxx/src/diagnostics/tests/CMakeLists.txt b/components/eamxx/src/diagnostics/tests/CMakeLists.txt index e2ed18b209a0..ac4f04834671 100644 --- a/components/eamxx/src/diagnostics/tests/CMakeLists.txt +++ b/components/eamxx/src/diagnostics/tests/CMakeLists.txt @@ -76,4 +76,6 @@ CreateDiagTest(atm_backtend "atm_backtend_test.cpp") CreateDiagTest(horiz_avg "horiz_avg_test.cpp") # Test for vertical contraction -CreateDiagTest(vert_contract "vert_contract_test.cpp") + +# Test zonal averaging +CreateDiagTest(zonal_avg "zonal_avg_test.cpp") diff --git a/components/eamxx/src/diagnostics/tests/zonal_avg_test.cpp b/components/eamxx/src/diagnostics/tests/zonal_avg_test.cpp new file mode 100644 index 000000000000..d26bba41ef66 --- /dev/null +++ b/components/eamxx/src/diagnostics/tests/zonal_avg_test.cpp @@ -0,0 +1,193 @@ +#include "catch2/catch.hpp" +#include "diagnostics/register_diagnostics.hpp" +#include "share/field/field_utils.hpp" +#include "share/grid/mesh_free_grids_manager.hpp" +#include "share/util/scream_setup_random_test.hpp" +#include "share/util/scream_universal_constants.hpp" + +namespace scream { + +std::shared_ptr create_gm(const ekat::Comm &comm, const int ncols, + const int nlevs) { + const int num_global_cols = ncols * comm.size(); + + using vos_t = std::vector; + ekat::ParameterList gm_params; + gm_params.set("grids_names", vos_t{"Point Grid"}); + auto &pl = gm_params.sublist("Point Grid"); + pl.set("type", "point_grid"); + pl.set("aliases", vos_t{"Physics"}); + pl.set("number_of_global_columns", num_global_cols); + pl.set("number_of_vertical_levels", nlevs); + + auto gm = create_mesh_free_grids_manager(comm, gm_params); + gm->build_grids(); + + return gm; +} + +TEST_CASE("zonal_avg") { + using namespace ShortFieldTagsNames; + using namespace ekat::units; + using TeamPolicy = Kokkos::TeamPolicy; + using TeamMember = typename TeamPolicy::member_type; + using KT = ekat::KokkosTypes; + using ESU = ekat::ExeSpaceUtils; + + // A numerical tolerance + auto tol = std::numeric_limits::epsilon() * 100; + + // A world comm + ekat::Comm comm(MPI_COMM_WORLD); + + // A time stamp + util::TimeStamp t0({2024, 1, 1}, {0, 0, 0}); + + // Create a grids manager - single column for these tests + constexpr int nlevs = 3; + constexpr int dim3 = 4; + const int ngcols = 6 * comm.size(); + const int nlats = 4; + + auto gm = create_gm(comm, ngcols, nlevs); + auto grid = gm->get_grid("Physics"); + + Field area = grid->get_geometry_data("area"); + auto area_view = area.get_view(); + + // Set latitude values + Field lat = gm->get_grid_nonconst("Physics")->create_geometry_data("lat", + grid->get_2d_scalar_layout(), Units::nondimensional()); + auto lat_view = lat.get_view(); + const Real lat_delta = sp(180.0) / nlats; + std::vector zonal_areas(nlats, 0.0); + for (int i=0; i < ngcols; i++) { + lat_view(i) = sp(-90.0) + (i % nlats + sp(0.5)) * lat_delta; + zonal_areas[i % nlats] += area_view[i]; + } + + // Input (randomized) qc + FieldLayout scalar1d_layout{{COL}, {ngcols}}; + FieldLayout scalar2d_layout{{COL, LEV}, {ngcols, nlevs}}; + FieldLayout scalar3d_layout{{COL, CMP, LEV}, {ngcols, dim3, nlevs}}; + + FieldIdentifier qc1_id("qc", scalar1d_layout, kg / kg, grid->name()); + FieldIdentifier qc2_fid("qc", scalar2d_layout, kg / kg, grid->name()); + FieldIdentifier qc3_fid("qc", scalar3d_layout, kg / kg, grid->name()); + + Field qc1(qc1_id); + Field qc2(qc2_fid); + Field qc3(qc3_fid); + + qc1.allocate_view(); + qc2.allocate_view(); + qc3.allocate_view(); + + // Construct random number generator stuff + using RPDF = std::uniform_real_distribution; + RPDF pdf(sp(0.0), sp(200.0)); + auto engine = scream::setup_random_test(); + + // Construct the Diagnostics + std::map> diags; + auto &diag_factory = AtmosphereDiagnosticFactory::instance(); + register_diagnostics(); + + ekat::ParameterList params; + REQUIRE_THROWS(diag_factory.create("ZonalAvgDiag", comm, + params)); // No 'field_name' parameter + + // Set time for qc and randomize its values + qc1.get_header().get_tracking().update_time_stamp(t0); + qc2.get_header().get_tracking().update_time_stamp(t0); + qc3.get_header().get_tracking().update_time_stamp(t0); + randomize(qc1, engine, pdf); + randomize(qc2, engine, pdf); + randomize(qc3, engine, pdf); + + // Create and set up the diagnostic + params.set("grid_name", grid->name()); + params.set("field_name", "qc"); + params.set("num_lat_vals", nlats); + auto diag1 = diag_factory.create("ZonalAvgDiag", comm, params); + auto diag2 = diag_factory.create("ZonalAvgDiag", comm, params); + auto diag3 = diag_factory.create("ZonalAvgDiag", comm, params); + diag1->set_grids(gm); + diag2->set_grids(gm); + diag3->set_grids(gm); + + // Test the zonal average of qc1 + diag1->set_required_field(qc1); + diag1->initialize(t0, RunType::Initial); + diag1->compute_diagnostic(); + auto diag1_field = diag1->get_diagnostic(); + + // Manual calculation + FieldLayout diag0_layout({CMP}, {nlats}, {ZonalAvgDiag::dim_name}); + FieldIdentifier diag0_id("qc_zonal_avg_manual", diag0_layout, kg / kg, + grid->name()); + Field diag0_field(diag0_id); + diag0_field.allocate_view(); + + // calculate the zonal average + auto qc1_view = qc1.get_view(); + auto diag0_view = diag0_field.get_view(); + for (int i=0; i < ngcols; i++) { + const int nlat = i % nlats; + diag0_view(nlat) += area_view(i) / zonal_areas[nlat] * qc1_view(i); + } + + // Compare + REQUIRE(views_are_equal(diag1_field, diag0_field)); + + // Try other known cases + // Set qc1_v to 1.0 to get zonal averages of 1.0/nlats + const Real zavg1 = sp(1.0); + qc1.deep_copy(zavg1); + diag1->compute_diagnostic(); + auto diag1_view_host = diag1_field.get_view(); + for (int nlat=0; nlat < nlats; nlat++) { + REQUIRE_THAT(diag1_view_host(nlat), Catch::Matchers::WithinRel(zavg1, tol)); + } + + // other diags + // Set qc2_v to 5.0 to get weighted average of 5.0 + const Real zavg2 = sp(5.0); + qc2.deep_copy(zavg2); + diag2->set_required_field(qc2); + diag2->initialize(t0, RunType::Initial); + diag2->compute_diagnostic(); + auto diag2_field = diag2->get_diagnostic(); + + auto diag2_view_host = diag2_field.get_view(); + for(int i = 0; i < nlevs; ++i) { + for (int nlat=0; nlat < nlats; nlat++) { + REQUIRE_THAT(diag2_view_host(nlat,i), Catch::Matchers::WithinRel(zavg2, tol)); + } + } + + // Try a random case with qc3 + FieldLayout diag3m_layout({CMP, CMP, LEV}, {nlats, dim3, nlevs}, + {ZonalAvgDiag::dim_name, e2str(CMP), e2str(LEV)}); + FieldIdentifier diag3m_id("qc_zonal_avg_manual", diag3m_layout, kg / kg, + grid->name()); + Field diag3m_field(diag3m_id); + diag3m_field.allocate_view(); + auto qc3_view = qc3.get_view(); + auto diag3m_view = diag3m_field.get_view(); + for (int i=0; i < ngcols; i++) { + const int nlat = i % nlats; + for (int j=0; j < dim3; j++) { + for (int k=0; k < nlevs; k++) { + diag3m_view(nlat,j,k) += area_view(i) / zonal_areas[nlat] * qc3_view(i,j,k); + } + } + } + diag3->set_required_field(qc3); + diag3->initialize(t0, RunType::Initial); + diag3->compute_diagnostic(); + auto diag3_field = diag3->get_diagnostic(); + REQUIRE(views_are_equal(diag3_field, diag3m_field)); +} + +} // namespace scream diff --git a/components/eamxx/src/diagnostics/zonal_avg.cpp b/components/eamxx/src/diagnostics/zonal_avg.cpp new file mode 100644 index 000000000000..519d825c5b09 --- /dev/null +++ b/components/eamxx/src/diagnostics/zonal_avg.cpp @@ -0,0 +1,193 @@ +#include "diagnostics/zonal_avg.hpp" + +#include "share/field/field_utils.hpp" + +namespace scream { + +void ZonalAvgDiag::compute_zonal_sum(const Field &result, + const Field &field, const Field &weight, const Field &lat, + const ekat::Comm *comm) +{ + auto result_layout = result.get_header().get_identifier().get_layout(); + const int lat_num = result_layout.dim(ZonalAvgDiag::dim_name); + const int ncols = field.get_header().get_identifier().get_layout().dim(0); + const Real lat_delta = sp(180.0) / lat_num; + + auto weight_view = weight.get_view(); + auto lat_view = lat.get_view(); + using KT = ekat::KokkosTypes; + using TeamPolicy = Kokkos::TeamPolicy; + using TeamMember = typename TeamPolicy::member_type; + using ESU = ekat::ExeSpaceUtils; + switch(result_layout.rank()) + { + case 1: { + auto field_view = field.get_view(); + auto result_view = result.get_view(); + TeamPolicy team_policy = ESU::get_default_team_policy(lat_num, ncols); + Kokkos::parallel_for("compute_zonal_sum_" + field.name(), team_policy, + KOKKOS_LAMBDA(const TeamMember& tm) { + const int lat_i = tm.league_rank(); + const Real lat_lower = sp(-90.0) + lat_i * lat_delta; + const Real lat_upper = lat_lower + lat_delta; + Kokkos::parallel_reduce(Kokkos::TeamVectorRange(tm, ncols), + KOKKOS_LAMBDA(int i, Real& val) { + // TODO: check if tenary is ok here (if not, multiply by flag) + int flag = (lat_lower <= lat_view(i)) && (lat_view(i) < lat_upper); + val += flag ? weight_view(i) * field_view(i) : sp(0.0); + }, result_view(lat_i)); + }); + } break; + case 2: { + const int d1 = result_layout.dim(1); + auto field_view = field.get_view(); + auto result_view = result.get_view(); + TeamPolicy team_policy = + ESU::get_default_team_policy(lat_num * d1, ncols); + Kokkos::parallel_for("compute_zonal_sum_" + field.name(), team_policy, + KOKKOS_LAMBDA(const TeamMember& tm) { + const int idx = tm.league_rank(); + const int d1_i = idx / lat_num; + const int lat_i = idx % lat_num; + const Real lat_lower = sp(-90.0) + lat_i * lat_delta; + const Real lat_upper = lat_lower + lat_delta; + Kokkos::parallel_reduce(Kokkos::TeamVectorRange(tm, ncols), + KOKKOS_LAMBDA(int i, Real& val) { + int flag = (lat_lower <= lat_view(i)) && (lat_view(i) < lat_upper); + // TODO: check if tenary is ok here (if not, multiply by flag) + val += flag ? weight_view(i) * field_view(i,d1_i) : sp(0.0); + }, result_view(lat_i,d1_i)); + }); + } break; + case 3: { + const int d1 = result_layout.dim(1); + const int d2 = result_layout.dim(2); + auto field_view = field.get_view(); + auto result_view = result.get_view(); + TeamPolicy team_policy = + ESU::get_default_team_policy(lat_num * d1 * d2, ncols); + Kokkos::parallel_for("compute_zonal_sum_" + field.name(), team_policy, + KOKKOS_LAMBDA(const TeamMember& tm) { + const int idx = tm.league_rank(); + const int d1_i = idx / (lat_num * d2); + const int idx2 = idx % (lat_num * d2); + const int d2_i = idx2 / lat_num; + const int lat_i = idx2 % lat_num; + const Real lat_lower = sp(-90.0) + lat_i * lat_delta; + const Real lat_upper = lat_lower + lat_delta; + Kokkos::parallel_reduce(Kokkos::TeamVectorRange(tm, ncols), + KOKKOS_LAMBDA(int i, Real& val) { + int flag = (lat_lower <= lat_view(i)) && (lat_view(i) < lat_upper); + // TODO: check if tenary is ok here (if not, multiply by flag) + val += flag ? weight_view(i) * field_view(i,d1_i,d2_i) : sp(0.0); + }, result_view(lat_i,d1_i,d2_i)); + }); + } break; + default: + EKAT_ERROR_MSG("Error! Unsupported field rank for zonal averages.\n"); + } + + if (comm) { + // TODO: use device-side MPI calls + // TODO: the dev ptr causes problems; revisit this later + // TODO: doing cuda-aware MPI allreduce would be ~10% faster + Kokkos::fence(); + result.sync_to_host(); + comm->all_reduce(result.template get_internal_view_data(), + result_layout.size(), MPI_SUM); + result.sync_to_dev(); + } +} + +ZonalAvgDiag::ZonalAvgDiag(const ekat::Comm &comm, + const ekat::ParameterList ¶ms) + : AtmosphereDiagnostic(comm, params) { + const auto &field_name = m_params.get("field_name"); + m_diag_name = field_name + "_zonal_avg"; + + m_lat_num = params.isParameter("num_lat_vals") ? + params.get("num_lat_vals") : 180; +} + +void ZonalAvgDiag::set_grids( + const std::shared_ptr grids_manager) { + const auto &field_name = m_params.get("field_name"); + const auto &grid_name = m_params.get("grid_name"); + const GridsManager::grid_ptr_type grid = grids_manager->get_grid("Physics"); + + add_field(field_name, grid_name); + + m_lat = grid->get_geometry_data("lat"); + // area will be scaled in initialize_impl + m_scaled_area = grid->get_geometry_data("area").clone(); +} + +void ZonalAvgDiag::initialize_impl(const RunType /*run_type*/) { + // TODO: auto everything + using FieldIdentifier = FieldHeader::identifier_type; + using FieldLayout = FieldIdentifier::layout_type; + using ShortFieldTagsNames::COL, ShortFieldTagsNames::CMP, + ShortFieldTagsNames::LEV; + const Field &field = get_fields_in().front(); + const FieldIdentifier &field_id = field.get_header().get_identifier(); + const FieldLayout &field_layout = field_id.get_layout(); + + EKAT_REQUIRE_MSG(field_layout.rank() >= 1 && field_layout.rank() <= 3, + "Error! Field rank not supported by ZonalAvgDiag.\n" + " - field name: " + field_id.name() + "\n" + " - field layout: " + field_layout.to_string() + "\n"); + EKAT_REQUIRE_MSG(field_layout.tags()[0] == COL, + "Error! ZonalAvgDiag diagnostic expects a layout starting " + "with the 'COL' tag.\n" + " - field name : " + field_id.name() + "\n" + " - field layout: " + field_layout.to_string() + "\n"); + + FieldLayout diagnostic_layout({CMP}, {m_lat_num}, {ZonalAvgDiag::dim_name}); + for (int idim=1; idim < field_layout.rank(); idim++) + { + diagnostic_layout.append_dim(field_layout.tag(idim), field_layout.dim(idim), + field_layout.name(idim)); + } + + FieldIdentifier diagnostic_id(m_diag_name, diagnostic_layout, + field_id.get_units(), field_id.get_grid_name()); + m_diagnostic_output = Field(diagnostic_id); + m_diagnostic_output.allocate_view(); + + // allocate zonal area + const FieldIdentifier &area_id = m_scaled_area.get_header().get_identifier(); + FieldLayout zonal_area_layout({CMP}, {m_lat_num}, {ZonalAvgDiag::dim_name}); + FieldIdentifier zonal_area_id("zonal area", zonal_area_layout, + area_id.get_units(), area_id.get_grid_name()); + Field zonal_area(zonal_area_id); + zonal_area.allocate_view(); + + // compute zonal area + FieldLayout ones_layout = area_id.get_layout().clone(); + FieldIdentifier ones_id("ones", ones_layout, area_id.get_units(), + area_id.get_grid_name()); + Field ones(ones_id); + ones.allocate_view(); + ones.deep_copy(1.0); + compute_zonal_sum(zonal_area, m_scaled_area, ones, m_lat, &m_comm); + + // scale area by 1 / zonal area + using RangePolicy = Kokkos::RangePolicy; + const Real lat_delta = sp(180.0) / m_lat_num; + const int ncols = field_layout.dim(0); + auto lat_view = m_lat.get_view(); + auto zonal_area_view = zonal_area.get_view(); + auto scaled_area_view = m_scaled_area.get_view(); + Kokkos::parallel_for("scale_area_by_zonal_area_" + field.name(), + RangePolicy(0, ncols), KOKKOS_LAMBDA(const int& i) { + const int lat_i = (lat_view(i) + sp(90.0)) / lat_delta; + scaled_area_view(i) /= zonal_area_view(lat_i); + }); +} + +void ZonalAvgDiag::compute_diagnostic_impl() { + const auto &field = get_fields_in().front(); + compute_zonal_sum(m_diagnostic_output, field, m_scaled_area, m_lat, &m_comm); +} + +} // namespace scream diff --git a/components/eamxx/src/diagnostics/zonal_avg.hpp b/components/eamxx/src/diagnostics/zonal_avg.hpp new file mode 100644 index 000000000000..827944d0e50e --- /dev/null +++ b/components/eamxx/src/diagnostics/zonal_avg.hpp @@ -0,0 +1,53 @@ +#ifndef EAMXX_ZONAL_AVERAGE_HPP +#define EAMXX_ZONAL_AVERAGE_HPP + +#include "share/atm_process/atmosphere_diagnostic.hpp" + +namespace scream { +// TODO: Update this comment +/* + * This diagnostic will calculate the area-weighted average of a field + * across the COL tag dimension, producing an N-1 dimensional field + * that is area-weighted average of the input field. + */ + +class ZonalAvgDiag : public AtmosphereDiagnostic { + + // TODO: comment this, noting it's a utility function that could exist elsewhere + static void compute_zonal_sum(const Field &result, const Field &field, + const Field &weight, const Field &lat, const ekat::Comm *comm = nullptr); + + public: + + inline static const std::string dim_name = "bin"; + + // Constructors + ZonalAvgDiag(const ekat::Comm &comm, const ekat::ParameterList ¶ms); + + // The name of the diagnostic + std::string name() const { return m_diag_name; } + + // Set the grid + void set_grids(const std::shared_ptr grids_manager); + + protected: +#ifdef KOKKOS_ENABLE_CUDA + public: +#endif + void compute_diagnostic_impl(); + + protected: + void initialize_impl(const RunType /*run_type*/); + + std::string m_diag_name; + + Field m_lat; + int m_lat_num; + + Field m_scaled_area; + +}; + +} // namespace scream + +#endif // EAMXX_ZONAL_AVERAGE_HPP diff --git a/components/eamxx/src/share/io/scorpio_output.cpp b/components/eamxx/src/share/io/scorpio_output.cpp index 7f168b4a9986..d2751894b8ab 100644 --- a/components/eamxx/src/share/io/scorpio_output.cpp +++ b/components/eamxx/src/share/io/scorpio_output.cpp @@ -790,7 +790,7 @@ void AtmosphereOutput::register_dimensions(const std::string& name) // If t==CMP, and the name stored in the layout is the default ("dim"), // we append also the extent, to allow different vector dims in the file - tag_name += tag_name=="dim" ? std::to_string(dims[i]) : ""; + tag_name += tags[i] == CMP ? std::to_string(dims[i]) : ""; auto is_partitioned = m_io_grid->get_partitioned_dim_tag()==tags[i]; int dim_len = is_partitioned @@ -873,6 +873,7 @@ void AtmosphereOutput::set_avg_cnt_tracking(const std::string& name, const Field // Now create and store a dev view to track the averaging count for this layout (if we are tracking) // We don't need to track average counts for files that are not tracking the time dim + using namespace ShortFieldTagsNames; const auto& avg_cnt_suffix = m_field_to_avg_cnt_suffix[name]; const auto size = layout.size(); const auto tags = layout.tags(); @@ -886,7 +887,7 @@ void AtmosphereOutput::set_avg_cnt_tracking(const std::string& name, const Field // If t==CMP, and the name stored in the layout is the default ("dim"), // we append also the extent, to allow different vector dims in the file - tag_name += tag_name=="dim" ? std::to_string(layout.dim(i)) : ""; + tag_name += t==CMP ? std::to_string(layout.dim(i)) : ""; avg_cnt_name += "_" + tag_name; } @@ -951,7 +952,7 @@ register_variables(const std::string& filename, auto tag_name = m_io_grid->has_special_tag_name(t) ? m_io_grid->get_special_tag_name(t) : layout.names()[i]; - if (tag_name=="dim") { + if (t==CMP) { tag_name += std::to_string(layout.dim(i)); } vec_of_dims.push_back(tag_name); // Add dimensions string to vector of dims. From 5895598f7bfc28660b0d9956a0885b218526a54a Mon Sep 17 00:00:00 2001 From: Chris Vogl Date: Fri, 18 Apr 2025 12:54:45 -0700 Subject: [PATCH 2/8] added FieldLayout::prepend_dim, along with utility private function FieldLayout::add to avoid duplication of code between append and prepend --- .../eamxx/src/share/field/field_layout.cpp | 40 ++++++++++++++++--- .../eamxx/src/share/field/field_layout.hpp | 8 +++- 2 files changed, 41 insertions(+), 7 deletions(-) diff --git a/components/eamxx/src/share/field/field_layout.cpp b/components/eamxx/src/share/field/field_layout.cpp index 3b2c4c195ec0..74a47e5a24f5 100644 --- a/components/eamxx/src/share/field/field_layout.cpp +++ b/components/eamxx/src/share/field/field_layout.cpp @@ -148,6 +148,27 @@ FieldLayout& FieldLayout::strip_dim (const int idim) { return *this; } +void FieldLayout::add_dim (const FieldTag t, const int extent, const std::string& name, + bool prepend_instead_of_append) +{ + if (prepend_instead_of_append) + { + m_tags.insert(m_tags.begin(), t); + m_names.insert(m_names.begin(), name); + m_dims.insert(m_dims.begin(), extent); + } + else + { + m_tags.push_back(t); + m_names.push_back(name); + m_dims.push_back(extent); + } + + ++m_rank; + set_extents(); + compute_type(); +} + FieldLayout& FieldLayout::append_dim (const FieldTag t, const int extent) { @@ -157,13 +178,20 @@ FieldLayout::append_dim (const FieldTag t, const int extent) FieldLayout& FieldLayout::append_dim (const FieldTag t, const int extent, const std::string& name) { - m_tags.push_back(t); - m_names.push_back(name); - m_dims.push_back(extent); + add_dim(t, extent, name); + return *this; +} + +FieldLayout& +FieldLayout::prepend_dim (const FieldTag t, const int extent) +{ + return prepend_dim(t,extent,e2str(t)); +} - ++m_rank; - set_extents(); - compute_type(); +FieldLayout& +FieldLayout::prepend_dim (const FieldTag t, const int extent, const std::string& name) +{ + add_dim(t, extent, name, true); return *this; } diff --git a/components/eamxx/src/share/field/field_layout.hpp b/components/eamxx/src/share/field/field_layout.hpp index f4d54168b866..4254712486b3 100644 --- a/components/eamxx/src/share/field/field_layout.hpp +++ b/components/eamxx/src/share/field/field_layout.hpp @@ -59,6 +59,10 @@ inline std::string e2str (const LayoutType lt) { */ class FieldLayout { + + void add_dim (const FieldTag t, const int extent, const std::string& name, + bool prepend_instead_of_append = false); + public: using extents_type = typename KokkosTypes::view_1d; @@ -132,6 +136,8 @@ class FieldLayout { FieldLayout& strip_dim (const int idim); FieldLayout& append_dim (const FieldTag t, const int extent); FieldLayout& append_dim (const FieldTag t, const int extent, const std::string& name); + FieldLayout& prepend_dim (const FieldTag t, const int extent); + FieldLayout& prepend_dim (const FieldTag t, const int extent, const std::string& name); FieldLayout& rename_dim (const int idim, const std::string& n); FieldLayout& rename_dim (const FieldTag tag, const std::string& n, const bool throw_if_not_found = true); FieldLayout& reset_dim (const int idim, const int extent); @@ -214,7 +220,7 @@ inline long long FieldLayout::size () const { return prod; } -inline FieldTag FieldLayout::tag (const int idim) const { +inline FieldTag FieldLayout::tag (const int idim) const { EKAT_REQUIRE_MSG (idim>=0 && idim Date: Fri, 18 Apr 2025 12:55:06 -0700 Subject: [PATCH 3/8] made use of FieldLayout::prepend_dim functionality in zonal average diagnostic for EAMxx --- components/eamxx/src/diagnostics/zonal_avg.cpp | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/components/eamxx/src/diagnostics/zonal_avg.cpp b/components/eamxx/src/diagnostics/zonal_avg.cpp index 519d825c5b09..be6cb256cdb0 100644 --- a/components/eamxx/src/diagnostics/zonal_avg.cpp +++ b/components/eamxx/src/diagnostics/zonal_avg.cpp @@ -142,13 +142,9 @@ void ZonalAvgDiag::initialize_impl(const RunType /*run_type*/) { " - field name : " + field_id.name() + "\n" " - field layout: " + field_layout.to_string() + "\n"); - FieldLayout diagnostic_layout({CMP}, {m_lat_num}, {ZonalAvgDiag::dim_name}); - for (int idim=1; idim < field_layout.rank(); idim++) - { - diagnostic_layout.append_dim(field_layout.tag(idim), field_layout.dim(idim), - field_layout.name(idim)); - } - + FieldLayout diagnostic_layout = + field_layout.clone().strip_dim(COL).prepend_dim({CMP}, {m_lat_num}, + {ZonalAvgDiag::dim_name}); FieldIdentifier diagnostic_id(m_diag_name, diagnostic_layout, field_id.get_units(), field_id.get_grid_name()); m_diagnostic_output = Field(diagnostic_id); From e6af439851f339ca4dbcbba3119ee5f0e370c416 Mon Sep 17 00:00:00 2001 From: Chris Vogl Date: Mon, 21 Apr 2025 12:38:14 -0700 Subject: [PATCH 4/8] added EAMxx zonal average diagnostic to scorpio and added to test homm_shoc_cld_p3_rrtgmp_pg2 --- .../src/diagnostics/register_diagnostics.hpp | 3 - .../src/diagnostics/tests/CMakeLists.txt | 3 +- .../src/diagnostics/tests/zonal_avg_test.cpp | 58 +++-- .../eamxx/src/diagnostics/zonal_avg.cpp | 222 +++++++++--------- .../eamxx/src/diagnostics/zonal_avg.hpp | 41 ++-- .../eamxx/src/share/field/field_layout.hpp | 7 +- .../eamxx/src/share/io/eamxx_io_utils.cpp | 7 + .../homme_shoc_cld_p3_rrtmgp_pg2/output.yaml | 2 + 8 files changed, 174 insertions(+), 169 deletions(-) diff --git a/components/eamxx/src/diagnostics/register_diagnostics.hpp b/components/eamxx/src/diagnostics/register_diagnostics.hpp index f16c04c4479c..67298fda51cb 100644 --- a/components/eamxx/src/diagnostics/register_diagnostics.hpp +++ b/components/eamxx/src/diagnostics/register_diagnostics.hpp @@ -55,11 +55,8 @@ inline void register_diagnostics () { diag_factory.register_product("AeroComCld",&create_atmosphere_diagnostic); diag_factory.register_product("AtmBackTendDiag",&create_atmosphere_diagnostic); diag_factory.register_product("HorizAvgDiag",&create_atmosphere_diagnostic); -<<<<<<< HEAD diag_factory.register_product("VertContractDiag",&create_atmosphere_diagnostic); -======= diag_factory.register_product("ZonalAvgDiag",&create_atmosphere_diagnostic); ->>>>>>> be844e83c7 (initial work on zonal diagnostic for EAMxx) } } // namespace scream diff --git a/components/eamxx/src/diagnostics/tests/CMakeLists.txt b/components/eamxx/src/diagnostics/tests/CMakeLists.txt index ac4f04834671..9eca934e0da4 100644 --- a/components/eamxx/src/diagnostics/tests/CMakeLists.txt +++ b/components/eamxx/src/diagnostics/tests/CMakeLists.txt @@ -76,6 +76,7 @@ CreateDiagTest(atm_backtend "atm_backtend_test.cpp") CreateDiagTest(horiz_avg "horiz_avg_test.cpp") # Test for vertical contraction +CreateDiagTest(vert_contract "vert_contract_test.cpp") # Test zonal averaging -CreateDiagTest(zonal_avg "zonal_avg_test.cpp") +CreateDiagTest(zonal_avg "zonal_avg_test.cpp" MPI_RANKS 1 ${SCREAM_TEST_MAX_RANKS}) diff --git a/components/eamxx/src/diagnostics/tests/zonal_avg_test.cpp b/components/eamxx/src/diagnostics/tests/zonal_avg_test.cpp index d26bba41ef66..21e5317c6b85 100644 --- a/components/eamxx/src/diagnostics/tests/zonal_avg_test.cpp +++ b/components/eamxx/src/diagnostics/tests/zonal_avg_test.cpp @@ -2,13 +2,12 @@ #include "diagnostics/register_diagnostics.hpp" #include "share/field/field_utils.hpp" #include "share/grid/mesh_free_grids_manager.hpp" -#include "share/util/scream_setup_random_test.hpp" -#include "share/util/scream_universal_constants.hpp" +#include "share/util/eamxx_setup_random_test.hpp" +#include "share/util/eamxx_universal_constants.hpp" namespace scream { -std::shared_ptr create_gm(const ekat::Comm &comm, const int ncols, - const int nlevs) { +std::shared_ptr create_gm(const ekat::Comm &comm, const int ncols, const int nlevs) { const int num_global_cols = ncols * comm.size(); using vos_t = std::vector; @@ -47,21 +46,21 @@ TEST_CASE("zonal_avg") { constexpr int nlevs = 3; constexpr int dim3 = 4; const int ngcols = 6 * comm.size(); - const int nlats = 4; + const int nlats = 4; auto gm = create_gm(comm, ngcols, nlevs); auto grid = gm->get_grid("Physics"); - Field area = grid->get_geometry_data("area"); + Field area = grid->get_geometry_data("area"); auto area_view = area.get_view(); // Set latitude values - Field lat = gm->get_grid_nonconst("Physics")->create_geometry_data("lat", - grid->get_2d_scalar_layout(), Units::nondimensional()); - auto lat_view = lat.get_view(); + Field lat = gm->get_grid_nonconst("Physics")->create_geometry_data( + "lat", grid->get_2d_scalar_layout(), Units::nondimensional()); + auto lat_view = lat.get_view(); const Real lat_delta = sp(180.0) / nlats; std::vector zonal_areas(nlats, 0.0); - for (int i=0; i < ngcols; i++) { + for (int i = 0; i < ngcols; i++) { lat_view(i) = sp(-90.0) + (i % nlats + sp(0.5)) * lat_delta; zonal_areas[i % nlats] += area_view[i]; } @@ -95,7 +94,7 @@ TEST_CASE("zonal_avg") { ekat::ParameterList params; REQUIRE_THROWS(diag_factory.create("ZonalAvgDiag", comm, - params)); // No 'field_name' parameter + params)); // No 'field_name' parameter // Set time for qc and randomize its values qc1.get_header().get_tracking().update_time_stamp(t0); @@ -108,7 +107,7 @@ TEST_CASE("zonal_avg") { // Create and set up the diagnostic params.set("grid_name", grid->name()); params.set("field_name", "qc"); - params.set("num_lat_vals", nlats); + params.set("number_of_zonal_bins", std::to_string(nlats)); auto diag1 = diag_factory.create("ZonalAvgDiag", comm, params); auto diag2 = diag_factory.create("ZonalAvgDiag", comm, params); auto diag3 = diag_factory.create("ZonalAvgDiag", comm, params); @@ -123,16 +122,16 @@ TEST_CASE("zonal_avg") { auto diag1_field = diag1->get_diagnostic(); // Manual calculation - FieldLayout diag0_layout({CMP}, {nlats}, {ZonalAvgDiag::dim_name}); - FieldIdentifier diag0_id("qc_zonal_avg_manual", diag0_layout, kg / kg, - grid->name()); + const std::string bin_dim_name = diag1_field.get_header().get_identifier().get_layout().name(0); + FieldLayout diag0_layout({CMP}, {nlats}, {bin_dim_name}); + FieldIdentifier diag0_id("qc_zonal_avg_manual", diag0_layout, kg / kg, grid->name()); Field diag0_field(diag0_id); diag0_field.allocate_view(); // calculate the zonal average - auto qc1_view = qc1.get_view(); + auto qc1_view = qc1.get_view(); auto diag0_view = diag0_field.get_view(); - for (int i=0; i < ngcols; i++) { + for (int i = 0; i < ngcols; i++) { const int nlat = i % nlats; diag0_view(nlat) += area_view(i) / zonal_areas[nlat] * qc1_view(i); } @@ -146,7 +145,7 @@ TEST_CASE("zonal_avg") { qc1.deep_copy(zavg1); diag1->compute_diagnostic(); auto diag1_view_host = diag1_field.get_view(); - for (int nlat=0; nlat < nlats; nlat++) { + for (int nlat = 0; nlat < nlats; nlat++) { REQUIRE_THAT(diag1_view_host(nlat), Catch::Matchers::WithinRel(zavg1, tol)); } @@ -160,26 +159,25 @@ TEST_CASE("zonal_avg") { auto diag2_field = diag2->get_diagnostic(); auto diag2_view_host = diag2_field.get_view(); - for(int i = 0; i < nlevs; ++i) { - for (int nlat=0; nlat < nlats; nlat++) { - REQUIRE_THAT(diag2_view_host(nlat,i), Catch::Matchers::WithinRel(zavg2, tol)); + for (int i = 0; i < nlevs; ++i) { + for (int nlat = 0; nlat < nlats; nlat++) { + REQUIRE_THAT(diag2_view_host(nlat, i), Catch::Matchers::WithinRel(zavg2, tol)); } } // Try a random case with qc3 FieldLayout diag3m_layout({CMP, CMP, LEV}, {nlats, dim3, nlevs}, - {ZonalAvgDiag::dim_name, e2str(CMP), e2str(LEV)}); - FieldIdentifier diag3m_id("qc_zonal_avg_manual", diag3m_layout, kg / kg, - grid->name()); + {bin_dim_name, e2str(CMP), e2str(LEV)}); + FieldIdentifier diag3m_id("qc_zonal_avg_manual", diag3m_layout, kg / kg, grid->name()); Field diag3m_field(diag3m_id); diag3m_field.allocate_view(); - auto qc3_view = qc3.get_view(); + auto qc3_view = qc3.get_view(); auto diag3m_view = diag3m_field.get_view(); - for (int i=0; i < ngcols; i++) { + for (int i = 0; i < ngcols; i++) { const int nlat = i % nlats; - for (int j=0; j < dim3; j++) { - for (int k=0; k < nlevs; k++) { - diag3m_view(nlat,j,k) += area_view(i) / zonal_areas[nlat] * qc3_view(i,j,k); + for (int j = 0; j < dim3; j++) { + for (int k = 0; k < nlevs; k++) { + diag3m_view(nlat, j, k) += area_view(i) / zonal_areas[nlat] * qc3_view(i, j, k); } } } @@ -190,4 +188,4 @@ TEST_CASE("zonal_avg") { REQUIRE(views_are_equal(diag3_field, diag3m_field)); } -} // namespace scream +} // namespace scream diff --git a/components/eamxx/src/diagnostics/zonal_avg.cpp b/components/eamxx/src/diagnostics/zonal_avg.cpp index be6cb256cdb0..17219ddc9443 100644 --- a/components/eamxx/src/diagnostics/zonal_avg.cpp +++ b/components/eamxx/src/diagnostics/zonal_avg.cpp @@ -4,87 +4,88 @@ namespace scream { -void ZonalAvgDiag::compute_zonal_sum(const Field &result, - const Field &field, const Field &weight, const Field &lat, - const ekat::Comm *comm) -{ - auto result_layout = result.get_header().get_identifier().get_layout(); - const int lat_num = result_layout.dim(ZonalAvgDiag::dim_name); - const int ncols = field.get_header().get_identifier().get_layout().dim(0); - const Real lat_delta = sp(180.0) / lat_num; - - auto weight_view = weight.get_view(); - auto lat_view = lat.get_view(); - using KT = ekat::KokkosTypes; +void ZonalAvgDiag::compute_zonal_sum(const Field &result, const Field &field, const Field &weight, + const Field &lat, const ekat::Comm *comm) { + auto result_layout = result.get_header().get_identifier().get_layout(); + const int num_zonal_bins = result_layout.dim(0); + const int ncols = field.get_header().get_identifier().get_layout().dim(0); + const Real lat_delta = sp(180.0) / num_zonal_bins; + + auto weight_view = weight.get_view(); + auto lat_view = lat.get_view(); + using KT = ekat::KokkosTypes; using TeamPolicy = Kokkos::TeamPolicy; using TeamMember = typename TeamPolicy::member_type; - using ESU = ekat::ExeSpaceUtils; - switch(result_layout.rank()) - { - case 1: { - auto field_view = field.get_view(); - auto result_view = result.get_view(); - TeamPolicy team_policy = ESU::get_default_team_policy(lat_num, ncols); - Kokkos::parallel_for("compute_zonal_sum_" + field.name(), team_policy, - KOKKOS_LAMBDA(const TeamMember& tm) { - const int lat_i = tm.league_rank(); + using ESU = ekat::ExeSpaceUtils; + switch (result_layout.rank()) { + case 1: { + auto field_view = field.get_view(); + auto result_view = result.get_view(); + TeamPolicy team_policy = ESU::get_default_team_policy(num_zonal_bins, ncols); + Kokkos::parallel_for( + "compute_zonal_sum_" + field.name(), team_policy, KOKKOS_LAMBDA(const TeamMember &tm) { + const int lat_i = tm.league_rank(); const Real lat_lower = sp(-90.0) + lat_i * lat_delta; const Real lat_upper = lat_lower + lat_delta; - Kokkos::parallel_reduce(Kokkos::TeamVectorRange(tm, ncols), - KOKKOS_LAMBDA(int i, Real& val) { - // TODO: check if tenary is ok here (if not, multiply by flag) - int flag = (lat_lower <= lat_view(i)) && (lat_view(i) < lat_upper); - val += flag ? weight_view(i) * field_view(i) : sp(0.0); - }, result_view(lat_i)); + Kokkos::parallel_reduce( + Kokkos::TeamVectorRange(tm, ncols), + KOKKOS_LAMBDA(int i, Real &val) { + // TODO: check if tenary is ok here (if not, multiply by flag) + int flag = (lat_lower <= lat_view(i)) && (lat_view(i) < lat_upper); + val += flag ? weight_view(i) * field_view(i) : sp(0.0); + }, + result_view(lat_i)); }); - } break; - case 2: { - const int d1 = result_layout.dim(1); - auto field_view = field.get_view(); - auto result_view = result.get_view(); - TeamPolicy team_policy = - ESU::get_default_team_policy(lat_num * d1, ncols); - Kokkos::parallel_for("compute_zonal_sum_" + field.name(), team_policy, - KOKKOS_LAMBDA(const TeamMember& tm) { - const int idx = tm.league_rank(); - const int d1_i = idx / lat_num; - const int lat_i = idx % lat_num; + } break; + case 2: { + const int d1 = result_layout.dim(1); + auto field_view = field.get_view(); + auto result_view = result.get_view(); + TeamPolicy team_policy = ESU::get_default_team_policy(num_zonal_bins * d1, ncols); + Kokkos::parallel_for( + "compute_zonal_sum_" + field.name(), team_policy, KOKKOS_LAMBDA(const TeamMember &tm) { + const int idx = tm.league_rank(); + const int d1_i = idx / num_zonal_bins; + const int lat_i = idx % num_zonal_bins; const Real lat_lower = sp(-90.0) + lat_i * lat_delta; const Real lat_upper = lat_lower + lat_delta; - Kokkos::parallel_reduce(Kokkos::TeamVectorRange(tm, ncols), - KOKKOS_LAMBDA(int i, Real& val) { - int flag = (lat_lower <= lat_view(i)) && (lat_view(i) < lat_upper); - // TODO: check if tenary is ok here (if not, multiply by flag) - val += flag ? weight_view(i) * field_view(i,d1_i) : sp(0.0); - }, result_view(lat_i,d1_i)); + Kokkos::parallel_reduce( + Kokkos::TeamVectorRange(tm, ncols), + KOKKOS_LAMBDA(int i, Real &val) { + int flag = (lat_lower <= lat_view(i)) && (lat_view(i) < lat_upper); + // TODO: check if tenary is ok here (if not, multiply by flag) + val += flag ? weight_view(i) * field_view(i, d1_i) : sp(0.0); + }, + result_view(lat_i, d1_i)); }); - } break; - case 3: { - const int d1 = result_layout.dim(1); - const int d2 = result_layout.dim(2); - auto field_view = field.get_view(); - auto result_view = result.get_view(); - TeamPolicy team_policy = - ESU::get_default_team_policy(lat_num * d1 * d2, ncols); - Kokkos::parallel_for("compute_zonal_sum_" + field.name(), team_policy, - KOKKOS_LAMBDA(const TeamMember& tm) { - const int idx = tm.league_rank(); - const int d1_i = idx / (lat_num * d2); - const int idx2 = idx % (lat_num * d2); - const int d2_i = idx2 / lat_num; - const int lat_i = idx2 % lat_num; + } break; + case 3: { + const int d1 = result_layout.dim(1); + const int d2 = result_layout.dim(2); + auto field_view = field.get_view(); + auto result_view = result.get_view(); + TeamPolicy team_policy = ESU::get_default_team_policy(num_zonal_bins * d1 * d2, ncols); + Kokkos::parallel_for( + "compute_zonal_sum_" + field.name(), team_policy, KOKKOS_LAMBDA(const TeamMember &tm) { + const int idx = tm.league_rank(); + const int d1_i = idx / (num_zonal_bins * d2); + const int idx2 = idx % (num_zonal_bins * d2); + const int d2_i = idx2 / num_zonal_bins; + const int lat_i = idx2 % num_zonal_bins; const Real lat_lower = sp(-90.0) + lat_i * lat_delta; const Real lat_upper = lat_lower + lat_delta; - Kokkos::parallel_reduce(Kokkos::TeamVectorRange(tm, ncols), - KOKKOS_LAMBDA(int i, Real& val) { - int flag = (lat_lower <= lat_view(i)) && (lat_view(i) < lat_upper); - // TODO: check if tenary is ok here (if not, multiply by flag) - val += flag ? weight_view(i) * field_view(i,d1_i,d2_i) : sp(0.0); - }, result_view(lat_i,d1_i,d2_i)); + Kokkos::parallel_reduce( + Kokkos::TeamVectorRange(tm, ncols), + KOKKOS_LAMBDA(int i, Real &val) { + int flag = (lat_lower <= lat_view(i)) && (lat_view(i) < lat_upper); + // TODO: check if tenary is ok here (if not, multiply by flag) + val += flag ? weight_view(i) * field_view(i, d1_i, d2_i) : sp(0.0); + }, + result_view(lat_i, d1_i, d2_i)); }); - } break; - default: - EKAT_ERROR_MSG("Error! Unsupported field rank for zonal averages.\n"); + } break; + default: + EKAT_ERROR_MSG("Error! Unsupported field rank for zonal averages.\n"); } if (comm) { @@ -93,89 +94,88 @@ void ZonalAvgDiag::compute_zonal_sum(const Field &result, // TODO: doing cuda-aware MPI allreduce would be ~10% faster Kokkos::fence(); result.sync_to_host(); - comm->all_reduce(result.template get_internal_view_data(), - result_layout.size(), MPI_SUM); + comm->all_reduce(result.template get_internal_view_data(), result_layout.size(), + MPI_SUM); result.sync_to_dev(); } } -ZonalAvgDiag::ZonalAvgDiag(const ekat::Comm &comm, - const ekat::ParameterList ¶ms) +ZonalAvgDiag::ZonalAvgDiag(const ekat::Comm &comm, const ekat::ParameterList ¶ms) : AtmosphereDiagnostic(comm, params) { - const auto &field_name = m_params.get("field_name"); - m_diag_name = field_name + "_zonal_avg"; - - m_lat_num = params.isParameter("num_lat_vals") ? - params.get("num_lat_vals") : 180; + const auto &field_name = m_params.get("field_name"); + const auto &num_bins_value = params.get("number_of_zonal_bins"); + m_diag_name = field_name + "_zonal_avg_with_" + num_bins_value + "_bins"; + m_num_zonal_bins = std::stoi(num_bins_value); } -void ZonalAvgDiag::set_grids( - const std::shared_ptr grids_manager) { +void ZonalAvgDiag::set_grids(const std::shared_ptr grids_manager) { const auto &field_name = m_params.get("field_name"); - const auto &grid_name = m_params.get("grid_name"); - const GridsManager::grid_ptr_type grid = grids_manager->get_grid("Physics"); + const auto &grid_name = m_params.get("grid_name"); add_field(field_name, grid_name); - - m_lat = grid->get_geometry_data("lat"); + const GridsManager::grid_ptr_type grid = grids_manager->get_grid(grid_name); + m_lat = grid->get_geometry_data("lat"); // area will be scaled in initialize_impl m_scaled_area = grid->get_geometry_data("area").clone(); } void ZonalAvgDiag::initialize_impl(const RunType /*run_type*/) { - // TODO: auto everything using FieldIdentifier = FieldHeader::identifier_type; - using FieldLayout = FieldIdentifier::layout_type; - using ShortFieldTagsNames::COL, ShortFieldTagsNames::CMP, - ShortFieldTagsNames::LEV; - const Field &field = get_fields_in().front(); + using FieldLayout = FieldIdentifier::layout_type; + using ShortFieldTagsNames::COL, ShortFieldTagsNames::CMP, ShortFieldTagsNames::LEV; + const Field &field = get_fields_in().front(); const FieldIdentifier &field_id = field.get_header().get_identifier(); const FieldLayout &field_layout = field_id.get_layout(); EKAT_REQUIRE_MSG(field_layout.rank() >= 1 && field_layout.rank() <= 3, "Error! Field rank not supported by ZonalAvgDiag.\n" - " - field name: " + field_id.name() + "\n" - " - field layout: " + field_layout.to_string() + "\n"); + " - field name: " + + field_id.name() + + "\n" + " - field layout: " + + field_layout.to_string() + "\n"); EKAT_REQUIRE_MSG(field_layout.tags()[0] == COL, "Error! ZonalAvgDiag diagnostic expects a layout starting " "with the 'COL' tag.\n" - " - field name : " + field_id.name() + "\n" - " - field layout: " + field_layout.to_string() + "\n"); + " - field name : " + + field_id.name() + + "\n" + " - field layout: " + + field_layout.to_string() + "\n"); FieldLayout diagnostic_layout = - field_layout.clone().strip_dim(COL).prepend_dim({CMP}, {m_lat_num}, - {ZonalAvgDiag::dim_name}); - FieldIdentifier diagnostic_id(m_diag_name, diagnostic_layout, - field_id.get_units(), field_id.get_grid_name()); + field_layout.clone().strip_dim(COL).prepend_dim({CMP}, {m_num_zonal_bins}, {"bin"}); + FieldIdentifier diagnostic_id(m_diag_name, diagnostic_layout, field_id.get_units(), + field_id.get_grid_name()); m_diagnostic_output = Field(diagnostic_id); m_diagnostic_output.allocate_view(); // allocate zonal area const FieldIdentifier &area_id = m_scaled_area.get_header().get_identifier(); - FieldLayout zonal_area_layout({CMP}, {m_lat_num}, {ZonalAvgDiag::dim_name}); - FieldIdentifier zonal_area_id("zonal area", zonal_area_layout, - area_id.get_units(), area_id.get_grid_name()); + FieldLayout zonal_area_layout({CMP}, {m_num_zonal_bins}, {"bin"}); + FieldIdentifier zonal_area_id("zonal area", zonal_area_layout, area_id.get_units(), + area_id.get_grid_name()); Field zonal_area(zonal_area_id); zonal_area.allocate_view(); // compute zonal area FieldLayout ones_layout = area_id.get_layout().clone(); - FieldIdentifier ones_id("ones", ones_layout, area_id.get_units(), - area_id.get_grid_name()); + FieldIdentifier ones_id("ones", ones_layout, area_id.get_units(), area_id.get_grid_name()); Field ones(ones_id); ones.allocate_view(); ones.deep_copy(1.0); compute_zonal_sum(zonal_area, m_scaled_area, ones, m_lat, &m_comm); // scale area by 1 / zonal area - using RangePolicy = Kokkos::RangePolicy; - const Real lat_delta = sp(180.0) / m_lat_num; - const int ncols = field_layout.dim(0); - auto lat_view = m_lat.get_view(); - auto zonal_area_view = zonal_area.get_view(); + using RangePolicy = Kokkos::RangePolicy; + const Real lat_delta = sp(180.0) / m_num_zonal_bins; + const int ncols = field_layout.dim(0); + auto lat_view = m_lat.get_view(); + auto zonal_area_view = zonal_area.get_view(); auto scaled_area_view = m_scaled_area.get_view(); - Kokkos::parallel_for("scale_area_by_zonal_area_" + field.name(), - RangePolicy(0, ncols), KOKKOS_LAMBDA(const int& i) { + Kokkos::parallel_for( + "scale_area_by_zonal_area_" + field.name(), RangePolicy(0, ncols), + KOKKOS_LAMBDA(const int &i) { const int lat_i = (lat_view(i) + sp(90.0)) / lat_delta; scaled_area_view(i) /= zonal_area_view(lat_i); }); @@ -186,4 +186,4 @@ void ZonalAvgDiag::compute_diagnostic_impl() { compute_zonal_sum(m_diagnostic_output, field, m_scaled_area, m_lat, &m_comm); } -} // namespace scream +} // namespace scream diff --git a/components/eamxx/src/diagnostics/zonal_avg.hpp b/components/eamxx/src/diagnostics/zonal_avg.hpp index 827944d0e50e..6c00f419c59d 100644 --- a/components/eamxx/src/diagnostics/zonal_avg.hpp +++ b/components/eamxx/src/diagnostics/zonal_avg.hpp @@ -4,23 +4,26 @@ #include "share/atm_process/atmosphere_diagnostic.hpp" namespace scream { -// TODO: Update this comment /* - * This diagnostic will calculate the area-weighted average of a field - * across the COL tag dimension, producing an N-1 dimensional field - * that is area-weighted average of the input field. + * This diagnostic will calculate area-weighted zonal averages of a field across + * the COL tag dimension producing an N dimensional field, where the COL tag + * dimension is replaced by a CMP tag dimension named "bin" that indicates which + * zonal band the average value corresponds to. */ class ZonalAvgDiag : public AtmosphereDiagnostic { - // TODO: comment this, noting it's a utility function that could exist elsewhere - static void compute_zonal_sum(const Field &result, const Field &field, - const Field &weight, const Field &lat, const ekat::Comm *comm = nullptr); - - public: - - inline static const std::string dim_name = "bin"; - + // Utility to compute the contraction of a field along its column dimension. + // This is equivalent to f_out = einsum('i,i...k->...k', weight, f_in). + // The implementation is such that: + // - all Field objects must be allocated + // - the first dimension for field, weight, and lat is for the columns (COL) + // - the first dimension for result is for the zonal bins (CMP,"bin") + // - field and result must be the same dimension, up to 3 + static void compute_zonal_sum(const Field &result, const Field &field, const Field &weight, + const Field &lat, const ekat::Comm *comm = nullptr); + +public: // Constructors ZonalAvgDiag(const ekat::Comm &comm, const ekat::ParameterList ¶ms); @@ -30,24 +33,22 @@ class ZonalAvgDiag : public AtmosphereDiagnostic { // Set the grid void set_grids(const std::shared_ptr grids_manager); - protected: +protected: #ifdef KOKKOS_ENABLE_CUDA - public: +public: #endif void compute_diagnostic_impl(); - protected: +protected: void initialize_impl(const RunType /*run_type*/); std::string m_diag_name; + int m_num_zonal_bins; Field m_lat; - int m_lat_num; - Field m_scaled_area; - }; -} // namespace scream +} // namespace scream -#endif // EAMXX_ZONAL_AVERAGE_HPP +#endif // EAMXX_ZONAL_AVERAGE_HPP diff --git a/components/eamxx/src/share/field/field_layout.hpp b/components/eamxx/src/share/field/field_layout.hpp index 4254712486b3..fd4a665c32fe 100644 --- a/components/eamxx/src/share/field/field_layout.hpp +++ b/components/eamxx/src/share/field/field_layout.hpp @@ -59,10 +59,6 @@ inline std::string e2str (const LayoutType lt) { */ class FieldLayout { - - void add_dim (const FieldTag t, const int extent, const std::string& name, - bool prepend_instead_of_append = false); - public: using extents_type = typename KokkosTypes::view_1d; @@ -159,6 +155,9 @@ class FieldLayout { protected: void compute_type (); void set_extents (); + void add_dim (const FieldTag t, const int extent, const std::string& name, + bool prepend_instead_of_append = false); + int m_rank; std::vector m_tags; diff --git a/components/eamxx/src/share/io/eamxx_io_utils.cpp b/components/eamxx/src/share/io/eamxx_io_utils.cpp index 7ad87ee06a00..2ef1cd7a2204 100644 --- a/components/eamxx/src/share/io/eamxx_io_utils.cpp +++ b/components/eamxx/src/share/io/eamxx_io_utils.cpp @@ -140,6 +140,7 @@ create_diagnostic (const std::string& diag_field_name, std::regex vert_layer ("(z|geopotential|height)_(mid|int)$"); std::regex horiz_avg ("([A-Za-z0-9_]+)_horiz_avg$"); std::regex vert_contract ("([A-Za-z0-9_]+)_vert_(avg|sum)(_((dp|dz)_weighted))?$"); + std::regex zonal_avg (R"(([A-Za-z0-9_]+)_zonal_avg_with_(\d+)_bins$)"); std::string diag_name; std::smatch matches; @@ -210,6 +211,12 @@ create_diagnostic (const std::string& diag_field_name, params.set("weighting_method", matches[5].str()); } } + else if (std::regex_search(diag_field_name,matches,zonal_avg)) { + diag_name = "ZonalAvgDiag"; + params.set("grid_name", grid->name()); + params.set("field_name", matches[1].str()); + params.set("number_of_zonal_bins", matches[2].str()); + } else { // No existing special regex matches, so we assume that the diag field name IS the diag name. diff --git a/components/eamxx/tests/multi-process/dynamics_physics/homme_shoc_cld_p3_rrtmgp_pg2/output.yaml b/components/eamxx/tests/multi-process/dynamics_physics/homme_shoc_cld_p3_rrtmgp_pg2/output.yaml index f948889d2a55..17720a133639 100644 --- a/components/eamxx/tests/multi-process/dynamics_physics/homme_shoc_cld_p3_rrtmgp_pg2/output.yaml +++ b/components/eamxx/tests/multi-process/dynamics_physics/homme_shoc_cld_p3_rrtmgp_pg2/output.yaml @@ -101,6 +101,8 @@ fields: - MeridionalVapFlux - PotentialTemperature_at_model_top - PotentialTemperature_at_500mb + - T_mid_zonal_avg_with_180_bins + - T_mid_zonal_avg_with_90_bins # GLL output for homme states. These # represent all current possible homme # states available. From 2d55ab766d1ebc22b0c665256afac9728d6de162 Mon Sep 17 00:00:00 2001 From: Chris Vogl Date: Wed, 21 May 2025 10:14:05 -0700 Subject: [PATCH 5/8] switched from using KOKKOS_LAMBDA for reductions in EAMxx zonal average diagnostic --- components/eamxx/src/diagnostics/zonal_avg.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/components/eamxx/src/diagnostics/zonal_avg.cpp b/components/eamxx/src/diagnostics/zonal_avg.cpp index 17219ddc9443..d7c453ecbdea 100644 --- a/components/eamxx/src/diagnostics/zonal_avg.cpp +++ b/components/eamxx/src/diagnostics/zonal_avg.cpp @@ -29,7 +29,7 @@ void ZonalAvgDiag::compute_zonal_sum(const Field &result, const Field &field, co const Real lat_upper = lat_lower + lat_delta; Kokkos::parallel_reduce( Kokkos::TeamVectorRange(tm, ncols), - KOKKOS_LAMBDA(int i, Real &val) { + [&](int i, Real &val) { // TODO: check if tenary is ok here (if not, multiply by flag) int flag = (lat_lower <= lat_view(i)) && (lat_view(i) < lat_upper); val += flag ? weight_view(i) * field_view(i) : sp(0.0); @@ -51,7 +51,7 @@ void ZonalAvgDiag::compute_zonal_sum(const Field &result, const Field &field, co const Real lat_upper = lat_lower + lat_delta; Kokkos::parallel_reduce( Kokkos::TeamVectorRange(tm, ncols), - KOKKOS_LAMBDA(int i, Real &val) { + [&](int i, Real &val) { int flag = (lat_lower <= lat_view(i)) && (lat_view(i) < lat_upper); // TODO: check if tenary is ok here (if not, multiply by flag) val += flag ? weight_view(i) * field_view(i, d1_i) : sp(0.0); @@ -76,7 +76,7 @@ void ZonalAvgDiag::compute_zonal_sum(const Field &result, const Field &field, co const Real lat_upper = lat_lower + lat_delta; Kokkos::parallel_reduce( Kokkos::TeamVectorRange(tm, ncols), - KOKKOS_LAMBDA(int i, Real &val) { + [&](int i, Real &val) { int flag = (lat_lower <= lat_view(i)) && (lat_view(i) < lat_upper); // TODO: check if tenary is ok here (if not, multiply by flag) val += flag ? weight_view(i) * field_view(i, d1_i, d2_i) : sp(0.0); From 23a74dff4ffa7fa66689e9e9f66152ea401d1c66 Mon Sep 17 00:00:00 2001 From: mahf708 Date: Wed, 21 May 2025 19:10:28 +0000 Subject: [PATCH 6/8] EAMxx: public methods in ZonalAvgDiag, fix tests --- .../src/diagnostics/tests/zonal_avg_test.cpp | 29 +++++++++---------- .../eamxx/src/diagnostics/zonal_avg.hpp | 25 ++++++++-------- .../eamxx/src/share/io/scorpio_output.cpp | 8 +++-- .../homme_shoc_cld_p3_rrtmgp_pg2/output.yaml | 2 -- 4 files changed, 32 insertions(+), 32 deletions(-) diff --git a/components/eamxx/src/diagnostics/tests/zonal_avg_test.cpp b/components/eamxx/src/diagnostics/tests/zonal_avg_test.cpp index 21e5317c6b85..0fcfb21f9333 100644 --- a/components/eamxx/src/diagnostics/tests/zonal_avg_test.cpp +++ b/components/eamxx/src/diagnostics/tests/zonal_avg_test.cpp @@ -28,10 +28,6 @@ std::shared_ptr create_gm(const ekat::Comm &comm, const int ncols, TEST_CASE("zonal_avg") { using namespace ShortFieldTagsNames; using namespace ekat::units; - using TeamPolicy = Kokkos::TeamPolicy; - using TeamMember = typename TeamPolicy::member_type; - using KT = ekat::KokkosTypes; - using ESU = ekat::ExeSpaceUtils; // A numerical tolerance auto tol = std::numeric_limits::epsilon() * 100; @@ -51,19 +47,20 @@ TEST_CASE("zonal_avg") { auto gm = create_gm(comm, ngcols, nlevs); auto grid = gm->get_grid("Physics"); - Field area = grid->get_geometry_data("area"); - auto area_view = area.get_view(); + Field area = grid->get_geometry_data("area"); + auto area_view_h = area.get_view(); // Set latitude values Field lat = gm->get_grid_nonconst("Physics")->create_geometry_data( "lat", grid->get_2d_scalar_layout(), Units::nondimensional()); - auto lat_view = lat.get_view(); + auto lat_view_h = lat.get_view(); const Real lat_delta = sp(180.0) / nlats; std::vector zonal_areas(nlats, 0.0); for (int i = 0; i < ngcols; i++) { - lat_view(i) = sp(-90.0) + (i % nlats + sp(0.5)) * lat_delta; - zonal_areas[i % nlats] += area_view[i]; + lat_view_h(i) = sp(-90.0) + (i % nlats + sp(0.5)) * lat_delta; + zonal_areas[i % nlats] += area_view_h[i]; } + lat.sync_to_dev(); // Input (randomized) qc FieldLayout scalar1d_layout{{COL}, {ngcols}}; @@ -129,12 +126,13 @@ TEST_CASE("zonal_avg") { diag0_field.allocate_view(); // calculate the zonal average - auto qc1_view = qc1.get_view(); - auto diag0_view = diag0_field.get_view(); + auto qc1_view_h = qc1.get_view(); + auto diag0_view_h = diag0_field.get_view(); for (int i = 0; i < ngcols; i++) { const int nlat = i % nlats; - diag0_view(nlat) += area_view(i) / zonal_areas[nlat] * qc1_view(i); + diag0_view_h(nlat) += area_view_h(i) / zonal_areas[nlat] * qc1_view_h(i); } + diag0_field.sync_to_dev(); // Compare REQUIRE(views_are_equal(diag1_field, diag0_field)); @@ -171,16 +169,17 @@ TEST_CASE("zonal_avg") { FieldIdentifier diag3m_id("qc_zonal_avg_manual", diag3m_layout, kg / kg, grid->name()); Field diag3m_field(diag3m_id); diag3m_field.allocate_view(); - auto qc3_view = qc3.get_view(); - auto diag3m_view = diag3m_field.get_view(); + auto qc3_view_h = qc3.get_view(); + auto diag3m_view_h = diag3m_field.get_view(); for (int i = 0; i < ngcols; i++) { const int nlat = i % nlats; for (int j = 0; j < dim3; j++) { for (int k = 0; k < nlevs; k++) { - diag3m_view(nlat, j, k) += area_view(i) / zonal_areas[nlat] * qc3_view(i, j, k); + diag3m_view_h(nlat, j, k) += area_view_h(i) / zonal_areas[nlat] * qc3_view_h(i, j, k); } } } + diag3m_field.sync_to_dev(); diag3->set_required_field(qc3); diag3->initialize(t0, RunType::Initial); diag3->compute_diagnostic(); diff --git a/components/eamxx/src/diagnostics/zonal_avg.hpp b/components/eamxx/src/diagnostics/zonal_avg.hpp index 6c00f419c59d..19ca5f66b8e5 100644 --- a/components/eamxx/src/diagnostics/zonal_avg.hpp +++ b/components/eamxx/src/diagnostics/zonal_avg.hpp @@ -13,16 +13,6 @@ namespace scream { class ZonalAvgDiag : public AtmosphereDiagnostic { - // Utility to compute the contraction of a field along its column dimension. - // This is equivalent to f_out = einsum('i,i...k->...k', weight, f_in). - // The implementation is such that: - // - all Field objects must be allocated - // - the first dimension for field, weight, and lat is for the columns (COL) - // - the first dimension for result is for the zonal bins (CMP,"bin") - // - field and result must be the same dimension, up to 3 - static void compute_zonal_sum(const Field &result, const Field &field, const Field &weight, - const Field &lat, const ekat::Comm *comm = nullptr); - public: // Constructors ZonalAvgDiag(const ekat::Comm &comm, const ekat::ParameterList ¶ms); @@ -37,11 +27,22 @@ class ZonalAvgDiag : public AtmosphereDiagnostic { #ifdef KOKKOS_ENABLE_CUDA public: #endif + void initialize_impl(const RunType /*run_type*/); void compute_diagnostic_impl(); -protected: - void initialize_impl(const RunType /*run_type*/); + // TODO: make it a local function in the cpp file + // Utility to compute the contraction of a field along its column dimension. + // This is equivalent to f_out = einsum('i,i...k->...k', weight, f_in). + // The implementation is such that: + // - all Field objects must be allocated + // - the first dimension for field, weight, and lat is for the columns (COL) + // - the first dimension for result is for the zonal bins (CMP,"bin") + // - field and result must be the same dimension, up to 3 + // TODO: make it a local function in the cpp file + static void compute_zonal_sum(const Field &result, const Field &field, const Field &weight, + const Field &lat, const ekat::Comm *comm = nullptr); +protected: std::string m_diag_name; int m_num_zonal_bins; diff --git a/components/eamxx/src/share/io/scorpio_output.cpp b/components/eamxx/src/share/io/scorpio_output.cpp index d2751894b8ab..c7bc76723ea6 100644 --- a/components/eamxx/src/share/io/scorpio_output.cpp +++ b/components/eamxx/src/share/io/scorpio_output.cpp @@ -790,7 +790,8 @@ void AtmosphereOutput::register_dimensions(const std::string& name) // If t==CMP, and the name stored in the layout is the default ("dim"), // we append also the extent, to allow different vector dims in the file - tag_name += tags[i] == CMP ? std::to_string(dims[i]) : ""; + // TODO: generalie this to all tags, for now hardcoding to dim and bin only + tag_name += (tag_name == "dim" or tag_name=="bin") ? std::to_string(dims[i]) : ""; auto is_partitioned = m_io_grid->get_partitioned_dim_tag()==tags[i]; int dim_len = is_partitioned @@ -887,7 +888,8 @@ void AtmosphereOutput::set_avg_cnt_tracking(const std::string& name, const Field // If t==CMP, and the name stored in the layout is the default ("dim"), // we append also the extent, to allow different vector dims in the file - tag_name += t==CMP ? std::to_string(layout.dim(i)) : ""; + // TODO: generalize this to all tags, for now hardcoding to dim and bin only + tag_name += (tag_name=="dim" or tag_name=="bin") ? std::to_string(layout.dim(i)) : ""; avg_cnt_name += "_" + tag_name; } @@ -952,7 +954,7 @@ register_variables(const std::string& filename, auto tag_name = m_io_grid->has_special_tag_name(t) ? m_io_grid->get_special_tag_name(t) : layout.names()[i]; - if (t==CMP) { + if (tag_name=="dim" or tag_name=="bin") { tag_name += std::to_string(layout.dim(i)); } vec_of_dims.push_back(tag_name); // Add dimensions string to vector of dims. diff --git a/components/eamxx/tests/multi-process/dynamics_physics/homme_shoc_cld_p3_rrtmgp_pg2/output.yaml b/components/eamxx/tests/multi-process/dynamics_physics/homme_shoc_cld_p3_rrtmgp_pg2/output.yaml index 17720a133639..f948889d2a55 100644 --- a/components/eamxx/tests/multi-process/dynamics_physics/homme_shoc_cld_p3_rrtmgp_pg2/output.yaml +++ b/components/eamxx/tests/multi-process/dynamics_physics/homme_shoc_cld_p3_rrtmgp_pg2/output.yaml @@ -101,8 +101,6 @@ fields: - MeridionalVapFlux - PotentialTemperature_at_model_top - PotentialTemperature_at_500mb - - T_mid_zonal_avg_with_180_bins - - T_mid_zonal_avg_with_90_bins # GLL output for homme states. These # represent all current possible homme # states available. From 055a241630e27781f54f8f5bbce4c28395ee8383 Mon Sep 17 00:00:00 2001 From: mahf708 Date: Fri, 23 May 2025 21:02:24 -0700 Subject: [PATCH 7/8] EAMxx: add docs for zonal avg diag --- .../docs/user/diags/field_contraction.md | 23 +++++++++++++++++++ .../eamxx/src/share/io/eamxx_io_utils.cpp | 11 +++++++-- 2 files changed, 32 insertions(+), 2 deletions(-) diff --git a/components/eamxx/docs/user/diags/field_contraction.md b/components/eamxx/docs/user/diags/field_contraction.md index abc23a18084b..9e019a56975b 100644 --- a/components/eamxx/docs/user/diags/field_contraction.md +++ b/components/eamxx/docs/user/diags/field_contraction.md @@ -3,6 +3,8 @@ In EAMxx, we can automatically calculate field reductions across the horizontal columns and across the model vertical levels. We call these horizontal and vertical reductions. +We can also automatically calculate zonal averages, +albeit with additional flexibility (see below). ## Horizontal reduction @@ -58,6 +60,25 @@ The supported weighting options for now are In the case of `pseudo_density`, the weighting is scaled by 1/g, where g is the gravitational acceleration, in units of m/s$^2$. +## Zonal reduction + +We currently have a utility to calculate zonal averages online. +To select the zonal average, you only need to suffix +a field name `X` with `_zonal_avg` and optionally the +number of bins `num_lat`. All zonal averages are calculated +using the area fraction in each bin as the weight. + +For 180 latitude bins (the default), the bins are defined +as follows: [-90, -89), [-89, -88), ..., [89, 90). +For 90 latitude bins, the bins are defined as follows: +[-90, -88), [-88, -85), ..., [88, 90). +And so on... + +| Reduction | Weight | Description | +| --------- | ------ | ----------- | +| `X_zonal_avg` | Area fraction | Average across 180 latitude bins | +| `X_zonal_avg_with_Y_bins` | Area fraction | Average across Y latitude bins | + ## Example ```yaml @@ -77,6 +98,8 @@ fields: - T_mid_vert_sum_dz_weighted # K * m - T_mid_vert_avg # K - T_mid_vert_sum # K + - T_mid_zonal_avg # K + - T_mid_zonal_avg_with_90_bins # K output_control: frequency: 1 frequency_units: nmonths diff --git a/components/eamxx/src/share/io/eamxx_io_utils.cpp b/components/eamxx/src/share/io/eamxx_io_utils.cpp index 2ef1cd7a2204..4fd4e091ee03 100644 --- a/components/eamxx/src/share/io/eamxx_io_utils.cpp +++ b/components/eamxx/src/share/io/eamxx_io_utils.cpp @@ -140,7 +140,7 @@ create_diagnostic (const std::string& diag_field_name, std::regex vert_layer ("(z|geopotential|height)_(mid|int)$"); std::regex horiz_avg ("([A-Za-z0-9_]+)_horiz_avg$"); std::regex vert_contract ("([A-Za-z0-9_]+)_vert_(avg|sum)(_((dp|dz)_weighted))?$"); - std::regex zonal_avg (R"(([A-Za-z0-9_]+)_zonal_avg_with_(\d+)_bins$)"); + std::regex zonal_avg (R"(([A-Za-z0-9_]+)_zonal_avg(_with_(\d+)_bins)?$)"); std::string diag_name; std::smatch matches; @@ -215,7 +215,14 @@ create_diagnostic (const std::string& diag_field_name, diag_name = "ZonalAvgDiag"; params.set("grid_name", grid->name()); params.set("field_name", matches[1].str()); - params.set("number_of_zonal_bins", matches[2].str()); + // The second match is optional + if (matches[2].matched) { + // note that the 3rd match is the number of bins + params.set("number_of_zonal_bins", matches[3].str()); + } else { + // set number_of_zonal_bins to default 180 + params.set("number_of_zonal_bins", "180"); + } } else { From ad24954eac49f8a2e6b76e548f97655b22d50ff2 Mon Sep 17 00:00:00 2001 From: mahf708 Date: Wed, 28 May 2025 08:56:56 -0700 Subject: [PATCH 8/8] EAMxx: force users to specify lat bins --- .../docs/user/diags/field_contraction.md | 20 +++++++-------- .../src/diagnostics/tests/zonal_avg_test.cpp | 25 ++++++++++++------- .../eamxx/src/diagnostics/zonal_avg.cpp | 2 +- .../eamxx/src/share/io/eamxx_io_utils.cpp | 11 ++------ 4 files changed, 28 insertions(+), 30 deletions(-) diff --git a/components/eamxx/docs/user/diags/field_contraction.md b/components/eamxx/docs/user/diags/field_contraction.md index 9e019a56975b..935e9d93855c 100644 --- a/components/eamxx/docs/user/diags/field_contraction.md +++ b/components/eamxx/docs/user/diags/field_contraction.md @@ -3,8 +3,7 @@ In EAMxx, we can automatically calculate field reductions across the horizontal columns and across the model vertical levels. We call these horizontal and vertical reductions. -We can also automatically calculate zonal averages, -albeit with additional flexibility (see below). +We can also automatically calculate zonal averages. ## Horizontal reduction @@ -63,21 +62,20 @@ where g is the gravitational acceleration, in units of m/s$^2$. ## Zonal reduction We currently have a utility to calculate zonal averages online. -To select the zonal average, you only need to suffix -a field name `X` with `_zonal_avg` and optionally the -number of bins `num_lat`. All zonal averages are calculated +To select the zonal average, you need to suffix +a field name `X` with `_zonal_avg` and the +number of bins `Y` as `_Y_bins`. All zonal averages are calculated using the area fraction in each bin as the weight. -For 180 latitude bins (the default), the bins are defined +For 180 latitude bins, the bins are defined as follows: [-90, -89), [-89, -88), ..., [89, 90). For 90 latitude bins, the bins are defined as follows: -[-90, -88), [-88, -85), ..., [88, 90). +[-90, -88), [-88, -86), ..., [88, 90). And so on... | Reduction | Weight | Description | | --------- | ------ | ----------- | -| `X_zonal_avg` | Area fraction | Average across 180 latitude bins | -| `X_zonal_avg_with_Y_bins` | Area fraction | Average across Y latitude bins | +| `X_zonal_avg_Y_bins` | Area fraction | Average across the zonal direction | ## Example @@ -98,8 +96,8 @@ fields: - T_mid_vert_sum_dz_weighted # K * m - T_mid_vert_avg # K - T_mid_vert_sum # K - - T_mid_zonal_avg # K - - T_mid_zonal_avg_with_90_bins # K + - T_mid_zonal_avg_180_bins # K + - T_mid_zonal_avg_90_bins # K output_control: frequency: 1 frequency_units: nmonths diff --git a/components/eamxx/src/diagnostics/tests/zonal_avg_test.cpp b/components/eamxx/src/diagnostics/tests/zonal_avg_test.cpp index 0fcfb21f9333..3e2150a65bfc 100644 --- a/components/eamxx/src/diagnostics/tests/zonal_avg_test.cpp +++ b/components/eamxx/src/diagnostics/tests/zonal_avg_test.cpp @@ -84,15 +84,6 @@ TEST_CASE("zonal_avg") { RPDF pdf(sp(0.0), sp(200.0)); auto engine = scream::setup_random_test(); - // Construct the Diagnostics - std::map> diags; - auto &diag_factory = AtmosphereDiagnosticFactory::instance(); - register_diagnostics(); - - ekat::ParameterList params; - REQUIRE_THROWS(diag_factory.create("ZonalAvgDiag", comm, - params)); // No 'field_name' parameter - // Set time for qc and randomize its values qc1.get_header().get_tracking().update_time_stamp(t0); qc2.get_header().get_tracking().update_time_stamp(t0); @@ -101,10 +92,26 @@ TEST_CASE("zonal_avg") { randomize(qc2, engine, pdf); randomize(qc3, engine, pdf); + // Construct the Diagnostics + std::map> diags; + auto &diag_factory = AtmosphereDiagnosticFactory::instance(); + register_diagnostics(); + // Create and set up the diagnostic + ekat::ParameterList params; + REQUIRE_THROWS(diag_factory.create("ZonalAvgDiag", comm, + params)); // Bad construction + params.set("grid_name", grid->name()); + REQUIRE_THROWS(diag_factory.create("ZonalAvgDiag", comm, + params)); // Still no field_name + params.set("field_name", "qc"); + REQUIRE_THROWS(diag_factory.create("ZonalAvgDiag", comm, + params)); // Still no number_of_zonal_bins + params.set("number_of_zonal_bins", std::to_string(nlats)); + // Now we should be good to go... auto diag1 = diag_factory.create("ZonalAvgDiag", comm, params); auto diag2 = diag_factory.create("ZonalAvgDiag", comm, params); auto diag3 = diag_factory.create("ZonalAvgDiag", comm, params); diff --git a/components/eamxx/src/diagnostics/zonal_avg.cpp b/components/eamxx/src/diagnostics/zonal_avg.cpp index d7c453ecbdea..ba5a3211f5d7 100644 --- a/components/eamxx/src/diagnostics/zonal_avg.cpp +++ b/components/eamxx/src/diagnostics/zonal_avg.cpp @@ -104,7 +104,7 @@ ZonalAvgDiag::ZonalAvgDiag(const ekat::Comm &comm, const ekat::ParameterList &pa : AtmosphereDiagnostic(comm, params) { const auto &field_name = m_params.get("field_name"); const auto &num_bins_value = params.get("number_of_zonal_bins"); - m_diag_name = field_name + "_zonal_avg_with_" + num_bins_value + "_bins"; + m_diag_name = field_name + "_zonal_avg_" + num_bins_value + "_bins"; m_num_zonal_bins = std::stoi(num_bins_value); } diff --git a/components/eamxx/src/share/io/eamxx_io_utils.cpp b/components/eamxx/src/share/io/eamxx_io_utils.cpp index 4fd4e091ee03..0639150e910b 100644 --- a/components/eamxx/src/share/io/eamxx_io_utils.cpp +++ b/components/eamxx/src/share/io/eamxx_io_utils.cpp @@ -140,7 +140,7 @@ create_diagnostic (const std::string& diag_field_name, std::regex vert_layer ("(z|geopotential|height)_(mid|int)$"); std::regex horiz_avg ("([A-Za-z0-9_]+)_horiz_avg$"); std::regex vert_contract ("([A-Za-z0-9_]+)_vert_(avg|sum)(_((dp|dz)_weighted))?$"); - std::regex zonal_avg (R"(([A-Za-z0-9_]+)_zonal_avg(_with_(\d+)_bins)?$)"); + std::regex zonal_avg (R"(([A-Za-z0-9_]+)_zonal_avg_(\d+)_bins$)"); std::string diag_name; std::smatch matches; @@ -215,14 +215,7 @@ create_diagnostic (const std::string& diag_field_name, diag_name = "ZonalAvgDiag"; params.set("grid_name", grid->name()); params.set("field_name", matches[1].str()); - // The second match is optional - if (matches[2].matched) { - // note that the 3rd match is the number of bins - params.set("number_of_zonal_bins", matches[3].str()); - } else { - // set number_of_zonal_bins to default 180 - params.set("number_of_zonal_bins", "180"); - } + params.set("number_of_zonal_bins", matches[2].str()); } else {