EAMxx: diagnose vertical derivatives

Author: mahf708

components/eamxx/docs/user/diags/vert_derivative.md

@@ -0,0 +1,53 @@
+# 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
+`dX_dp_vert_derivative` or `dX_dz_vert_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 weighting two finite differences.
+The input field `X` must be defined on the midpoints, and its
+difference is defined as the difference between two consecutive
+levels, pointing downwards. That is, for a pressure coordinate,
+
+$$
+\nabla_{p} X = w  \frac{x}{a} + (1-w)\frac{y}{b}
+$$
+
+where
+
+- $x = X_{i, k} - X_{i, k-1}$ and $a = p_{i,k} - p_{i,k-1}$
+- $y = X_{i, k+1} - X_{i, k}$ and $b = p_{i,k+1} - p_{i,k}$
+- $w = b / (a + b)$
+
+and $x = 0$, $w=0$ for $k=0$ (highest level)
+and $y = 0$, $w=1$ for $k=K$ (lowest level).
+
+## 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
+      - dT_mid_dp_vert_derivative  # K / Pa
+      - dT_mid_dz_vert_derivative  # K / m
+output_control:
+  frequency: 1
+  frequency_units: nmonths
+```

components/eamxx/mkdocs.yml

@@ -18,6 +18,7 @@ nav:
     - 'Diagnostics':
       - 'Overview': 'user/diags/index.md'
       - 'Field contraction diagnostics': 'user/diags/field_contraction.md'
+      - 'Vertical divergence diagnostics': 'user/diags/vert_derivative.md'
     - 'Presentations': 'user/presentations.md'
   - 'Developer Guide':
     - 'Quick-start Guide': 'developer/dev_quickstart.md'

components/eamxx/src/diagnostics/CMakeLists.txt

@@ -23,6 +23,7 @@ set(DIAGNOSTIC_SRCS
   water_path.cpp
   wind_speed.cpp
   vert_contract.cpp
+  vert_derivative.cpp
   zonal_avg.cpp
 )
 

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"
 
 namespace scream {
@@ -56,6 +57,7 @@ inline void register_diagnostics () {
   diag_factory.register_product("AtmBackTendDiag",&create_atmosphere_diagnostic<AtmBackTendDiag>);
   diag_factory.register_product("HorizAvgDiag",&create_atmosphere_diagnostic<HorizAvgDiag>);
   diag_factory.register_product("VertContractDiag",&create_atmosphere_diagnostic<VertContractDiag>);
+  diag_factory.register_product("VertDerivativeDiag",&create_atmosphere_diagnostic<VertDerivativeDiag>);
   diag_factory.register_product("ZonalAvgDiag",&create_atmosphere_diagnostic<ZonalAvgDiag>);
 }
 

components/eamxx/src/diagnostics/tests/CMakeLists.txt

@@ -78,5 +78,8 @@ CreateDiagTest(horiz_avg "horiz_avg_test.cpp")
 # Test for vertical contraction
 CreateDiagTest(vert_contract "vert_contract_test.cpp")
 
+# 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})

components/eamxx/src/diagnostics/tests/vert_derivative_test.cpp

@@ -0,0 +1,163 @@
+#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<GridsManager> 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<std::string>;
+  ekat::ParameterList gm_params;
+  gm_params.set("grids_names", vos_t{"Point Grid"});
+  auto &pl = gm_params.sublist("Point Grid");
+  pl.set<std::string>("type", "point_grid");
+  pl.set("aliases", vos_t{"Physics"});
+  pl.set<int>("number_of_global_columns", num_global_cols);
+  pl.set<int>("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 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 pm_fid("p_mid", scalar2d_layout, Pa, grid->name());
+  FieldIdentifier zm_fid("z_mid", scalar2d_layout, m, grid->name());
+
+  Field fin2(fin2_fid);
+  Field pm(pm_fid);
+  // Field dz(dz_fid);
+
+  fin2.allocate_view();
+  pm.allocate_view();
+  // zm.allocate_view();
+
+  // Construct random number generator stuff
+  using RPDF = std::uniform_real_distribution<Real>;
+  RPDF pdf(sp(0.0), sp(200.0));
+  auto engine = scream::setup_random_test();
+
+  // Construct the diagnostics factory
+  std::map<std::string, std::shared_ptr<AtmosphereDiagnostic>> 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);
+  pm.get_header().get_tracking().update_time_stamp(t0);
+  // zm.get_header().get_tracking().update_time_stamp(t0);
+  randomize(fin2, engine, pdf);
+  randomize(pm, engine, pdf);
+  // randomize(zm, engine, pdf);
+
+  // Create and set up the diagnostic
+  params.set("grid_name", grid->name());
+  params.set<std::string>("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<std::string>("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<std::string>("derivative_method", "dp");
+  auto dp_vert_derivative = diag_factory.create("VertDerivativeDiag", comm, params);
+
+  // TODO: for some reason the z_mid field keeps getting set to 0
+  // TODO: as a workaround, just calculate z_mid here (sigh...)
+  // TODO: commenting out all the zm stuff for now
+  // // dz_vert_derivative
+  // params.set<std::string>("derivative_method", "dz");
+  // auto dz_vert_derivative = diag_factory.create("VertDerivativeDiag", comm, params);
+
+  dp_vert_derivative->set_grids(gm);
+  // dz_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
+    pm.sync_to_host();
+    fin2.sync_to_host();
+    diag1_m.sync_to_host();
+
+    auto pm_v      = pm.get_view<const Real **, Host>();
+    auto fin2_v    = fin2.get_view<const Real **, Host>();
+    auto diag1_m_v = diag1_m.get_view<Real **, Host>();
+
+    for (int icol = 0; icol < ngcols; ++icol) {
+      // handle boundary point
+      diag1_m_v(icol, 0) = (fin2_v(icol, 1) - fin2_v(icol, 0)) / (pm_v(icol, 1) - pm_v(icol, 0));
+      // interior points
+      for (int ilev = 1; ilev < nlevs - 1; ++ilev) {
+        auto pa               = pm_v(icol, ilev) - pm_v(icol, ilev - 1);
+        auto fa               = fin2_v(icol, ilev) - fin2_v(icol, ilev - 1);
+        auto pb               = pm_v(icol, ilev + 1) - pm_v(icol, ilev);
+        auto fb               = fin2_v(icol, ilev + 1) - fin2_v(icol, ilev);
+        auto ww               = pb / (pa + pb);
+        diag1_m_v(icol, ilev) = ww * (fa / pa) + (1 - ww) * (fb / pb);
+      }
+      // handle boundary point
+      diag1_m_v(icol, nlevs - 1) = (fin2_v(icol, nlevs - 1) - fin2_v(icol, nlevs - 2)) /
+                                   (pm_v(icol, nlevs - 1) - pm_v(icol, nlevs - 2));
+    }
+    diag1_m.sync_to_dev();
+
+    // Calculate weighted avg through diagnostics
+    dp_vert_derivative->set_required_field(fin2);
+    dp_vert_derivative->set_required_field(pm);
+    dp_vert_derivative->initialize(t0, RunType::Initial);
+    dp_vert_derivative->compute_diagnostic();
+    auto dp_vert_derivative_f = dp_vert_derivative->get_diagnostic();
+
+    REQUIRE(views_are_equal(dp_vert_derivative_f, diag1_m));
+  }
+
+  // TODO: add SECTION("dz_vert_derivative")
+}
+
+} // namespace scream