add entrainemnt budget to eamxx

components/eamxx/src/diagnostics/CMakeLists.txt

@@ -20,6 +20,7 @@ set(DIAGNOSTIC_SRCS
   aodvis.cpp
   number_path.cpp
   aerocom_cld.cpp
+  entrainment_budget.cpp
 )
 
 add_library(diagnostics ${DIAGNOSTIC_SRCS})

components/eamxx/src/diagnostics/entrainment_budget.cpp

@@ -0,0 +1,268 @@
+#include "diagnostics/entrainment_budget.hpp"
+
+#include <ekat/kokkos/ekat_kokkos_utils.hpp>
+
+#include "diagnostics/entrainment_budget_util.hpp"
+#include "ekat/ekat_workspace.hpp"
+#include "share/util/scream_universal_constants.hpp"
+
+namespace scream {
+
+EntrainmentBudget::EntrainmentBudget(const ekat::Comm &comm, const ekat::ParameterList &params)
+    : AtmosphereDiagnostic(comm, params) {
+  // Nothing to do here
+}
+
+void EntrainmentBudget::set_grids(
+    const std::shared_ptr<const GridsManager> grids_manager) {
+  using namespace ekat::units;
+  using namespace ShortFieldTagsNames;
+
+  auto grid             = grids_manager->get_grid("Physics");
+  const auto &grid_name = grid->name();
+
+  const auto nondim = Units::nondimensional();
+
+  // Set the index map and units map
+  EntrainmentBudgetDiagUtil eadu;
+  m_index_map = eadu.index_map;
+  m_units_map = eadu.units_map;
+  m_ndiag     = eadu.size;
+
+  if(eadu.pblinvalg == "temperature-inversion-only") {
+    m_pblinvalg = 1;
+  } else {
+    EKAT_REQUIRE_MSG(
+        eadu.pblinvalg != "temperature-inversion-only",
+        "Error! Invalid pblinvalg. Only temperature-inversion-only "
+        "is currently supported.\n");
+  }
+
+  // Ensure m_index_map and m_units_map match
+  EKAT_REQUIRE_MSG(m_index_map.size() == m_units_map.size(),
+                   "Error! Some inconsistency in EntrainmentBudget: index and units "
+                   "maps do not match!\n");
+  // Ensure m_index_map and m_ndiag match
+  EKAT_REQUIRE_MSG(static_cast<int>(m_index_map.size()) == m_ndiag,
+                   "Error! Some inconsistency in EntrainmentBudget: m_ndiag and index "
+                   "map do not match!\n");
+
+  m_ncols = grid->get_num_local_dofs();
+  m_nlevs = grid->get_num_vertical_levels();
+
+  // Define layouts we need (both inputs and outputs)
+  FieldLayout scalar2d_layout{{COL, LEV}, {m_ncols, m_nlevs}};
+  FieldLayout vector1d_layout{{COL, CMP}, {m_ncols, m_ndiag}};
+
+  // The fields required for this diagnostic to be computed
+  // Get qc and qv
+  add_field<Required>("qc", scalar2d_layout, kg / kg, grid_name);
+  add_field<Required>("qv", scalar2d_layout, kg / kg, grid_name);
+  // Get T_mid, p_mid
+  add_field<Required>("T_mid", scalar2d_layout, K, grid_name);
+  add_field<Required>("p_mid", scalar2d_layout, Pa, grid_name);
+  // Get pseudo_density
+  add_field<Required>("pseudo_density", scalar2d_layout, Pa, grid_name);
+
+  // Construct and allocate the output field
+  FieldIdentifier fid("EntrainmentBudget", vector1d_layout, nondim, grid_name);
+  m_diagnostic_output = Field(fid);
+  m_diagnostic_output.allocate_view();
+
+  // Self-document the outputs to parse in post-processing
+  using stratt_t = std::map<std::string, std::string>;
+  auto d         = get_diagnostic();
+  auto &metadata =
+      d.get_header().get_extra_data<stratt_t>("io: string attributes");
+  for(const auto &it : m_index_map) {
+    metadata[it.first] =
+        std::to_string(it.second) + " (" + m_units_map[it.first] + ")";
+  }
+}
+
+void EntrainmentBudget::initialize_impl(const RunType /*run_type*/) {
+  // Field qt will have units and layout similar to qc, qv
+  const auto &qv   = get_field_in("qv");
+  const auto &qvid = qv.get_header().get_identifier();
+  const auto &qvgn = qvid.get_grid_name();
+  const auto &qvlo = qvid.get_layout();
+  FieldIdentifier qf_prev("qtot_prev", qvlo.clone(), qvid.get_units(), qvgn);
+  m_prev_qt = Field(qf_prev);
+  m_prev_qt.allocate_view();
+  // Field tl will have units of and layout similar to T_mid
+  const auto &tm   = get_field_in("T_mid");
+  const auto &tmid = tm.get_header().get_identifier();
+  const auto &tmgn = tmid.get_grid_name();
+  const auto &tmlo = tmid.get_layout();
+  FieldIdentifier tf_prev("tliq_prev", tmlo.clone(), tmid.get_units(), tmgn);
+  m_prev_tl = Field(tf_prev);
+  m_prev_tl.allocate_view();
+}
+
+void EntrainmentBudget::calc_tl_qt(const view_2d &tm_v, const view_2d &pm_v,
+                          const view_2d &qv_v, const view_2d &qc_v,
+                          const view_2d &tl_v, const view_2d &qt_v) {
+  int ncols = m_ncols;
+  int nlevs = m_nlevs;
+  Kokkos::parallel_for(
+      Kokkos::RangePolicy<>(0, ncols * nlevs), KOKKOS_LAMBDA(const int &idx) {
+        const int icol   = idx / ncols;
+        const int jlev   = idx % nlevs;
+        qt_v(icol, jlev) = qc_v(icol, jlev) + qv_v(icol, jlev);
+        tl_v(icol, jlev) = PF::calculate_thetal_from_theta(
+            PF::calculate_theta_from_T(tm_v(icol, jlev), pm_v(icol, jlev)),
+            tm_v(icol, jlev), qc_v(icol, jlev));
+      });
+}
+
+void EntrainmentBudget::init_timestep(const util::TimeStamp &start_of_step) {
+  m_start_t = start_of_step;
+
+  const auto &tm_v = get_field_in("T_mid").get_view<Real **>();
+  const auto &pm_v = get_field_in("p_mid").get_view<Real **>();
+  const auto &qv_v = get_field_in("qv").get_view<Real **>();
+  const auto &qc_v = get_field_in("qc").get_view<Real **>();
+
+  const auto &m_prev_qt_v = m_prev_qt.get_view<Real **>();
+  const auto &m_prev_tl_v = m_prev_tl.get_view<Real **>();
+
+  calc_tl_qt(tm_v, pm_v, qv_v, qc_v, m_prev_tl_v, m_prev_qt_v);
+}
+
+void EntrainmentBudget::compute_diagnostic_impl() {
+  using PC  = scream::physics::Constants<Real>;
+  using KT  = KokkosTypes<DefaultDevice>;
+  using MT  = typename KT::MemberType;
+  using ESU = ekat::ExeSpaceUtils<typename KT::ExeSpace>;
+
+  constexpr Real g = PC::gravit;
+
+  // Before doing anything, subview the out field for each variable
+  auto out = m_diagnostic_output.get_view<Real **>();
+
+  auto o_pm_hplus = ekat::subview_1(out, m_index_map["p+"]);
+  auto o_tl_hplus = ekat::subview_1(out, m_index_map["tl+"]);
+  auto o_tl_caret = ekat::subview_1(out, m_index_map["tl^"]);
+  auto o_tl_ttend = ekat::subview_1(out, m_index_map["tl_ttend"]);
+  auto o_qt_hplus = ekat::subview_1(out, m_index_map["qt+"]);
+  auto o_qt_caret = ekat::subview_1(out, m_index_map["qt^"]);
+  auto o_qt_ttend = ekat::subview_1(out, m_index_map["qt_ttend"]);
+
+  const auto &qc_v = get_field_in("qc").get_view<Real **>();
+  const auto &qv_v = get_field_in("qv").get_view<Real **>();
+  const auto &tm_v = get_field_in("T_mid").get_view<Real **>();
+  const auto &pm_v = get_field_in("p_mid").get_view<Real **>();
+  const auto &pd_v = get_field_in("pseudo_density").get_view<Real **>();
+
+  const auto &prev_qtot_v = m_prev_qt.get_view<Real **>();
+  const auto &prev_tliq_v = m_prev_tl.get_view<Real **>();
+
+  view_2d qt_v("qt_v", m_ncols, m_nlevs);
+  view_2d tl_v("tl_v", m_ncols, m_nlevs);
+  calc_tl_qt(tm_v, pm_v, qv_v, qc_v, tl_v, qt_v);
+
+  const auto &curr_ts =
+      get_field_in("qc").get_header().get_tracking().get_time_stamp();
+
+  std::int64_t dt = curr_ts - m_start_t;
+
+  const int num_levs  = m_nlevs;
+  const int pblinvalg = m_pblinvalg;
+  const auto policy   = ESU::get_default_team_policy(m_ncols, m_nlevs);
+
+  constexpr int wsms = 1;
+  using WSMgr        = ekat::WorkspaceManager<Real, DefaultDevice>;
+  WSMgr wsm(num_levs, wsms, policy);
+
+  Kokkos::parallel_for(
+      "Compute " + name(), policy, KOKKOS_LAMBDA(const MT &team) {
+        const int icol     = team.league_rank();
+        const auto qc_icol = ekat::subview(qc_v, icol);
+        const auto qv_icol = ekat::subview(qv_v, icol);
+        const auto tm_icol = ekat::subview(tm_v, icol);
+        const auto pm_icol = ekat::subview(pm_v, icol);
+        const auto pd_icol = ekat::subview(pd_v, icol);
+
+        const auto qt_icol = ekat::subview(qt_v, icol);
+        const auto tl_icol = ekat::subview(tl_v, icol);
+
+        auto prev_qtot_icol = ekat::subview(prev_qtot_v, icol);
+        auto prev_tliq_icol = ekat::subview(prev_tliq_v, icol);
+
+        auto ws = wsm.get_workspace(team);
+        ekat::Unmanaged<WSMgr::view_1d<Real>> tm_grad;
+        ws.take_many_contiguous_unsafe<wsms>({"tm_grad"}, {&tm_grad});
+
+        // We first want to find the PBL inversion. There are three methods to
+        // do so. All our methods here rely on the *gradient* of state fields
+        // (for now). First, we can simply find the first level from the surface
+        // that has a a temperature "inversion" (temperature goes positive
+        // instead of negative). Second, we can find the level which has the
+        // biggest positive jump in theta_l. Third, we can find the level which
+        // has the biggest negative jump in qt.
+
+        // We apply the first method only for now for now.
+        // We aim to find the minimum temperature gradient,
+        // indicating a positive temperature jump, but pressure
+        // is decreasing, so a negative gradient.
+        int opt_tm_grad_lev = 0;
+        if(pblinvalg == 1) {
+          // Find tm_grad = d(T_mid)/d(p_mid)
+          Kokkos::parallel_for(Kokkos::TeamVectorRange(team, 1, num_levs),
+                               [&](int k) {
+                                 tm_grad(k) = (tm_icol(k - 1) - tm_icol(k)) /
+                                              (pm_icol(k - 1) - pm_icol(k));
+                               });
+          team.team_barrier();
+
+          // Find minimum gradient, because d(p_mid) < 0 in definition above
+          // Starting from the surface, and ensuring p_mid > 70000.0 Pa,
+          // to avoid resolving to some odd place higher up in the atmosphere.
+          using minloc_t       = Kokkos::MinLoc<Real, int>;
+          using minloc_value_t = typename minloc_t::value_type;
+          minloc_value_t minloc;
+          Kokkos::parallel_reduce(
+              Kokkos::TeamVectorRange(team, 1, num_levs),
+              [&](const int &k, minloc_value_t &result) {
+                if(tm_grad(k) < result.val  && pm_icol(k) > 70000.0) {
+                  result.val = tm_grad(k);
+                  result.loc = k;
+                }
+              },
+              minloc_t(minloc));
+          opt_tm_grad_lev = minloc.loc;
+
+          // for(int k = num_levs; k > 1; k--) {
+          //   if(tm_grad(k) < opt_tm_grad && pm_icol(k) > 70000.0) {
+          //     opt_tm_grad     = tm_grad(k);
+          //     opt_tm_grad_lev = k;  // this is *inside* PBL
+          //   }
+          // }
+        }
+
+        // Save some outputs just above the "mixed" PBL
+        o_pm_hplus(icol) = pm_icol(opt_tm_grad_lev - 1);
+        o_tl_hplus(icol) = tl_icol(opt_tm_grad_lev - 1);
+        o_qt_hplus(icol) = qt_icol(opt_tm_grad_lev - 1);
+
+        // Now only need to compute below from opt_tm_grad_lev to num_levs
+        Kokkos::parallel_for(
+            Kokkos::TeamVectorRange(team, opt_tm_grad_lev, num_levs),
+            [&](int k) {
+              // The time-based tendencies
+              auto tl_tend = (tl_icol(k) - prev_tliq_icol(k)) / dt;
+              auto qt_tend = (qt_icol(k) - prev_qtot_icol(k)) / dt;
+              // the integrated (averaged) quantities
+              o_tl_caret(icol) += tl_icol(k) * pd_icol(k) / g;
+              o_tl_ttend(icol) += tl_tend * pd_icol(k) / g;
+              o_qt_caret(icol) += qt_icol(k) * pd_icol(k) / g;
+              o_qt_ttend(icol) += qt_tend * pd_icol(k) / g;
+            });
+        team.team_barrier();
+
+        // release stuff from wsm
+        ws.release_many_contiguous<wsms>({&tm_grad});
+      });
+}
+
+}  // namespace scream

