diff --git a/components/eamxx/docs/user/diags/vert_derivative.md b/components/eamxx/docs/user/diags/vert_derivative.md new file mode 100644 index 000000000000..622a3c26b004 --- /dev/null +++ b/components/eamxx/docs/user/diags/vert_derivative.md @@ -0,0 +1,75 @@ +# Vertical derivative diagnostics + +We can automatically calculate the vertical derivative of a given field +with respect to a vertical coordinate (pressure or height). +In other words, the vertical derivative is the derivative of the field +along a vertical coordinate. +To automatically get the derivative of a field X, we can request +`X_pvert_derivative` or `X_zvert_derivative` +in the output yaml files. + +*WARNING: support for the dz-based derivative is experimental!* + +## Brief description of algorithm + +For simplicity, the field `X` can only have two dimensions +(column and level). This covers most fields of interest, but +notably does not cover higher-dimensioned fields in radiation and MAM. +We calculate the derivative by finding the difference in the +input field `X`, which must be provided on midpoints, by interpolating +its values at interfaces weighted by the pressure differential. +That is, for a pressure-based derivative + +$$ +\nabla_{p} X = \frac{\Delta X}{\Delta p} +$$ + +and for height-based derivative + +$$ +\nabla_{p} X = \frac{\Delta X}{\Delta z} +$$ + +where + +$$ +\Delta X = \tilde{X}_{i,k+1} - \tilde{X}_{i,k} +$$ + +$$ +\tilde{X}_{i,k+1} = \frac{\left( + X_{i,k} \times \Delta p_{i,k+1} + + X_{i,k+1} \times \Delta p_{i,k} + \right)} + {\left( + \Delta p_{i,k+1} + \Delta p_{i,k} + \right)} +$$ + +Here, $\Delta p$ is the "pseudo density" pressure differential in Pa; +and $\Delta z$ is the height of the layer in m + +Note that the definition $\tilde{X}_{i,k+1}$ assumes that all +fields are weighted by the pressure differential, which is a +prognostic parameter of the dynamics core solver. Also note that +when $\Delta p_{i,k+1} \gg \Delta p_{i,k}$, +$\tilde{X}_{i,k+1}$ is closer to $X_{i,k}$ than $X_{i,k+1}$. + +## Example + +```yaml +%YAML 1.1 +--- +filename_prefix: monthly.outputs +averaging_type: average +max_snapshots_per_file: 1 +fields: + physics_pg2: + field_names: + # in this example, we use T_mid in units of K + - T_mid_pvert_derivative # K / Pa + - T_mid_zvert_derivative # K / m +output_control: + frequency: 1 + frequency_units: nmonths +``` diff --git a/components/eamxx/mkdocs.yml b/components/eamxx/mkdocs.yml index 72f82f4b3a01..1183f70cd383 100644 --- a/components/eamxx/mkdocs.yml +++ b/components/eamxx/mkdocs.yml @@ -20,6 +20,7 @@ nav: - 'Field contraction': 'user/diags/field_contraction.md' - 'Conditional sampling': 'user/diags/conditional_sampling.md' - 'Binary arithmetics': 'user/diags/binary_ops.md' + - 'Vertical derivative': 'user/diags/vert_derivative.md' - 'Presentations': 'user/presentations.md' - 'IO Aliases': 'user/io_aliases.md' - 'Developer Guide': diff --git a/components/eamxx/src/diagnostics/CMakeLists.txt b/components/eamxx/src/diagnostics/CMakeLists.txt index 8b740f852905..0cf2178ade05 100644 --- a/components/eamxx/src/diagnostics/CMakeLists.txt +++ b/components/eamxx/src/diagnostics/CMakeLists.txt @@ -24,6 +24,7 @@ set(DIAGNOSTIC_SRCS water_path.cpp wind_speed.cpp vert_contract.cpp + vert_derivative.cpp zonal_avg.cpp conditional_sampling.cpp ) diff --git a/components/eamxx/src/diagnostics/register_diagnostics.hpp b/components/eamxx/src/diagnostics/register_diagnostics.hpp index bf108a9a16db..b7e4a3aa4d6b 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/vert_derivative.hpp" #include "diagnostics/zonal_avg.hpp" #include "diagnostics/conditional_sampling.hpp" #include "diagnostics/binary_ops.hpp" @@ -58,6 +59,7 @@ inline void register_diagnostics () { diag_factory.register_product("AtmBackTendDiag",&create_atmosphere_diagnostic); diag_factory.register_product("HorizAvgDiag",&create_atmosphere_diagnostic); diag_factory.register_product("VertContractDiag",&create_atmosphere_diagnostic); + diag_factory.register_product("VertDerivativeDiag",&create_atmosphere_diagnostic); diag_factory.register_product("ZonalAvgDiag",&create_atmosphere_diagnostic); diag_factory.register_product("ConditionalSampling",&create_atmosphere_diagnostic); diag_factory.register_product("BinaryOpsDiag", &create_atmosphere_diagnostic); diff --git a/components/eamxx/src/diagnostics/tests/CMakeLists.txt b/components/eamxx/src/diagnostics/tests/CMakeLists.txt index 1c7cb7419f8e..0dd755b1e27b 100644 --- a/components/eamxx/src/diagnostics/tests/CMakeLists.txt +++ b/components/eamxx/src/diagnostics/tests/CMakeLists.txt @@ -78,6 +78,9 @@ CreateDiagTest(horiz_avg "horiz_avg_test.cpp" MPI_RANKS 1 ${SCREAM_TEST_MAX_RANK # Test for vertical contraction CreateDiagTest(vert_contract "vert_contract_test.cpp" MPI_RANKS 1 ${SCREAM_TEST_MAX_RANKS}) +# Test for vertical derivative +CreateDiagTest(vert_derivative "vert_derivative_test.cpp") + # Test zonal averaging CreateDiagTest(zonal_avg zonal_avg_test.cpp MPI_RANKS 1 ${SCREAM_TEST_MAX_RANKS}) diff --git a/components/eamxx/src/diagnostics/tests/vert_derivative_test.cpp b/components/eamxx/src/diagnostics/tests/vert_derivative_test.cpp new file mode 100644 index 000000000000..33138eee983b --- /dev/null +++ b/components/eamxx/src/diagnostics/tests/vert_derivative_test.cpp @@ -0,0 +1,166 @@ +#include "catch2/catch.hpp" +#include "diagnostics/register_diagnostics.hpp" +#include "physics/share/physics_constants.hpp" +#include "share/field/field_utils.hpp" +#include "share/grid/mesh_free_grids_manager.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) { + 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("vert_derivative") { + using namespace ShortFieldTagsNames; + using namespace ekat::units; + + // 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 = 150; + const int ngcols = 195 * comm.size(); + + auto gm = create_gm(comm, ngcols, nlevs); + auto grid = gm->get_grid("Physics"); + + // Input (randomized) qc + FieldLayout scalar1d_layout{{LEV}, {nlevs}}; + FieldLayout scalar2d_layout{{COL, LEV}, {ngcols, nlevs}}; + + FieldIdentifier fin2_fid("qc", scalar2d_layout, kg / kg, grid->name()); + FieldIdentifier dp_fid("pseudo_density", scalar2d_layout, Pa, grid->name()); + + Field fin2(fin2_fid); + Field dp(dp_fid); + + fin2.allocate_view(); + dp.allocate_view(); + + // Construct random number generator stuff + using RPDF = std::uniform_real_distribution; + RPDF dp_pdf(10,2000); + auto engine = scream::setup_random_test(); + + // Construct the diagnostics factory + std::map> diags; + auto &diag_factory = AtmosphereDiagnosticFactory::instance(); + register_diagnostics(); + + ekat::ParameterList params; + // instantiation works because we don't do anything in the constructor + auto bad_diag = diag_factory.create("VertDerivativeDiag", comm, params); + SECTION("bad_diag") { + // this will throw because no field_name was provided + REQUIRE_THROWS(bad_diag->set_grids(gm)); + } + + fin2.get_header().get_tracking().update_time_stamp(t0); + dp.get_header().get_tracking().update_time_stamp(t0); + // instead of random fin2, let's just a predictable one: + fin2.sync_to_host(); + auto fin2_v_host_assign = fin2.get_view(); + for (int icol = 0; icol < ngcols; ++icol) { + for (int ilev = 0; ilev < nlevs; ++ilev) { + fin2_v_host_assign(icol, ilev) = 1.0 + icol + ilev; + } + } + fin2.sync_to_dev(); + randomize(dp, engine, dp_pdf); + + // Create and set up the diagnostic + params.set("grid_name", grid->name()); + params.set("field_name", "qc"); + SECTION("bad_diag_2") { + // this will throw because no derivative_method was provided + auto bad_diag_2 = diag_factory.create("VertDerivativeDiag", comm, params); + REQUIRE_THROWS(bad_diag_2->set_grids(gm)); + } + + SECTION("bad_diag_3") { + // this will throw because bad derivative_method was provided + params.set("derivative_method", "xyz"); + auto bad_diag_3 = diag_factory.create("VertDerivativeDiag", comm, params); + REQUIRE_THROWS(bad_diag_3->set_grids(gm)); + } + + // dp_vert_derivative + params.set("derivative_method", "p"); + auto dp_vert_derivative = diag_factory.create("VertDerivativeDiag", comm, params); + + dp_vert_derivative->set_grids(gm); + + // Fields for manual calculation + FieldIdentifier diag1_fid("qc_vert_derivative_manual", scalar2d_layout, kg / kg, grid->name()); + Field diag1_m(diag1_fid); + diag1_m.allocate_view(); + + SECTION("dp_vert_derivative") { + // calculate dp_vert_div manually + dp.sync_to_host(); + fin2.sync_to_host(); + diag1_m.sync_to_host(); + + auto dp_v = dp.get_view(); + auto fin2_v = fin2.get_view(); + auto diag1_m_v = diag1_m.get_view(); + + for (int icol = 0; icol < ngcols; ++icol) { + for (int ilev = 0; ilev < nlevs; ++ilev) { + auto fa1 = (ilev < nlevs - 1) ? (fin2_v(icol, ilev + 1) * dp_v(icol, ilev) + + fin2_v(icol, ilev) * dp_v(icol, ilev + 1)) / + (dp_v(icol, ilev) + dp_v(icol, ilev + 1)) : fin2_v(icol, nlevs-1); + auto fa0 = (ilev > 0 ) ? (fin2_v(icol, ilev) * dp_v(icol, ilev - 1) + + fin2_v(icol, ilev - 1) * dp_v(icol, ilev)) / + (dp_v(icol, ilev - 1) + dp_v(icol, ilev)) : fin2_v(icol, 0); + diag1_m_v(icol, ilev) = (fa1 - fa0) / dp_v(icol, ilev); + } + } + diag1_m.sync_to_dev(); + + // Calculate weighted avg through diagnostics + dp_vert_derivative->set_required_field(fin2); + dp_vert_derivative->set_required_field(dp); + dp_vert_derivative->initialize(t0, RunType::Initial); + dp_vert_derivative->compute_diagnostic(); + auto dp_vert_derivative_f = dp_vert_derivative->get_diagnostic(); + + // Check that the two are the same manually using the tolerance + dp_vert_derivative_f.sync_to_host(); + auto view_deriv_f_d = dp_vert_derivative_f.get_view(); + auto view_deriv_f_m = diag1_m.get_view(); + for (int icol = 0; icol < ngcols; ++icol) { + for (int ilev = 0; ilev < nlevs; ++ilev) { + // TODO: the calculations are sometimes diverging because of fp issues in the test on h100 gcc13 + // TODO: so a cop-out, just test if the abs diff is within 1e-5 + REQUIRE(std::abs(view_deriv_f_d(icol, ilev) - view_deriv_f_m(icol, ilev)) < 1e-5); + } + } + } + + // TODO: add SECTION("dz_vert_derivative") +} + +} // namespace scream diff --git a/components/eamxx/src/diagnostics/vert_derivative.cpp b/components/eamxx/src/diagnostics/vert_derivative.cpp new file mode 100644 index 000000000000..53c32032f375 --- /dev/null +++ b/components/eamxx/src/diagnostics/vert_derivative.cpp @@ -0,0 +1,123 @@ +#include "diagnostics/vert_derivative.hpp" + +#include "share/util/eamxx_common_physics_functions.hpp" + +#include + +namespace scream { + +VertDerivativeDiag::VertDerivativeDiag(const ekat::Comm &comm, const ekat::ParameterList ¶ms) + : AtmosphereDiagnostic(comm, params) {} + +void VertDerivativeDiag::set_grids(const std::shared_ptr grids_manager) { + using namespace ShortFieldTagsNames; + using namespace ekat::units; + + const auto &fn = m_params.get("field_name"); + const auto &gn = m_params.get("grid_name"); + const auto g = grids_manager->get_grid(gn); + + add_field(fn, gn); + + // we support p or z + m_derivative_method = m_params.get("derivative_method"); + EKAT_REQUIRE_MSG(m_derivative_method == "p" || m_derivative_method == "z", + "Error! VertDerivativeDiag only supports 'p' or 'z' as derivative_method.\n" + " - derivative_method: " + + m_derivative_method + "\n"); + m_diag_name = fn + "_" + m_derivative_method + "vert_derivative"; + + auto scalar3d = g->get_3d_scalar_layout(true); + add_field("pseudo_density", scalar3d, Pa, gn); + if (m_derivative_method == "z") { + add_field("dz", scalar3d, m, gn); + } +} + +void VertDerivativeDiag::initialize_impl(const RunType /*run_type*/) { + using namespace ShortFieldTagsNames; + using namespace ekat::units; + + const auto &f = get_fields_in().front(); + const auto &fid = f.get_header().get_identifier(); + const auto &layout = fid.get_layout(); + + // TODO: support higher-dimensioned input fields + EKAT_REQUIRE_MSG(layout.rank() >= 2 && layout.rank() <= 2, + "Error! Field rank not supported by VertDerivativeDiag.\n" + " - field name: " + fid.name() + "\n" + " - field layout: " + layout.to_string() + "\n"); + EKAT_REQUIRE_MSG(layout.tags().back() == LEV, + "Error! VertDerivativeDiag diagnostic expects a layout ending " + "with the 'LEV' tag.\n" + " - field name : " + fid.name() + "\n" + " - field layout: " + layout.to_string() + "\n"); + + ekat::units::Units diag_units = fid.get_units(); + + if (m_derivative_method == "p") { + diag_units = fid.get_units() / Pa; + } else if (m_derivative_method == "z") { + diag_units = fid.get_units() / m; + } + + FieldIdentifier d_fid(m_diag_name, layout, diag_units, fid.get_grid_name()); + m_diagnostic_output = Field(d_fid); + m_diagnostic_output.allocate_view(); +} + +void VertDerivativeDiag::compute_diagnostic_impl() { + // f: input field; o: output field; d: denominator field + const auto &f = get_fields_in().front(); + const auto &o = m_diagnostic_output; + // keep dp around + const auto &dp = get_field_in("pseudo_density"); + + // TODO: support higher-dimensioned input fields + auto f2d = f.get_view(); + auto o2d = o.get_view(); + + auto dp2d = dp.get_view(); + + using KT = KokkosTypes; + using MT = typename KT::MemberType; + using TPF = ekat::TeamPolicyFactory; + const int ncols = f.get_header().get_identifier().get_layout().dim(0); + const int nlevs = f.get_header().get_identifier().get_layout().dim(1); + const auto policy = TPF::get_default_team_policy(ncols, nlevs); + + auto d_v = (m_derivative_method == "z") ? get_field_in("dz").get_view() : dp2d; + + Kokkos::parallel_for( + "Compute df / denominator for " + m_diagnostic_output.name(), policy, + KOKKOS_LAMBDA(const MT &team) { + const int icol = team.league_rank(); + auto f_icol = ekat::subview(f2d, icol); // field at midpoint + auto o_icol = ekat::subview(o2d, icol); // output at midnpoint + auto d_icol = ekat::subview(d_v, icol); // recall denominator is already a difference of interfaces + auto dpicol = ekat::subview(dp2d, icol); // in case of z deriv, d_icol and dpicol are not the same + + Kokkos::parallel_for(Kokkos::TeamVectorRange(team, nlevs), [&](const int ilev) { + // general formula for interface (assuming all fields are weighted by dp): + // f_int(k+1) = (f_m(k+1)dp(k) + f_m(k)dp(k+1)) / (dp(k) + dp(k+1)) + // if dp(k) << dp(k+1), the interface is closer to fm(k) than fm(k+1) + // boundary points: assume constant extrapolation (i.e., f_int(0)=f_mid(0)) + + auto f_int_kp1 = + (ilev < nlevs - 1) + ? (f_icol(ilev + 1) * dpicol(ilev) + f_icol(ilev) * dpicol(ilev + 1)) / + (dpicol(ilev) + dpicol(ilev + 1)) + : f_icol(nlevs - 1); + + auto f_int_kp0 = + (ilev > 0) + ? (f_icol(ilev) * dpicol(ilev - 1) + f_icol(ilev - 1) * dpicol(ilev)) / + (dpicol(ilev - 1) + dpicol(ilev)) + : f_icol(0); + + o_icol(ilev) = (f_int_kp1 - f_int_kp0) / d_icol(ilev); + }); + }); +} + +} // namespace scream diff --git a/components/eamxx/src/diagnostics/vert_derivative.hpp b/components/eamxx/src/diagnostics/vert_derivative.hpp new file mode 100644 index 000000000000..0567a086fe36 --- /dev/null +++ b/components/eamxx/src/diagnostics/vert_derivative.hpp @@ -0,0 +1,40 @@ + +#ifndef EAMXX_VERT_DERIVATIVE_HPP +#define EAMXX_VERT_DERIVATIVE_HPP + +#include "share/atm_process/atmosphere_diagnostic.hpp" + +namespace scream { + +/* + * This diagnostic will calculate the derivative of a field X in + * the vertical direction such that dX/dp or dX/dz + */ + +class VertDerivativeDiag : public AtmosphereDiagnostic { +public: + // Constructors + VertDerivativeDiag(const ekat::Comm &comm, const ekat::ParameterList ¶ms); + + // The name of the diagnostic CLASS (not the computed field) + std::string name() const { return "VertDerivativeDiag"; } + + // Set the grid + void set_grids(const std::shared_ptr grids_manager); + +protected: +#ifdef KOKKOS_ENABLE_CUDA +public: +#endif + void compute_diagnostic_impl(); + void initialize_impl(const RunType /*run_type*/); + + // Name of each field (because the diagnostic impl is generic) + std::string m_diag_name; + // Name of derivative method (differential dp or dz) + std::string m_derivative_method; +}; + +} // namespace scream + +#endif // EAMXX_VERT_DERIVATIVE_HPP diff --git a/components/eamxx/src/share/io/eamxx_io_utils.cpp b/components/eamxx/src/share/io/eamxx_io_utils.cpp index 54c9bd64f930..ba8723525b81 100644 --- a/components/eamxx/src/share/io/eamxx_io_utils.cpp +++ b/components/eamxx/src/share/io/eamxx_io_utils.cpp @@ -148,6 +148,7 @@ create_diagnostic (const std::string& diag_field_name, std::regex zonal_avg (R"()" + generic_field + R"(_zonal_avg_(\d+)_bins$)"); std::regex conditional_sampling (R"()" + generic_field + R"(_where_)" + generic_field + R"(_(gt|ge|eq|ne|le|lt)_([+-]?\d+(?:\.\d+)?)$)"); std::regex binary_ops (generic_field + "_" "(plus|minus|times|over)" + "_" + generic_field + "$"); + std::regex vert_derivative (generic_field + "_(p|z)vert_derivative$"); std::string diag_name; std::smatch matches; @@ -218,6 +219,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,vert_derivative)) { + diag_name = "VertDerivativeDiag"; + params.set("grid_name", grid->name()); + params.set("field_name", matches[1].str()); + params.set("derivative_method", matches[2].str()); + } else if (std::regex_search(diag_field_name,matches,zonal_avg)) { diag_name = "ZonalAvgDiag"; params.set("grid_name", grid->name());