|
| 1 | +#include "catch2/catch.hpp" |
| 2 | +#include "diagnostics/register_diagnostics.hpp" |
| 3 | +#include "share/field/field_utils.hpp" |
| 4 | +#include "share/grid/mesh_free_grids_manager.hpp" |
| 5 | +#include "share/util/scream_setup_random_test.hpp" |
| 6 | +#include "share/util/scream_universal_constants.hpp" |
| 7 | + |
| 8 | +namespace scream { |
| 9 | + |
| 10 | +std::shared_ptr<GridsManager> create_gm(const ekat::Comm &comm, const int ncols, |
| 11 | + const int nlevs) { |
| 12 | + const int num_global_cols = ncols * comm.size(); |
| 13 | + |
| 14 | + using vos_t = std::vector<std::string>; |
| 15 | + ekat::ParameterList gm_params; |
| 16 | + gm_params.set("grids_names", vos_t{"Point Grid"}); |
| 17 | + auto &pl = gm_params.sublist("Point Grid"); |
| 18 | + pl.set<std::string>("type", "point_grid"); |
| 19 | + pl.set("aliases", vos_t{"Physics"}); |
| 20 | + pl.set<int>("number_of_global_columns", num_global_cols); |
| 21 | + pl.set<int>("number_of_vertical_levels", nlevs); |
| 22 | + |
| 23 | + auto gm = create_mesh_free_grids_manager(comm, gm_params); |
| 24 | + gm->build_grids(); |
| 25 | + |
| 26 | + return gm; |
| 27 | +} |
| 28 | + |
| 29 | +TEST_CASE("horiz_avg") { |
| 30 | + using namespace ShortFieldTagsNames; |
| 31 | + using namespace ekat::units; |
| 32 | + using TeamPolicy = Kokkos::TeamPolicy<Field::device_t::execution_space>; |
| 33 | + using TeamMember = typename TeamPolicy::member_type; |
| 34 | + using KT = ekat::KokkosTypes<DefaultDevice>; |
| 35 | + using ESU = ekat::ExeSpaceUtils<typename KT::ExeSpace>; |
| 36 | + |
| 37 | + // A numerical tolerance |
| 38 | + auto tol = std::numeric_limits<Real>::epsilon() * 100; |
| 39 | + |
| 40 | + // A world comm |
| 41 | + ekat::Comm comm(MPI_COMM_WORLD); |
| 42 | + |
| 43 | + // A time stamp |
| 44 | + util::TimeStamp t0({2024, 1, 1}, {0, 0, 0}); |
| 45 | + |
| 46 | + // Create a grids manager - single column for these tests |
| 47 | + constexpr int nlevs = 3; |
| 48 | + constexpr int dim3 = 4; |
| 49 | + const int ngcols = 6 * comm.size(); |
| 50 | + |
| 51 | + auto gm = create_gm(comm, ngcols, nlevs); |
| 52 | + auto grid = gm->get_grid("Physics"); |
| 53 | + |
| 54 | + // Input (randomized) qc |
| 55 | + FieldLayout scalar1d_layout{{COL}, {ngcols}}; |
| 56 | + FieldLayout scalar2d_layout{{COL, LEV}, {ngcols, nlevs}}; |
| 57 | + FieldLayout scalar3d_layout{{COL, CMP, LEV}, {ngcols, dim3, nlevs}}; |
| 58 | + |
| 59 | + FieldIdentifier qc1_fid("qc", scalar1d_layout, kg / kg, grid->name()); |
| 60 | + FieldIdentifier qc2_fid("qc", scalar2d_layout, kg / kg, grid->name()); |
| 61 | + FieldIdentifier qc3_fid("qc", scalar3d_layout, kg / kg, grid->name()); |
| 62 | + |
| 63 | + Field qc1(qc1_fid); |
| 64 | + Field qc2(qc2_fid); |
| 65 | + Field qc3(qc3_fid); |
| 66 | + |
| 67 | + qc1.allocate_view(); |
| 68 | + qc2.allocate_view(); |
| 69 | + qc3.allocate_view(); |
| 70 | + |
| 71 | + // Construct random number generator stuff |
| 72 | + using RPDF = std::uniform_real_distribution<Real>; |
| 73 | + RPDF pdf(sp(0.0), sp(200.0)); |
| 74 | + |
| 75 | + auto engine = scream::setup_random_test(); |
| 76 | + |
| 77 | + // Construct the Diagnostics |
| 78 | + std::map<std::string, std::shared_ptr<AtmosphereDiagnostic>> diags; |
| 79 | + auto &diag_factory = AtmosphereDiagnosticFactory::instance(); |
| 80 | + register_diagnostics(); |
| 81 | + |
| 82 | + ekat::ParameterList params; |
| 83 | + REQUIRE_THROWS(diag_factory.create("HorizAvgDiag", comm, |
| 84 | + params)); // No 'field_name' parameter |
| 85 | + |
| 86 | + // Set time for qc and randomize its values |
| 87 | + qc1.get_header().get_tracking().update_time_stamp(t0); |
| 88 | + qc2.get_header().get_tracking().update_time_stamp(t0); |
| 89 | + qc3.get_header().get_tracking().update_time_stamp(t0); |
| 90 | + randomize(qc1, engine, pdf); |
| 91 | + randomize(qc2, engine, pdf); |
| 92 | + randomize(qc3, engine, pdf); |
| 93 | + |
| 94 | + // Create and set up the diagnostic |
| 95 | + params.set("grid_name", grid->name()); |
| 96 | + params.set<std::string>("field_name", "qc"); |
| 97 | + auto diag1 = diag_factory.create("HorizAvgDiag", comm, params); |
| 98 | + auto diag2 = diag_factory.create("HorizAvgDiag", comm, params); |
| 99 | + auto diag3 = diag_factory.create("HorizAvgDiag", comm, params); |
| 100 | + diag1->set_grids(gm); |
| 101 | + diag2->set_grids(gm); |
| 102 | + diag3->set_grids(gm); |
| 103 | + |
| 104 | + auto area = grid->get_geometry_data("area"); |
| 105 | + |
| 106 | + diag1->set_required_field(qc1); |
| 107 | + diag1->initialize(t0, RunType::Initial); |
| 108 | + |
| 109 | + diag1->compute_diagnostic(); |
| 110 | + auto diag1_f = diag1->get_diagnostic(); |
| 111 | + |
| 112 | + FieldIdentifier diag0_fid("qc_horiz_avg_manual", |
| 113 | + scalar1d_layout.clone().strip_dim(COL), kg / kg, |
| 114 | + grid->name()); |
| 115 | + Field diag0(diag0_fid); |
| 116 | + diag0.allocate_view(); |
| 117 | + auto diag0_v = diag0.get_view<Real>(); |
| 118 | + |
| 119 | + auto qc1_v = qc1.get_view<Real *>(); |
| 120 | + auto area_v = area.get_view<const Real *>(); |
| 121 | + |
| 122 | + // calculate total area |
| 123 | + Real atot = field_sum<Real>(area, &comm); |
| 124 | + // calculate weighted avg |
| 125 | + Real wavg = sp(0.0); |
| 126 | + Kokkos::parallel_reduce( |
| 127 | + "HorizAvgDiag::compute_diagnostic_impl::weighted_sum", ngcols, |
| 128 | + KOKKOS_LAMBDA(const int icol, Real &local_wavg) { |
| 129 | + local_wavg += (area_v[icol] / atot) * qc1_v[icol]; |
| 130 | + }, |
| 131 | + wavg); |
| 132 | + Kokkos::deep_copy(diag0_v, wavg); |
| 133 | + |
| 134 | + diag1_f.sync_to_host(); |
| 135 | + auto diag1_v_h = diag1_f.get_view<Real, Host>(); |
| 136 | + REQUIRE(diag1_v_h() == wavg); |
| 137 | + |
| 138 | + // Try known cases |
| 139 | + // Set qc1_v to 1.0 to get weighted average of 1.0 |
| 140 | + wavg = sp(1.0); |
| 141 | + Kokkos::deep_copy(qc1_v, wavg); |
| 142 | + diag1->compute_diagnostic(); |
| 143 | + auto diag1_v2_host = diag1_f.get_view<Real, Host>(); |
| 144 | + REQUIRE_THAT(diag1_v2_host(), |
| 145 | + Catch::Matchers::WithinRel( |
| 146 | + wavg, tol)); // Catch2's floating point comparison |
| 147 | + |
| 148 | + // other diags |
| 149 | + // Set qc2_v to 5.0 to get weighted average of 5.0 |
| 150 | + wavg = sp(5.0); |
| 151 | + auto qc2_v = qc2.get_view<Real **>(); |
| 152 | + Kokkos::deep_copy(qc2_v, wavg); |
| 153 | + |
| 154 | + diag2->set_required_field(qc2); |
| 155 | + diag2->initialize(t0, RunType::Initial); |
| 156 | + diag2->compute_diagnostic(); |
| 157 | + auto diag2_f = diag2->get_diagnostic(); |
| 158 | + |
| 159 | + auto diag2_v_host = diag2_f.get_view<Real *, Host>(); |
| 160 | + |
| 161 | + for(int i = 0; i < nlevs; ++i) { |
| 162 | + REQUIRE_THAT(diag2_v_host(i), Catch::Matchers::WithinRel(wavg, tol)); |
| 163 | + } |
| 164 | + |
| 165 | + auto qc3_v = qc3.get_view<Real ***>(); |
| 166 | + FieldIdentifier diag3_manual_fid("qc_horiz_avg_manual", |
| 167 | + scalar3d_layout.clone().strip_dim(COL), |
| 168 | + kg / kg, grid->name()); |
| 169 | + Field diag3_manual(diag3_manual_fid); |
| 170 | + diag3_manual.allocate_view(); |
| 171 | + auto diag3_manual_v = diag3_manual.get_view<Real **>(); |
| 172 | + // calculate diag3_manual by hand |
| 173 | + auto p = ESU::get_default_team_policy(dim3 * nlevs, ngcols); |
| 174 | + Kokkos::parallel_for( |
| 175 | + "HorizAvgDiag::compute_diagnostic_impl::manual_diag3", p, |
| 176 | + KOKKOS_LAMBDA(const TeamMember &m) { |
| 177 | + const int idx = m.league_rank(); |
| 178 | + const int j = idx / nlevs; |
| 179 | + const int k = idx % nlevs; |
| 180 | + Real sum = sp(0.0); |
| 181 | + Kokkos::parallel_reduce( |
| 182 | + Kokkos::TeamThreadRange(m, ngcols), |
| 183 | + [&](const int icol, Real &accum) { |
| 184 | + accum += (area_v(icol) / atot) * qc3_v(icol, j, k); |
| 185 | + }, |
| 186 | + sum); |
| 187 | + Kokkos::single(Kokkos::PerTeam(m), |
| 188 | + [&]() { diag3_manual_v(j, k) = sum; }); |
| 189 | + }); |
| 190 | + diag3->set_required_field(qc3); |
| 191 | + diag3->initialize(t0, RunType::Initial); |
| 192 | + diag3->compute_diagnostic(); |
| 193 | + auto diag3_f = diag3->get_diagnostic(); |
| 194 | + REQUIRE(views_are_equal(diag3_f, diag3_manual)); |
| 195 | +} |
| 196 | + |
| 197 | +} // namespace scream |
0 commit comments