components/eamxx/src/diagnostics/entrainment_budget.hpp

@@ -0,0 +1,69 @@
+#ifndef EAMXX_ENTRAINMENT_BUDGET_DIAG
+#define EAMXX_ENTRAINMENT_BUDGET_DIAG
+
+#include "share/atm_process/atmosphere_diagnostic.hpp"
+#include "share/util/scream_common_physics_functions.hpp"
+
+namespace scream {
+
+/*
+ * This diagnostic will compute extra ACI diagnostics.
+ */
+
+class EntrainmentBudget : public AtmosphereDiagnostic {
+ public:
+  using PF      = scream::PhysicsFunctions<DefaultDevice>;
+  using KT      = KokkosTypes<DefaultDevice>;
+  using view_2d = typename KT::template view_2d<Real>;
+
+  // Constructors
+  EntrainmentBudget(const ekat::Comm &comm, const ekat::ParameterList &params);
+
+  // The name of the diagnostic
+  std::string name() const { return "EntrainmentBudget"; };
+
+  // Set the grid
+  void set_grids(
+      const std::shared_ptr<const GridsManager> grids_manager) override;
+
+ protected:
+#ifdef KOKKOS_ENABLE_CUDA
+ public:
+#endif
+  void compute_diagnostic_impl();
+
+  // Let's override the init time step method
+  void init_timestep(const util::TimeStamp &start_of_step) override;
+
+  // Let's override the initialize method to set the fields below
+  void initialize_impl(const RunType /*run_type*/) override;
+
+  // A function to calculate the tracked fields of interest
+  void calc_tl_qt(const view_2d &tm_v, const view_2d &pm_v, const view_2d &qv_v,
+                  const view_2d &qc_v, const view_2d &tl_v,
+                  const view_2d &qt_v);
+
+  // Store fields at init_timestep
+  // That is before anything takes place in atm
+  Field m_prev_qt;
+  Field m_prev_tl;
+  util::TimeStamp m_start_t;
+
+  // Grid info
+  int m_ncols;
+  int m_nlevs;
+
+  // How many diags we have
+  int m_ndiag;
+
+  // Attribute maps for self-documentation
+  std::map<std::string, int> m_index_map;
+  std::map<std::string, std::string> m_units_map;
+
+  // Which PBL inversion algorithm to use
+  int m_pblinvalg;
+};
+
+}  // namespace scream
+
+#endif  // EAMXX_ENTRAINMENT_BUDGET_DIAG

