Merge pull request #3108 from eagles-project/oscar/find_season_index

MAM4xx: Add reader for season_wes, which is part of the microphysics interface.

Author: jgfouca

components/eamxx/cime_config/namelist_defaults_scream.xml

@@ -295,6 +295,8 @@ be lost if SCREAM_HACK_XML is not enabled.
       <mam4_xs_long_file type="file" doc="File containing photolysis data"> ${DIN_LOC_ROOT}/atm/scream/mam4xx/photolysis/temp_prs_GT200nm_JPL10_c130206.nc</mam4_xs_long_file>
       <!--Elevated emissions-->
       <elevated_emiss_ymd type="integer"> 20100101 </elevated_emiss_ymd>
+      <!--Dry gas deposition-->
+      <mam4_season_wes_file type="file" doc="File containing season_wes data"> ${DIN_LOC_ROOT}/atm/scream/mam4xx/drydep/season_wes.nc</mam4_season_wes_file>
       <!-- For all other grids -->
       <mam4_so2_elevated_emiss_file_name type="file" doc="elevated emissions for so2">${DIN_LOC_ROOT}/atm/scream/mam4xx/emissions/ne30pg2/elevated/cmip6_mam4_so2_elev_1x1_2010_clim_ne30pg2_c20241008.nc</mam4_so2_elevated_emiss_file_name>
       <mam4_so4_a1_elevated_emiss_file_name type="file" doc="elevated emissions for so4_a1">${DIN_LOC_ROOT}/atm/scream/mam4xx/emissions/ne30pg2/elevated/cmip6_mam4_so4_a1_elev_1x1_2010_clim_ne30pg2_c20241008.nc</mam4_so4_a1_elevated_emiss_file_name>

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

@@ -3,6 +3,7 @@
 // impl namespace for some driver level functions for microphysics
 
 #include "readfiles/photo_table_utils.cpp"
+#include "readfiles/find_season_index_utils.hpp"
 #include "physics/rrtmgp/shr_orb_mod_c2f.hpp"
 
 namespace scream {
@@ -332,6 +333,14 @@ void MAMMicrophysics::set_grids(
         "MAX_NUM_ELEVATED_EMISSIONS_FIELDS in tracer_reader_utils.hpp \n");
 
   }  // Tracer external forcing data
+
+  {
+    const std::string season_wes_file = m_params.get<std::string>("mam4_season_wes_file");
+    const auto& clat = col_latitudes_;
+    mam_coupling::find_season_index_reader(season_wes_file,
+                                         clat,
+                                         index_season_lai_);
+  }
 }  // set_grids
 
 // ================================================================

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

@@ -25,6 +25,8 @@ class MAMMicrophysics final : public scream::AtmosphereProcess {
 
   using view_1d_host = typename KT::view_1d<Real>::HostMirror;
 
+  using view_int_2d       = typename KT::template view_2d<int>;
+
   // a thread team dispatched to a single vertical column
   using ThreadTeam = mam4::ThreadTeam;
 
@@ -239,6 +241,8 @@ class MAMMicrophysics final : public scream::AtmosphereProcess {
   view_1d_host acos_cosine_zenith_host_;
   view_1d acos_cosine_zenith_;
 
+  view_int_2d index_season_lai_;
+
 };  // MAMMicrophysics
 
 }  // namespace scream

components/eamxx/src/physics/mam/readfiles/find_season_index_utils.hpp

