EAMxx: calculate vertical divergence in diagnostics

Author: mahf708

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

@@ -0,0 +1,45 @@
+# vertical divergence diagnostics
+
+We can automatically calculate the vertical divergence of a given field
+with respect to the vertical coordinate (pressure or height).
+In other words, the vertical divergence is the derivative of the field
+along a vertical coordinate.
+To automatically get the divergence of a field X, we can request
+`dX_dp_vert_divergence` or `X_dz_vert_divergence`
+in the output yaml files.
+
+## Brief description
+
+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.
+We also assume that the highest level (top of the atmosphere)
+has a zero derivative.
+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_{\text{vert}} X = \frac{X_{i, k} - X_{i, k-1}}{\Delta p_{k}}
+$$
+
+for all $k>0$ (at $k=0$, the derivative is zero).
+
+## 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_divergence  # K / Pa
+      - dT_mid_dz_vert_divergence  # 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_divergence.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_divergence.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_divergence.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("VertDivergenceDiag",&create_atmosphere_diagnostic<VertDivergenceDiag>);
   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 divergence
+CreateDiagTest(vert_divergence "vert_divergence_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_divergence_test.cpp

@@ -0,0 +1,157 @@
+#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_divergence") {
+  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 dp_fid("pseudo_density", scalar2d_layout, Pa, grid->name());
+  FieldIdentifier dz_fid("dz", scalar2d_layout, m, grid->name());
+
+  Field fin2(fin2_fid);
+  Field dp(dp_fid);
+  // Field dz(dz_fid);
+
+  fin2.allocate_view();
+  dp.allocate_view();
+  // dz.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("VertDivergenceDiag", 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);
+  // dz.get_header().get_tracking().update_time_stamp(t0);
+  randomize(fin2, engine, pdf);
+  randomize(dp, engine, pdf);
+  // randomize(dz, 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 divergence_method was provided
+    auto bad_diag_2 = diag_factory.create("VertDivergenceDiag", comm, params);
+    REQUIRE_THROWS(bad_diag_2->set_grids(gm));
+  }
+
+  SECTION("bad_diag_3") {
+    // this will throw because bad divergence_method was provided
+    params.set<std::string>("divergence_method", "xyz");
+    auto bad_diag_3 = diag_factory.create("VertDivergenceDiag", comm, params);
+    REQUIRE_THROWS(bad_diag_3->set_grids(gm));
+  }
+
+  // dp_vert_divergence
+  params.set<std::string>("divergence_method", "dp");
+  auto dp_vert_divergence = diag_factory.create("VertDivergenceDiag", comm, params);
+
+
+  // TODO: for some reason the dz field keeps getting set to 0
+  // TODO: as a workaround, just calculate dz here (sigh...)
+  // TODO: commenting out all the dz stuff for now
+  // // dz_vert_divergence
+  // params.set<std::string>("divergence_method", "dz");
+  // auto dz_vert_divergence = diag_factory.create("VertDivergenceDiag", comm, params);
+
+  dp_vert_divergence->set_grids(gm);
+  // dz_vert_divergence->set_grids(gm);
+
+  // Fields for manual calculation
+  FieldIdentifier diag1_fid("qc_vert_divergence_manual", scalar2d_layout, kg / kg, grid->name());
+  Field diag1_m(diag1_fid);
+  diag1_m.allocate_view();
+
+  SECTION("dp_vert_divergence") {
+    // calculate dp_vert_div manually 
+    dp.sync_to_host();
+    fin2.sync_to_host();
+    diag1_m.sync_to_host();
+
+    auto dp_v = dp.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) {
+        // first one is zero by design
+        diag1_m_v(icol, 0) = 0;
+        for (int ilev = 1; ilev < nlevs; ++ilev) {
+            diag1_m_v(icol, ilev) = (fin2_v(icol, ilev) - fin2_v(icol, ilev-1)) / dp_v(icol, ilev);
+        }
+    }
+    diag1_m.sync_to_dev();    
+
+    // Calculate weighted avg through diagnostics
+    dp_vert_divergence->set_required_field(fin2);
+    dp_vert_divergence->set_required_field(dp);
+    dp_vert_divergence->initialize(t0, RunType::Initial);
+    dp_vert_divergence->compute_diagnostic();
+    auto dp_vert_divergence_f = dp_vert_divergence->get_diagnostic();
+
+    REQUIRE(views_are_equal(dp_vert_divergence_f, diag1_m));
+  }
+  
+  // TODO: add SECTION("dz_vert_divergence") 
+
+}
+
+}  // namespace scream

components/eamxx/src/diagnostics/vert_divergence.cpp

