Merge pull request #347 from eagles-project/mingxuanwupnnl/eamxx_mam_microphys_for_mamxx

Mingxuanwupnnl/eamxx mam microphys for mamxx

Author: odiazib

src/mam4xx/gas_chem_mechanism.hpp

@@ -13,6 +13,7 @@ constexpr int nzcnt = 32;     // number of non-zero matrix entries
 constexpr int clscnt4 = 30;   // number of species in implicit class
 constexpr int extcnt = 9;     // number of species with external forcing
 constexpr int nfs = 8;        // number of fixed species
+constexpr int indexm = 0;     // index of total atm density in invariant array
 constexpr int permute_4[gas_pcnst] = {0,  1,  2,  3,  4,  5,  6,  7,  8,  9,
                                       10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
                                       20, 21, 22, 23, 24, 25, 26, 27, 28, 29};

src/mam4xx/mo_photo.hpp

@@ -72,8 +72,9 @@ inline PhotoTableData create_photo_table_data(int nw, int nt, int np_xs,
 
   // create views
   table_data.rsf_tab =
-      View5D("photo_table_data.rsf_tab", table_data.nw, table_data.numalb,
-             table_data.numcolo3, table_data.numsza, table_data.nump);
+      View5D("photo_table_data.rsf", table_data.nw, table_data.nump,
+             table_data.numsza, table_data.numcolo3, table_data.numalb);
+
   table_data.xsqy = View4D("photo_table_data.xsqy", table_data.numj,
                            table_data.nw, table_data.nt, table_data.np_xs);
   table_data.sza = View1D("photo_table_data.sza", table_data.numsza);
@@ -92,7 +93,7 @@ inline PhotoTableData create_photo_table_data(int nw, int nt, int np_xs,
   table_data.prs = View1D("photo_table_data.prs", table_data.np_xs);
   table_data.dprs = View1D("photo_table_data.dprs", table_data.np_xs - 1);
   table_data.pht_alias_mult_1 = View1D("photo_table_data.pht_alias_mult_1", 2);
-  table_data.lng_indexer = ViewInt1D("photo_table_data.lng_indexer", phtcnt);
+  table_data.lng_indexer = ViewInt1D("photo_table_data.lng_indexer", 1);
 
   return table_data;
 }
@@ -105,6 +106,30 @@ struct PhotoTableWorkArrays {
   View1D psum_l;
   View1D psum_u;
 };
+inline int get_photo_table_work_len(const PhotoTableData &photo_table_data) {
+  return pver * photo_table_data.numj +                /*lng_prates*/
+         pver * photo_table_data.nw +                  /*rsf*/
+         photo_table_data.numj * photo_table_data.nw + /*xswk*/
+         2 * photo_table_data.nw /*psum_l + psum_u*/;
+} // get_photo_table_work_len
+KOKKOS_INLINE_FUNCTION
+void set_photo_table_work_arrays(const PhotoTableData &photo_table_data,
+                                 const View1D &work,
+                                 PhotoTableWorkArrays &photo_table_work) {
+
+  auto work_ptr = (Real *)work.data();
+  photo_table_work.lng_prates = View2D(work_ptr, photo_table_data.numj, pver);
+  work_ptr += pver * photo_table_data.numj;
+  photo_table_work.rsf = View2D(work_ptr, photo_table_data.nw, pver);
+  work_ptr += pver * photo_table_data.nw;
+  photo_table_work.xswk =
+      View2D(work_ptr, photo_table_data.numj, photo_table_data.nw);
+  work_ptr += photo_table_data.numj * photo_table_data.nw;
+  photo_table_work.psum_l = View1D(work_ptr, photo_table_data.nw);
+  work_ptr += photo_table_data.nw;
+  photo_table_work.psum_u = View1D(work_ptr, photo_table_data.nw);
+  work_ptr += photo_table_data.nw;
+} // set_photo_table_work_arrays
 
 KOKKOS_INLINE_FUNCTION
 void cloud_mod(const Real zen_angle, const Real *clouds, const Real *lwc,
@@ -359,6 +384,7 @@ void calc_sum_wght(const Real dels[3], const Real wrk0, // in
   const Real wght_1_1_1 = dels[0] * wrk1;
 
   for (int wn = 0; wn < nw; wn++) {
+
     psum[wn] = wght_0_0_0 * rsf_tab(wn, iz, is, iv, ial) +
                wght_0_0_1 * rsf_tab(wn, iz, is, iv, ialp1) +
                wght_0_1_0 * rsf_tab(wn, iz, is, ivp1, ial) +
@@ -678,7 +704,7 @@ void table_photo(const View2D &photo, // out
                  const ConstColumnView &pmid, const ConstColumnView &pdel,
                  const ConstColumnView &temper, // in
                  const ColumnView &colo3_in, const Real zen_angle,
-                 const Real srf_alb, const ColumnView &lwc,
+                 const Real srf_alb, const ConstColumnView &lwc,
                  const ConstColumnView &clouds, // in
                  const Real esfact, const PhotoTableData &table_data,
                  PhotoTableWorkArrays &work_arrays) {
@@ -763,6 +789,7 @@ void table_photo(const View2D &photo, // out
       }   // end if
     }     // end mm
   }
+
   // } // end col_loop
 }
 

src/mam4xx/mo_setext.hpp

@@ -15,11 +15,14 @@ using View3D = DeviceType::view_3d<Real>;
 constexpr int pver = mam4::nlev;
 constexpr int extfrc_cnt = 9;
 constexpr int extcnt = 9; //, & ! number of species with external forcing
-
+// MAX_NUM_SECTIONS: Maximum number of sections in forcing data. Increase this
+// number if needed.
+constexpr int MAX_NUM_SECTIONS = 4;
 struct Forcing {
+  // This index is in Fortran format. i.e. starts in 1
   int frc_ndx;
   bool file_alt_data;
-  View2D fields_data;
+  View1D fields_data[MAX_NUM_SECTIONS];
   int nsectors;
 };
 
@@ -58,12 +61,12 @@ void extfrc_set(const Forcing *forcings, const View2D &frcing) {
         for (int kk = 0; kk < pver; ++kk) {
           // frcing(:ncol,:,nn) = frcing(:ncol,:,nn) + &
           // forcings(mm)%fields(isec)%data(:ncol,pver:1:-1,lchnk)
-          frcing(kk, nn) += forcing_mm.fields_data(isec, pver - 1 - kk);
+          frcing(kk, nn) += forcing_mm.fields_data[isec](pver - 1 - kk);
         } // kk
       } else {
         // // forcings(mm)%fields(isec)%data(:ncol,:,lchnk)
         for (int kk = 0; kk < pver; ++kk) {
-          frcing(kk, nn) += forcing_mm.fields_data(isec, kk);
+          frcing(kk, nn) += forcing_mm.fields_data[isec](kk);
         }
       }
     } // isec
@@ -139,4 +142,4 @@ void setext(const Forcing *forcings,
 
 } // end namespace mam4
 
-#endif
+#endif

src/mam4xx/mo_setinv.hpp

@@ -4,9 +4,9 @@
 #include <haero/atmosphere.hpp>
 
 #include <mam4xx/aero_config.hpp>
+#include <mam4xx/conversions.hpp>
 #include <mam4xx/gas_chem_mechanism.hpp>
 #include <mam4xx/mam4_types.hpp>
-
 namespace mam4 {
 namespace mo_setinv {
 // number of invariants
@@ -41,13 +41,9 @@ struct Config {
   Config &operator=(const Config &) = default;
 };
 
-// NOTE: vmr isn't actually used in the fortran code, but I've kept it here to
-// keep the function signature the same
 KOKKOS_INLINE_FUNCTION
 void setinv_single_level(Real invariants[nfs], const Real tfld,
-                         const Real h2ovmr,
-                         const Real vmr[AeroConfig::num_gas_phase_species()],
-                         const Real pmid,
+                         const Real h2ovmr, const Real pmid,
                          const Real cnst_offline[num_tracer_cnst],
                          const Config config_) {
   // -----------------------------------------------------------------
@@ -90,58 +86,37 @@ void setinv_single_level(Real invariants[nfs], const Real tfld,
 } // end setinv_single_level()
 
 KOKKOS_INLINE_FUNCTION
-void setinv(const ThreadTeam &team, const ColumnView invariants[nfs],
-            const ColumnView &tfld, const ColumnView &h2ovmr,
-            const ColumnView vmr[AeroConfig::num_gas_phase_species()],
-            const ColumnView cnst_offline[num_tracer_cnst],
-            const ColumnView &pmid) {
+void setinv(const ThreadTeam &team, const View2D &invariants,
+            const ConstColumnView &tfld, const ConstColumnView &qv,
+            const View1D cnst_offline[num_tracer_cnst],
+            const ConstColumnView &pmid) {
 
   Config setinv_config_;
   constexpr int nk = mam4::nlev;
+  constexpr Real mwh2o = haero::Constants::molec_weight_h2o;
 
   Kokkos::parallel_for(Kokkos::TeamVectorRange(team, nk), [&](int k) {
     const Real tfld_k = tfld(k);
-    const Real h2ovmr_k = h2ovmr(k);
+    const Real qv_k = qv(k);
     const Real pmid_k = pmid(k);
+
     Real invariants_k[nfs];
-    const int nspec = AeroConfig::num_gas_phase_species();
-    Real vmr_k[nspec];
-    for (int i = 0; i < nspec; ++i) {
-      vmr_k[i] = vmr[i](k);
+    for (int i = 0; i < nfs; ++i) {
+      invariants_k[i] = invariants(k, i);
     }
+
     Real cnst_offline_k[num_tracer_cnst];
     for (int i = 0; i < num_tracer_cnst; ++i) {
       cnst_offline_k[i] = cnst_offline[i](k);
     }
-    setinv_single_level(invariants_k, tfld_k, h2ovmr_k, vmr_k, pmid_k,
-                        cnst_offline_k, setinv_config_);
-    for (int i = 0; i < nfs; ++i) {
-      invariants[i](k) = invariants_k[i];
-    }
-  }); // end kokkos::parfor(k)
-} // end setinv_nlev()
 
-KOKKOS_INLINE_FUNCTION
-void setinv(const ThreadTeam &team, const ColumnView invariants[nfs],
-            const ColumnView &tfld, const ColumnView &h2ovmr, const View2D &vmr,
-            const View2D &cnst_offline, const ColumnView &pmid) {
-
-  Config setinv_config_;
-  constexpr int nk = mam4::nlev;
-
-  Kokkos::parallel_for(Kokkos::TeamVectorRange(team, nk), [&](int k) {
-    const Real tfld_k = tfld(k);
-    const Real h2ovmr_k = h2ovmr(k);
-    const Real pmid_k = pmid(k);
-    Real invariants_k[nfs];
-    View1D vmr_k = Kokkos::subview(vmr, k, Kokkos::ALL());
-    View1D cnst_offline_k = Kokkos::subview(cnst_offline, k, Kokkos::ALL());
+    Real h2ovmr_k = conversions::vmr_from_mmr(qv_k, mwh2o);
 
-    setinv_single_level(invariants_k, tfld_k, h2ovmr_k, vmr_k.data(), pmid_k,
-                        cnst_offline_k.data(), setinv_config_);
+    setinv_single_level(invariants_k, tfld_k, h2ovmr_k, pmid_k, cnst_offline_k,
+                        setinv_config_);
 
     for (int i = 0; i < nfs; ++i) {
-      invariants[i](k) = invariants_k[i];
+      invariants(k, i) = invariants_k[i];
     }
   }); // end kokkos::parfor(k)
 } // end setinv_nlev()

src/validation/mo_setext/extfrc_set.cpp

@@ -29,20 +29,22 @@ void extfrc_set(Ensemble *ensemble) {
       forcing_mm.file_alt_data = int(input.get_array(
           "forcings" + std::to_string(i) + "_file_alt_data")[0]);
 
-      forcing_mm.fields_data = View2D("data", forcing_mm.nsectors, pver);
+      EKAT_ASSERT_MSG(
+          MAX_NUM_SECTIONS >= int(forcing_mm.nsectors),
+          "Error! The number of sections is greater than MAX_NUM_SECTIONS..\n");
 
       for (int isec = 1; isec <= forcing_mm.nsectors; ++isec) {
-
         auto label = "forcings" + std::to_string(i) + "_fields" +
                      std::to_string(isec) + "_data";
         const auto data1 = input.get_array(label);
 
         View1DHost forcings_fields_data_host =
             View1DHost((Real *)data1.data(), data1.size());
 
-        const View1D data_isec =
-            Kokkos::subview(forcing_mm.fields_data, isec - 1, Kokkos::ALL());
+        View1D data_isec("data", pver);
         Kokkos::deep_copy(data_isec, forcings_fields_data_host);
+        forcing_mm.fields_data[isec - 1] = data_isec;
+
       } // isec
 
       forcings[i - 1] = forcing_mm;
@@ -61,4 +63,4 @@ void extfrc_set(Ensemble *ensemble) {
 
     output.set("frcing", frcing_out);
   });
-}
+}

src/validation/mo_setinv/setinv_test_nlev.cpp

@@ -14,6 +14,10 @@ using namespace mam4;
 using namespace haero;
 void setinv_test_nlev(Ensemble *ensemble) {
   ensemble->process([=](const Input &input, Output &output) {
+    using View1D = DeviceType::view_1d<Real>;
+    using View2D = DeviceType::view_2d<Real>;
+    using View2DHost = typename HostType::view_2d<Real>;
+    using View1DHost = typename HostType::view_1d<Real>;
     // Ensemble parameters
     // Declare array of strings for input names
     std::string input_arrays[] = {
@@ -33,71 +37,58 @@ void setinv_test_nlev(Ensemble *ensemble) {
 
     const Real tfld_in = input.get_array("tfld")[0];
     const Real h2ovmr_in = input.get_array("h2ovmr")[0];
-    // NOTE: vmr turns out to be unused, but still in the fxn signature for now
-    auto vmr_in = input.get_array("vmr");
     const Real pmid_in = input.get_array("pmid")[0];
     auto c_off_in = input.get_array("cnst_offline_yaml");
 
     const int nlev = mam4::nlev;
-    const int nspec = AeroConfig::num_gas_phase_species();
     const int nfs = mam4::mo_setinv::nfs;
     const int num_tracer_cnst = mam4::mo_setinv::num_tracer_cnst;
 
     ColumnView tfld = haero::testing::create_column_view(nlev);
-    auto tfld_h = Kokkos::create_mirror_view(tfld);
-    ColumnView h2ovmr = haero::testing::create_column_view(nlev);
-    auto h2ovmr_h = Kokkos::create_mirror_view(h2ovmr);
+    ColumnView qv = haero::testing::create_column_view(nlev);
     ColumnView pmid = haero::testing::create_column_view(nlev);
-    auto pmid_h = Kokkos::create_mirror_view(pmid);
 
-    using View1DHost = typename HostType::view_1d<Real>;
-    ColumnView invariants[nfs];
-    View1DHost invariants_h[nfs];
-    for (int i = 0; i < nfs; ++i) {
-      invariants[i] = haero::testing::create_column_view(nlev);
-      invariants_h[i] = Kokkos::create_mirror_view(invariants[i]);
-    }
+    View2D invariants("invariants", nlev, nfs);
+    auto invariants_h = Kokkos::create_mirror_view(invariants);
 
-    using View2D = DeviceType::view_2d<Real>;
-    using View2DHost = typename HostType::view_2d<Real>;
-    View2D vmr("vmr", nlev, nspec);
-    View2DHost vmr_h("vmr_h", nlev, nspec);
-    View2D c_off("c_off", nlev, num_tracer_cnst);
-    View2DHost c_off_h("c_off_h", nlev, num_tracer_cnst);
+    View1D c_off[num_tracer_cnst];
+    for (int i = 0; i < num_tracer_cnst; ++i) {
+      c_off[i] = View1D("c_off", nlev);
+    } //
+
+    View2DHost c_off_h("c_off_h", num_tracer_cnst, nlev);
+
+    constexpr Real mwh2o = Constants::molec_weight_h2o;
+    Real qv_k_in = conversions::mmr_from_vmr(h2ovmr_in, mwh2o);
 
     for (int k = 0; k < nlev; ++k) {
-      tfld_h(k) = tfld_in;
-      h2ovmr_h(k) = h2ovmr_in;
-      pmid_h(k) = pmid_in;
-      for (int i = 0; i < nspec; ++i) {
-        vmr_h(k, i) = vmr_in[i];
-      }
       for (int i = 0; i < num_tracer_cnst; ++i) {
-        c_off_h(k, i) = c_off_in[i];
+        c_off_h(i, k) = c_off_in[i];
       }
     }
 
-    Kokkos::deep_copy(tfld, tfld_h);
-    Kokkos::deep_copy(h2ovmr, h2ovmr_h);
-    Kokkos::deep_copy(pmid, pmid_h);
-    Kokkos::deep_copy(vmr, vmr_h);
-    Kokkos::deep_copy(c_off, c_off_h);
+    Kokkos::deep_copy(tfld, tfld_in);
+    Kokkos::deep_copy(qv, qv_k_in);
+    Kokkos::deep_copy(pmid, pmid_in);
+    for (int i = 0; i < num_tracer_cnst; ++i) {
+      const auto c_off_h_at_i = ekat::subview(c_off_h, i);
+      Kokkos::deep_copy(c_off[i], c_off_h_at_i);
+    }
 
     // Single-column dispatch.
     auto team_policy = ThreadTeamPolicy(1u, Kokkos::AUTO);
     Kokkos::parallel_for(
         team_policy, KOKKOS_LAMBDA(const ThreadTeam &team) {
-          mam4::mo_setinv::setinv(team, invariants, tfld, h2ovmr, vmr, c_off,
-                                  pmid);
+          mam4::mo_setinv::setinv(team, invariants, tfld, qv, c_off, pmid);
         });
 
     std::vector<Real> invariants_out(nfs);
     for (int i = 0; i < nfs; ++i) {
-      Kokkos::deep_copy(invariants_h[i], invariants[i]);
-      invariants_out[i] = invariants_h[i](0);
+      Kokkos::deep_copy(invariants_h, invariants);
+      invariants_out[i] = invariants_h(0, i);
       for (int k = 0; k < nlev; ++k) {
         // make sure every level got the same answer
-        EKAT_ASSERT(invariants_h[i](0) == invariants_h[i](k));
+        EKAT_ASSERT(invariants_h(0, i) == invariants_h(k, i));
       }
     }
 

src/validation/mo_setinv/setinv_test_single_level.cpp

@@ -33,8 +33,6 @@ void setinv_test_single_level(Ensemble *ensemble) {
 
     const Real tfld = input.get_array("tfld")[0];
     const Real h2ovmr = input.get_array("h2ovmr")[0];
-    // NOTE: vmr turns out to be unused, but still in the fxn signature for now
-    auto vmr = input.get_array("vmr");
     const Real pmid = input.get_array("pmid")[0];
     const int nfs = input.get_array("nfs")[0];
     auto cnst_offline = input.get_array("cnst_offline_yaml");
@@ -43,9 +41,8 @@ void setinv_test_single_level(Ensemble *ensemble) {
 
     Real invariants[nfs];
 
-    mam4::mo_setinv::setinv_single_level(invariants, tfld, h2ovmr, vmr.data(),
-                                         pmid, cnst_offline.data(),
-                                         setinv_config_);
+    mam4::mo_setinv::setinv_single_level(invariants, tfld, h2ovmr, pmid,
+                                         cnst_offline.data(), setinv_config_);
 
     std::vector<Real> inv_out;
     for (int i = 0; i < nfs; ++i) {