components/eamxx/src/diagnostics/entrainment_budget_util.hpp

@@ -0,0 +1,45 @@
+#ifndef EAMXX_ENTRAINMENT_BUDGET_DIAG_UTIL
+#define EAMXX_ENTRAINMENT_BUDGET_DIAG_UTIL
+
+#include <map>
+
+namespace scream {
+class EntrainmentBudgetDiagUtil {
+ public:
+  std::map<std::string, int> index_map;
+  std::map<std::string, std::string> units_map;
+  int size;
+  std::string pblinvalg;
+
+  EntrainmentBudgetDiagUtil() {
+    // Start (post) incrementing size from 0
+    size = 0;
+
+    // pblp
+    index_map["p+"] = size++;
+    units_map["p+"] = "Pa";
+    // tl+
+    index_map["tl+"] = size++;
+    units_map["tl+"] = "K";
+    // tl^
+    index_map["tl^"] = size++;
+    units_map["tl^"] = "K";
+    // tl_ttend
+    index_map["tl_ttend"] = size++;
+    units_map["tl_ttend"] = "K/s";
+    // qt+
+    index_map["qt+"] = size++;
+    units_map["qt+"] = "kg/kg";
+    // qt^
+    index_map["qt^"] = size++;
+    units_map["qt^"] = "kg/kg";
+    // qt_ttend
+    index_map["qt_ttend"] = size++;
+    units_map["qt_ttend"] = "kg/kg/s";
+    // pblinvalg
+    pblinvalg = "temperature-inversion-only";
+  }
+};
+}  // namespace scream
+
+#endif  // EAMXX_ENTRAINMENT_BUDGET_DIAG_UTIL