@@ -0,0 +1,144 @@
+#include "diagnostics/vert_divergence.hpp"
+
+#include "share/util/eamxx_common_physics_functions.hpp"
+
+namespace scream {
+
+VertDivergenceDiag::VertDivergenceDiag(const ekat::Comm &comm,
+                                   const ekat::ParameterList &params)
+    : AtmosphereDiagnostic(comm, params) {}
+
+void VertDivergenceDiag::set_grids(
+    const std::shared_ptr<const GridsManager> grids_manager) {
+  using namespace ShortFieldTagsNames;
+  using namespace ekat::units;
+
+  const auto &fn = m_params.get<std::string>("field_name");
+  const auto &gn = m_params.get<std::string>("grid_name");
+  const auto g   = grids_manager->get_grid(gn);
+
+  add_field<Required>(fn, gn);
+
+  // we support dp or dz
+  m_divergence_method = m_params.get<std::string>("divergence_method");
+  EKAT_REQUIRE_MSG(
+      m_divergence_method == "dp" || m_divergence_method == "dz",
+      "Error! VertDivergenceDiag only supports 'dp' or 'dz' as divergence_method.\n"
+      " - divergence_method: " + m_divergence_method + "\n");
+  m_diag_name = "d" + fn + m_divergence_method + "_vert_divergence";
+
+  auto scalar3d = g->get_3d_scalar_layout(true);
+  if (m_divergence_method == "dp") {
+    add_field<Required>("pseudo_density", scalar3d, Pa, gn); 
+  }
+  else if (m_divergence_method == "dz")
+  {
+    add_field<Required>("pseudo_density", scalar3d, Pa, gn);
+    add_field<Required>("qv", scalar3d, kg / kg, gn);
+    add_field<Required>("p_mid", scalar3d, Pa, gn);
+    add_field<Required>("T_mid", scalar3d, K, gn);
+  }
+}
+
+void VertDivergenceDiag::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 VertDivergenceDiag.\n"
+                   " - field name: " + fid.name() + "\n"
+                   " - field layout: " + layout.to_string() + "\n");
+  EKAT_REQUIRE_MSG(layout.tags().back() == LEV,
+                   "Error! VertDivergenceDiag 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();
+
+  m_denominator = get_field_in("pseudo_density").clone("denominator");
+
+  if (m_divergence_method == "dp") {
+    diag_units = fid.get_units() / Pa;
+  }
+  else if (m_divergence_method == "dz") {
+    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 VertDivergenceDiag::compute_diagnostic_impl() {
+  const auto &f = get_fields_in().front();
+  const auto &d = m_diagnostic_output;
+
+  // TODO: support higher-dimensioned input fields
+  auto f2d = f.get_view<const Real**>(); 
+  auto d2d = d.get_view<Real**>();
+
+  using CO = scream::ColumnOps<DefaultDevice,Real>;
+  using KT  = KokkosTypes<DefaultDevice>;
+  using MT  = typename KT::MemberType;
+  using ESU = ekat::ExeSpaceUtils<typename KT::ExeSpace>;
+  using PF = scream::PhysicsFunctions<DefaultDevice>;
+  const int ncols = m_denominator.get_header().get_identifier().get_layout().dim(0);
+  const int nlevs = m_denominator.get_header().get_identifier().get_layout().dim(1);
+  const auto policy = ESU::get_default_team_policy(ncols, nlevs);
+
+  // get the denominator first
+  if (m_divergence_method == "dp") {
+    m_denominator.update(get_field_in("pseudo_density"), sp(1.0), sp(0.0));
+  }
+  else if (m_divergence_method == "dz") {
+    // TODO: for some reason the dz field keeps getting set to 0
+    // TODO: as a workaround, just calculate dz here (sigh...)
+    // m_denominator.update(get_field_in("dz"), 1.0, 0.0);
+
+    auto dz_v = m_denominator.get_view<Real**>();
+    auto dp_v = get_field_in("pseudo_density").get_view<const Real**>();
+    auto pm_v = get_field_in("p_mid").get_view<const Real**>();
+    auto tm_v = get_field_in("T_mid").get_view<const Real**>();
+    auto qv_v = get_field_in("qv").get_view<const Real**>();
+    Kokkos::parallel_for(
+      "Compute dz for " + m_diagnostic_output.name(), policy, KOKKOS_LAMBDA(const MT &team) {
+        const int icol = team.league_rank();
+        auto dz_icol = ekat::subview(dz_v, icol);
+        auto dp_icol = ekat::subview(dp_v, icol);
+        auto pm_icol = ekat::subview(pm_v, icol);
+        auto tm_icol = ekat::subview(tm_v, icol);
+        auto qv_icol = ekat::subview(qv_v, icol);
+        PF::calculate_dz(team, dp_icol, pm_icol, tm_icol, qv_icol, dz_icol);
+      });
+  }
+
+  auto d_v = m_denominator.get_view<Real**>();
+  // calculate df / denominator (setting the first level to be zero)
+  // TODO: perhaps we could calculate the values of the field at interfaces using col ops?
+  // TODO: though not sure it will be worth it...
+  Kokkos::parallel_for(
+    "Compute df / denominator for " + m_diagnostic_output.name(), policy, KOKKOS_LAMBDA(const MT &team) {
+      const int icol = team.league_rank();
+      auto df_icol = ekat::subview(f2d, icol);
+      auto dd_icol = ekat::subview(d2d, icol);
+      auto dv_icol = ekat::subview(d_v, icol);
+      
+      Kokkos::parallel_for(Kokkos::TeamVectorRange(team, nlevs), [&](const int ilev) {
+          if (ilev == 0) {
+            dd_icol(ilev) = 0;
+          }
+          else {
+            dd_icol(ilev) = (df_icol(ilev) - df_icol(ilev-1)) / dv_icol(ilev);
+          }
+      });
+    });
+}
+
+}  // namespace scream