use vert_contraction() for diagnostic outputs--almost finished

Author: mjschmidt271

components/eamxx/src/physics/mam/eamxx_mam_microphysics_process_interface.cpp

@@ -188,15 +188,53 @@ void MAMMicrophysics::set_grids(
   add_field<Updated>("constituent_fluxes", scalar2d_pcnst, kg / m2 / s,
                      grid_name);
 
+  // FIXME: moved this to class
   // Number of externally forced chemical species
-  constexpr int extcnt = mam4::gas_chemistry::extcnt;
+  // constexpr int extcnt = mam4::gas_chemistry::extcnt;
 
-  FieldLayout scalar3d_extcnt = grid_->get_3d_vector_layout(true, extcnt, "ext_cnt");
+  FieldLayout scalar3d_extcnt = grid_->get_3d_vector_layout(true, extcnt_, "ext_cnt");
 
   // Register computed fields for external forcing
   // - extfrc: 3D instantaneous forcing rate [kg/m³/s]
   add_field<Computed>("mam4_external_forcing", scalar3d_extcnt, kg / m3 / s, grid_name);
 
+  // extents are [ncol, nlev, gas_pcnst, nq<...>]
+  FieldLayout tensor3d_gas_tend =
+      grid_->get_3d_tensor_layout(true, {num_gas_spec_, mam4::microphysics::nqtendaa()},
+                                  {"num_gas_spec", "nqtendaa"});
+  FieldLayout tensor3d_gas_tend_cw =
+      grid_->get_3d_tensor_layout(true, {num_gas_spec_, mam4::microphysics::nqqcwtendaa()},
+                                  {"num_gas_spec", "nqqcwtendaa"});
+  FieldLayout tensor2d_gas_tend =
+      grid_->get_3d_tensor_layout(true, {num_gas_spec_, mam4::microphysics::nqtendaa()},
+                                  {"num_gas_spec", "nqtendaa"});
+  FieldLayout tensor2d_gas_tend_cw =
+      grid_->get_3d_tensor_layout(true, {num_gas_spec_, mam4::microphysics::nqqcwtendaa()},
+                                  {"num_gas_spec", "nqqcwtendaa"});
+  FieldLayout vector3d_gas_tend = grid_->get_3d_vector_layout(true, num_gas_spec_, "num_gas_spec");
+
+  // fields for holding diagnostic quantities relating to gas-species tendencies
+  // NOTE: denoting this as non-dimensional because the original fortran MAM used these to store
+  // different "tracer mixing ratios" that could be [mol/mol], [#/kmol], etc.
+  // consider breaking this into separate 3D-vector fields?
+  add_field<Computed>("gas_spec_tendencies", tensor3d_gas_tend, nondim, grid_name);
+  add_field<Computed>("gas_spec_tendencies_cw", tensor3d_gas_tend_cw, nondim, grid_name);
+  // these are the column-integrated fields returned by vert_contraction()
+  add_field<Computed>("col_int_gas_spec_tend", tensor2d_gas_tend, nondim, grid_name);
+  add_field<Computed>("col_int_gas_spec_tend_cw", tensor2d_gas_tend_cw, nondim, grid_name);
+  // and finally, the weights used by vert_contraction()
+  add_field<Computed>("wts_gas_spec_tends", vector3d_gas_tend, nondim, grid_name);
+
+  // ===============================================================================================
+  // mam::microphysics wants the full column of these tendencies
+  // TODO: is it better for each column's threadTeam to have its own 3d view,
+  //       or subview into a 4d view that contains all 'ncol' columns of these?
+  // view_3d qgcm_tendaa("diag_tendency-col", nlev, num_gas_spec_, mam4::microphysics::nqtendaa());
+  // view_3d qqcwgcm_tendaa("diag_tendency_cw-col", nlev, num_gas_spec_, mam4::microphysics::nqqcwtendaa());
+  // Kokkos::deep_copy(qgcm_tendaa, 0.0);
+  // Kokkos::deep_copy(qqcwgcm_tendaa, 0.0);
+  // ===============================================================================================
+
   // Creating a Linoz reader and setting Linoz parameters involves reading data
   // from a file and configuring the necessary parameters for the Linoz model.
   {
@@ -387,7 +425,8 @@ void MAMMicrophysics::init_buffers(const ATMBufferManager &buffer_manager) {
 int MAMMicrophysics::get_len_temporary_views() {
   const int photo_table_len = get_photo_table_work_len(photo_table_);
   const int sethet_work_len = mam4::mo_sethet::get_total_work_len_sethet();
-  constexpr int extcnt      = mam4::gas_chemistry::extcnt;
+  // FIXME: moved this to class
+  // constexpr int extcnt      = mam4::gas_chemistry::extcnt;
   int work_len              = 0;
   // work_photo_table_
   work_len += ncol_ * photo_table_len;
@@ -398,13 +437,14 @@ int MAMMicrophysics::get_len_temporary_views() {
   // invariants_
   work_len += ncol_ * nlev_ * mam4::gas_chemistry::nfs;
   // extfrc_
-  work_len += ncol_ * nlev_ * extcnt;
+  work_len += ncol_ * nlev_ * extcnt_;
   return work_len;
 }
 void MAMMicrophysics::init_temporary_views() {
   const int photo_table_len = get_photo_table_work_len(photo_table_);
   const int sethet_work_len = mam4::mo_sethet::get_total_work_len_sethet();
-  constexpr int extcnt      = mam4::gas_chemistry::extcnt;
+  // FIXME: moved this to class
+  // constexpr int extcnt      = mam4::gas_chemistry::extcnt;
   auto work_ptr             = (Real *)buffer_.temporary_views.data();
 
   work_photo_table_ = view_2d(work_ptr, ncol_, photo_table_len);
@@ -416,8 +456,8 @@ void MAMMicrophysics::init_temporary_views() {
   work_ptr += ncol_ * nlev_ * mam4::mo_photo::phtcnt;
   invariants_ = view_3d(work_ptr, ncol_, nlev_, mam4::gas_chemistry::nfs);
   work_ptr += ncol_ * nlev_ * mam4::gas_chemistry::nfs;
-  extfrc_ = view_3d(work_ptr, ncol_, nlev_, extcnt);
-  work_ptr += ncol_ * nlev_ * extcnt;
+  extfrc_ = view_3d(work_ptr, ncol_, nlev_, extcnt_);
+  work_ptr += ncol_ * nlev_ * extcnt_;
 
   // Error check
   // NOTE: workspace_provided can be larger than workspace_used, but let's try
@@ -774,18 +814,18 @@ void MAMMicrophysics::run_impl(const double dt) {
     Kokkos::deep_copy(acos_cosine_zenith_, acos_cosine_zenith_host_);
   }
   const auto zenith_angle = acos_cosine_zenith_;
-  constexpr int gas_pcnst = mam_coupling::gas_pcnst();
 
-  const auto &extfrc   = extfrc_;
-  const auto &forcings = forcings_;
-  constexpr int extcnt = mam4::gas_chemistry::extcnt;
+  const auto &extfrc               = extfrc_;
+  const auto &forcings             = forcings_;
+  // FIXME: moved this to class
+  // constexpr int extcnt             = mam4::gas_chemistry::extcnt;
 
   const int offset_aerosol = mam4::utils::gasses_start_ind();
-  Real adv_mass_kg_per_moles[gas_pcnst];
+  Real adv_mass_kg_per_moles[num_gas_spec_];
   // NOTE: Making copies of clsmap_4 and permute_4 to fix undefined arrays on
   // the device.
-  int clsmap_4[gas_pcnst], permute_4[gas_pcnst];
-  for(int i = 0; i < gas_pcnst; ++i) {
+  int clsmap_4[num_gas_spec_], permute_4[num_gas_spec_];
+  for(int i = 0; i < num_gas_spec_; ++i) {
     // NOTE: state_q is kg/kg-dry-air; adv_mass is in g/mole.
     // Convert adv_mass to kg/mole as vmr_from_mmr function uses
     // molec_weight_dry_air with kg/mole units
@@ -803,13 +843,22 @@ void MAMMicrophysics::run_impl(const double dt) {
   const auto &index_season_lai = index_season_lai_;
   const int pcnst              = mam4::pcnst;
 
-  Kokkos::Array<Real, gas_pcnst> molar_mass_g_per_mol_tmp;
-  for (int i = 0; i < gas_pcnst; ++i) {
+  Kokkos::Array<Real, num_gas_spec_> molar_mass_g_per_mol_tmp;
+  for (int i = 0; i < num_gas_spec_; ++i) {
     molar_mass_g_per_mol_tmp[i] = mam4::gas_chemistry::adv_mass[i];  // host-only access
   }
 
+  // these are computed in the column loop and used afterward, so declare here
+  // TODO: consider enabling/disabling this based on whether output is required--this would
+  // require another version of perform_atmospheric_chemistry_and_microphysics(), potentially
+  // templated on the forthcoming diagnostics struct
+  // nonetheless, skip for now
+  // Slice into the diagnostic fields to pass a single column to mam4xx::microphysics
+  auto gas_spec_tend = get_field_out("gas_spec_tendencies");
+  auto gas_spec_tend_cw = get_field_out("gas_spec_tendencies_cw");
+  auto gas_spec_tend_wts = get_field_out("wts_gas_spec_tends");
   //NOTE: we need to initialize photo_rates_
-  Kokkos::deep_copy(photo_rates_,0.0);
+  Kokkos::deep_copy(photo_rates_, 0.0);
   // loop over atmosphere columns and compute aerosol microphyscs
   Kokkos::parallel_for(
       "MAMMicrophysics::run_impl", policy,
@@ -833,9 +882,9 @@ void MAMMicrophysics::run_impl(const double dt) {
             mam_coupling::aerosols_for_column(dry_aero, icol);
 
         const auto invariants_icol = ekat::subview(invariants, icol);
-        mam4::mo_setext::Forcing forcings_in[extcnt];
+        mam4::mo_setext::Forcing forcings_in[extcnt_];
 
-        for(int i = 0; i < extcnt; ++i) {
+        for(int i = 0; i < extcnt_; ++i) {
           const int nsectors       = forcings[i].nsectors;
           const int frc_ndx        = forcings[i].frc_ndx;
           const auto file_alt_data = forcings[i].file_alt_data;
@@ -921,20 +970,19 @@ void MAMMicrophysics::run_impl(const double dt) {
             }
           }
         }
+
+        auto gas_spec_tend_V = gas_spec_tend.get_view<Real****>();
+        auto gas_spec_tend_cw_V = gas_spec_tend_cw.get_view<Real****>();
+        auto gas_spec_tend_wts_V = gas_spec_tend_wts.get_view<Real***>();
+        // NOTE: these are called "qgcm_tendaa" and "qqcwgcm_tendaa" in mam4xx
+        auto gas_spec_tend_col = ekat::subview(gas_spec_tend_V, icol);
+        auto gas_spec_tend_cw_col = ekat::subview(gas_spec_tend_cw_V, icol);
+
         // These output values need to be put somewhere:
-        const auto aqso4_flx_col = ekat::subview(aqso4_flx, icol);  // deposition flux of so4 [mole/mole/s]
-        const auto aqh2so4_flx_col = ekat::subview(aqh2so4_flx, icol);  // deposition flux of h2so4 [mole/mole/s]
-        Real dflx_col[gas_pcnst] = {};  // deposition velocity [1/cm/s]
-        Real dvel_col[gas_pcnst] = {};  // deposition flux [1/cm^2/s]
-        // mam::microphysics wants the full column of these tendencies
-        // TODO: is it better for each column's threadTeam to have its own 3d view,
-        //       or subview into a 4d view that contains all 'ncol' columns of these?
-        view_3d qgcm_tendaa("diag_tendency-col", nlev, gas_pcnst, mam4::microphysics::nqtendaa());
-        view_3d qqcwgcm_tendaa("diag_tendency_cw-col", nlev, gas_pcnst, mam4::microphysics::nqqcwtendaa());
-        Kokkos::deep_copy(qgcm_tendaa, 0.0);
-        Kokkos::deep_copy(qqcwgcm_tendaa, 0.0);
+        Real dflx_col[num_gas_spec_] = {};  // deposition velocity [1/cm/s]
+        Real dvel_col[num_gas_spec_] = {};  // deposition flux [1/cm^2/s]
         // Output: values are dvel, dflx
-        // Diagnostic Output: qgcm_tendaa, qqcwgcm_tendaa
+        // Diagnostic Output: gas_spec_tend_col, gas_spec_tend_cw_col
         // Input/Output: progs::stateq, progs::qqcw
         team.team_barrier();
         mam4::microphysics::perform_atmospheric_chemistry_and_microphysics(
@@ -955,31 +1003,41 @@ void MAMMicrophysics::run_impl(const double dt) {
             dvel_col, dflx_col, qgcm_tendaa, qqcwgcm_tendaa, progs);
 
         team.team_barrier();
+        const auto pdel_col = ekat::subview(dry_atm.p_del, icol);
+        auto wts_col = ekat::subview(gas_spec_tend_wts_V, icol);
+        set_vert_contraction_weights(wts_col, molar_mass_g_per_mol_tmp, pdel_col, nlev_);
+
         // Update constituent fluxes with gas drydep fluxes (dflx)
         // FIXME: Possible units mismatch (dflx is in kg/cm2/s but
         // constituent_fluxes is kg/m2/s) (Following mimics Fortran code
         // behavior but we should look into it)
         Kokkos::parallel_for(Kokkos::TeamVectorRange(team, offset_aerosol, pcnst), [&](int ispc) {
           constituent_fluxes(icol, ispc) -= dflx_col[ispc - offset_aerosol];
         });
-
       });  // parallel_for for the column loop
   Kokkos::fence();
 
+  // the variables in these column-integrated tendencies correspond to these MAM diagnostics:
+  // ['<xyz>_sfgaex1', '<xyz>_sfgaex2', '<xyz>_sfnnuc1', '<xyz>_sfcoag1']
+  auto col_int_gas_spec_tend = get_field_out("col_int_gas_spec_tend");
+  auto col_int_gas_spec_tend_cw = get_field_out("col_int_gas_spec_tend_cw");
+  vert_contraction<Real>(col_int_gas_spec_tend, gas_spec_tend, gas_spec_tend_wts);
+  vert_contraction<Real>(col_int_gas_spec_tend_cw, gas_spec_tend_cw, gas_spec_tend_wts);
+
   auto extfrc_fm = get_field_out("mam4_external_forcing").get_view<Real***>();
 
   // Avogadro's number [molecules/mol]
   const Real Avogadro = haero::Constants::avogadro;
   // Mapping from external forcing species index to physics constituent index
   // NOTE: These indices should match the species in extfrc_lst
   // TODO: getting rid of hard-coded indices
-  Kokkos::Array<int, extcnt> extfrc_pcnst_index = {3, 6, 14, 27, 28, 13, 18, 30, 5};
+  Kokkos::Array<int, extcnt_> extfrc_pcnst_index = {3, 6, 14, 27, 28, 13, 18, 30, 5};
 
   // Transpose extfrc_ from internal layout [ncol][nlev][extcnt]
   // to output layout [ncol][extcnt][nlev]
   // This aligns with expected field storage in the EAMxx infrastructure.
   Kokkos::parallel_for("transpose_extfrc",
-    Kokkos::MDRangePolicy<Kokkos::Rank<3>>({0,0,0}, {ncol, extcnt, nlev}),
+    Kokkos::MDRangePolicy<Kokkos::Rank<3>>({0,0,0}, {ncol, extcnt_, nlev}),
     KOKKOS_LAMBDA(const int i, const int j, const int k) {
       const int pcnst_idx = extfrc_pcnst_index[j];
       const Real molar_mass_g_per_mol = molar_mass_g_per_mol_tmp[pcnst_idx]; // g/mol

components/eamxx/src/physics/mam/eamxx_mam_microphysics_process_interface.hpp

@@ -4,6 +4,7 @@
 #include <physics/mam/eamxx_mam_generic_process_interface.hpp>
 #include <physics/mam/mam_coupling.hpp>
 #include <share/util/eamxx_common_physics_functions.hpp>
+#include <share/field/field_utils.hpp>
 
 #include "readfiles/tracer_reader_utils.hpp"
 // For calling MAM4 processes
@@ -60,58 +61,11 @@ class MAMMicrophysics final : public MAMGenericInterface {
   // Finalize
   void finalize_impl(){/*Do nothing*/};
 
-  void process_diagnostic_tendencies(view_3d &qgcm_tendaa, view_3d &qqcwgcm_tendaa,
-                                  const Kokkos::Array<Real, mam_coupling::gas_pcnst()> &adv_mass,
-                                  const const_view_1d &pdel, const int ncol, const int nlev) {
-    // HACK: these are placeholders until we establish a way of requesting diagnostic output
-    // flag: write the full columns of mean tracer mixing ratios
-    constexpr bool write_full_col_gas_spec = true;
-    // flag: write the column-integrated tendencies
-    constexpr bool write_coltend_gas_spec = true;
-    const int gas_pcnst = mam_coupling::gas_pcnst();
-    const int nq = mam4::microphysics::nqtendaa();
-    const int nqqcw = mam4::microphysics::nqqcwtendaa();
-
-    if (write_full_col_gas_spec) {
-      // TODO: write out qgcm_tendaa and qqcwgcm_tendaa...
-    }
-
-    if (write_coltend_gas_spec) {
-      using physconst = scream::physics::Constants<Real>;
-      static constexpr Real gravity = physconst::gravit;
-      static constexpr Real mw_dry_air = physconst::MWdry;
-      // first calculate the tendencies
-      view_2d q_coltendaa("diag_tends", gas_pcnst, nq);
-      view_2d qqcw_coltendaa("diag_tends-cw", gas_pcnst, nqqcw);
-      Kokkos::deep_copy(q_coltendaa, 0.0);
-      Kokkos::deep_copy(qqcw_coltendaa, 0.0);
-
-      // TODO: can probably do this using vert_contract.hpp, though the messiness
-      // of the weights could mean it's not worth it
-      Kokkos::parallel_for("calc_diagnostic_tendencies",
-        Kokkos::MDRangePolicy<Kokkos::Rank<3>>({0,0,0}, {nlev, gas_pcnst, nq}),
-        KOKKOS_LAMBDA(const int k, const int spec, const int itend) {
-          Real pdel_fac = pdel(k) / gravity;
-          // const auto q_k = ekat::subview(qgcm_tendaa, k);
-          q_coltendaa(spec, itend) = q_coltendaa(spec, itend)
-              + qgcm_tendaa(k, spec, itend) * pdel_fac * (adv_mass[spec] / mw_dry_air);
-      });
-      // NOTE: looping over the third dimension (nqqcw) may be pointless
-      // here because nqqcw == 1 appears to be hard-coded
-      Kokkos::parallel_for("calc_diagnostic_tendencies",
-        Kokkos::MDRangePolicy<Kokkos::Rank<3>>({0,0,0}, {nlev, gas_pcnst, nqqcw}),
-        KOKKOS_LAMBDA(const int k, const int spec, const int itend) {
-          Real pdel_fac = pdel(k) / gravity;
-          // const auto q_k = ekat::subview(qqcwgcm_tendaa, k);
-          qqcw_coltendaa(spec, itend) = qqcw_coltendaa(spec, itend)
-              + qqcwgcm_tendaa(k, spec, itend) * pdel_fac * (adv_mass[spec] / mw_dry_air);
-      });
-
-      // TODO: write q_coltendaa and qqcw_coltendaa out...
-    }
-  }
-
  private:
+  // number of species involved in gas phase chemistry (gases + aerosols)
+  static constexpr int num_gas_spec_ = mam_coupling::gas_pcnst();
+  // number of species with external forcing
+  static constexpr int extcnt_ = mam4::gas_chemistry::extcnt;
   // The orbital year, used for zenith angle calculations:
   // If > 0, use constant orbital year for duration of simulation
   // If < 0, use year from timestamp for orbital parameters
@@ -211,6 +165,17 @@ class MAMMicrophysics final : public MAMGenericInterface {
   void init_temporary_views();
   int len_temporary_views_{0};
 
+  void set_vert_contraction_weights(view_2d wts, const Kokkos::Array<Real, num_gas_spec_> adv_mass,
+                                    const const_view_1d pdel, const int nlev) {
+    using physconst = scream::physics::Constants<Real>;
+    static constexpr Real gravity = physconst::gravit;
+    static constexpr Real mw_dry_air = physconst::MWdry;
+    Kokkos::parallel_for("calc_diagnostic_tendencies",
+      Kokkos::MDRangePolicy<Kokkos::Rank<2>>({0,0}, {nlev, num_gas_spec_}),
+      KOKKOS_LAMBDA(const int k, const int spec) {
+        wts(k, spec) = pdel(k) / gravity * adv_mass[spec] / mw_dry_air;
+    });
+  }
 };  // MAMMicrophysics
 
 }  // namespace scream