@@ -0,0 +1,75 @@
+#ifndef EAMXX_MAM_FIND_SEASON_INDEX_UTILS
+#define EAMXX_MAM_FIND_SEASON_INDEX_UTILS
+
+#include <ekat/kokkos/ekat_kokkos_utils.hpp>
+#include <mam4xx/mam4.hpp>
+
+#include "share/io/scorpio_input.hpp"
+#include "share/io/scream_scorpio_interface.hpp"
+
+namespace scream::mam_coupling {
+
+// views for single- and multi-column data
+
+using const_view_1d = typename KT::template view_1d<const Real>;
+using view_int_2d   = typename KT::template view_2d<int>;
+
+using view_1d_host     = typename KT::view_1d<Real>::HostMirror;
+using view_int_3d_host = typename KT::view_3d<int>::HostMirror;
+using view_int_2d_host = typename KT::view_2d<int>::HostMirror;
+
+/**
+ * @brief Reads the season index from the given file and computes the season
+ * indices based on latitudes.
+ *
+ * @param[in] season_wes_file The path to the season_wes.nc file.
+ * @param[in] clat A 1D view of latitude values in degrees.
+ * @param[out] index_season_lai A 2D view to store the computed season indices.
+ * Note that indices are in C++ (starting from zero).
+ */
+
+inline void find_season_index_reader(const std::string &season_wes_file,
+                                     const const_view_1d &clat,
+                                     view_int_2d &index_season_lai) {
+  const int plon = clat.extent(0);
+  scorpio::register_file(season_wes_file, scorpio::Read);
+
+  const int nlat_lai = scorpio::get_dimlen(season_wes_file, "lat");
+  const int npft_lai = scorpio::get_dimlen(season_wes_file, "pft");
+
+  view_1d_host lat_lai("lat_lai", nlat_lai);
+  view_int_2d_host wk_lai_temp("wk_lai", npft_lai, nlat_lai);
+  view_int_3d_host wk_lai("wk_lai", nlat_lai, npft_lai, 12);
+
+  scorpio::read_var(season_wes_file, "lat", lat_lai.data());
+
+  Kokkos::MDRangePolicy<Kokkos::HostSpace::execution_space, Kokkos::Rank<2>>
+      policy_wk_lai({0, 0}, {nlat_lai, npft_lai});
+
+  // loop over time to get all 12 instantence of season_wes
+  for(int itime = 0; itime < 12; ++itime) {
+    scorpio::read_var(season_wes_file, "season_wes", wk_lai_temp.data(), itime);
+    // copy data from wk_lai_temp to wk_lai.
+    // NOTE: season_wes has different layout that wk_lai
+    Kokkos::parallel_for("copy_to_wk_lai", policy_wk_lai,
+                         [&](const int j, const int k) {
+                           wk_lai(j, k, itime) = wk_lai_temp(k, j);
+                         });
+    Kokkos::fence();
+  }
+  scorpio::release_file(season_wes_file);
+
+  index_season_lai = view_int_2d("index_season_lai", plon, 12);
+  const view_int_2d_host index_season_lai_host =
+      Kokkos::create_mirror_view(index_season_lai);
+
+  const view_1d_host clat_host = Kokkos::create_mirror_view(clat);
+  Kokkos::deep_copy(clat_host, clat);
+
+  // Computation is performed on the host
+  mam4::mo_drydep::find_season_index(clat_host, lat_lai, nlat_lai, wk_lai,
+                                     index_season_lai_host);
+  Kokkos::deep_copy(index_season_lai, index_season_lai_host);
+}
+}  // namespace scream::mam_coupling
+#endif  //

components/eamxx/tests/single-process/mam/aero_microphys/CMakeLists.txt

@@ -39,6 +39,7 @@ set (TEST_INPUT_FILES
      scream/mam4xx/emissions/ne2np4/elevated/cmip6_mam4_num_a2_elev_ne2np4_2010_clim_c20240823.nc
      scream/mam4xx/emissions/ne2np4/elevated/cmip6_mam4_num_a4_elev_ne2np4_2010_clim_c20240823.nc
      scream/mam4xx/emissions/ne2np4/elevated/cmip6_mam4_soag_elev_ne2np4_2010_clim_c20240823.nc
+     scream/mam4xx/drydep/season_wes.nc
 )
 foreach (file IN ITEMS ${TEST_INPUT_FILES})
   GetInputFile(${file})

components/eamxx/tests/single-process/mam/aero_microphys/input.yaml

@@ -37,7 +37,7 @@ atmosphere_processes:
    mam4_num_a2_elevated_emiss_file_name : ${SCREAM_DATA_DIR}/mam4xx/emissions/ne2np4/elevated/cmip6_mam4_num_a2_elev_ne2np4_2010_clim_c20240823.nc
    mam4_num_a4_elevated_emiss_file_name : ${SCREAM_DATA_DIR}/mam4xx/emissions/ne2np4/elevated/cmip6_mam4_num_a4_elev_ne2np4_2010_clim_c20240823.nc
    mam4_soag_elevated_emiss_file_name : ${SCREAM_DATA_DIR}/mam4xx/emissions/ne2np4/elevated/cmip6_mam4_soag_elev_ne2np4_2010_clim_c20240823.nc
-
+   mam4_season_wes_file : ${SCREAM_DATA_DIR}/mam4xx/drydep/season_wes.nc
 grids_manager:
   Type: Mesh Free
   geo_data_source: IC_FILE