Merge branch 'bartgol/eamxx/iop-remapper' into next (PR #6914)

Author: bartgol

components/eamxx/src/control/atmosphere_driver.cpp

@@ -1001,11 +1001,14 @@ void AtmosphereDriver::restart_model ()
       continue;
     }
     const auto& restart_group = it.second->get_groups_info().at("RESTART");
-    std::vector<std::string> fnames;
+    std::vector<Field> fields;
     for (const auto& fn : restart_group->m_fields_names) {
-      fnames.push_back(fn);
+      fields.push_back(it.second->get_field(fn));
+    }
+    read_fields_from_file (fields,it.second->get_grid(),filename);
+    for (auto& f : fields) {
+      f.get_header().get_tracking().update_time_stamp(m_current_ts);
     }
-    read_fields_from_file (fnames,it.second->get_grid(),filename,m_current_ts);
   }
 
   // Restart the num steps counter in the atm time stamp
@@ -1267,17 +1270,20 @@ void AtmosphereDriver::set_initial_conditions ()
     // Now loop over all grids, and load from file the needed fields on each grid (if any).
     const auto& file_name = ic_pl.get<std::string>("Filename");
     m_atm_logger->info("    [EAMxx] IC filename: " + file_name);
-    for (const auto& it : m_field_mgrs) {
-      const auto& grid_name = it.first;
+    for (const auto& [grid_name,fm] : m_field_mgrs) {
+      std::vector<Field> ic_fields;
+      for (const auto& fn : ic_fields_names[grid_name]) {
+        ic_fields.push_back(fm->get_field(fn));
+      }
       if (not m_iop_data_manager) {
-        read_fields_from_file (ic_fields_names[grid_name],it.second->get_grid(),file_name,m_current_ts);
+        read_fields_from_file (ic_fields,fm->get_grid(),file_name);
       } else {
         // For IOP enabled, we load from file and copy data from the closest
         // lat/lon column to every other column
-        m_iop_data_manager->read_fields_from_file_for_iop(file_name,
-                                             ic_fields_names[grid_name],
-                                             m_current_ts,
-                                             it.second);
+        m_iop_data_manager->read_fields_from_file_for_iop(file_name,ic_fields,fm->get_grid());
+      }
+      for (auto& f : ic_fields) {
+        f.get_header().get_tracking().update_time_stamp(m_current_ts);
       }
     }
   }
@@ -1342,29 +1348,33 @@ void AtmosphereDriver::set_initial_conditions ()
     m_atm_logger->info("    [EAMxx] Reading topography from file ...");
     const auto& file_name = ic_pl.get<std::string>("topography_filename");
     m_atm_logger->info("        filename: " + file_name);
-    for (const auto& it : m_field_mgrs) {
-      const auto& grid_name = it.first;
+    for (const auto& [grid_name,fm] : m_field_mgrs) {
+      std::vector<Field> topo_fields;
+      int nfields = topography_eamxx_fields_names[grid_name].size();
+      for (int i=0; i<nfields; ++i) {
+        auto eamxx_fname = topography_eamxx_fields_names[grid_name][i];
+        auto file_fname  = topography_file_fields_names[grid_name][i];
+        topo_fields.push_back(fm->get_field(eamxx_fname).alias(file_fname));
+      }
+
       if (not m_iop_data_manager) {
         // Topography files always use "ncol_d" for the GLL grid value of ncol.
         // To ensure we read in the correct value, we must change the name for that dimension
-        auto io_grid = it.second->get_grid();
+        auto io_grid = fm->get_grid();
         if (grid_name=="Physics GLL") {
           using namespace ShortFieldTagsNames;
           auto grid = io_grid->clone(io_grid->name(),true);
           grid->reset_field_tag_name(COL,"ncol_d");
           io_grid = grid;
         }
-        read_fields_from_file (topography_file_fields_names[grid_name],
-                               topography_eamxx_fields_names[grid_name],
-                               io_grid,file_name,m_current_ts);
+        read_fields_from_file (topo_fields,io_grid,file_name);
       } else {
         // For IOP enabled, we load from file and copy data from the closest
         // lat/lon column to every other column
-        m_iop_data_manager->read_fields_from_file_for_iop(file_name,
-                                             topography_file_fields_names[grid_name],
-                                             topography_eamxx_fields_names[grid_name],
-                                             m_current_ts,
-                                             it.second);
+        m_iop_data_manager->read_fields_from_file_for_iop(file_name,topo_fields,fm->get_grid());
+      }
+      for (auto& f : topo_fields) {
+        f.get_header().get_tracking().update_time_stamp(m_current_ts);
       }
     }
     // Store in provenance list, for later usage in output file metadata
@@ -1469,83 +1479,17 @@ void AtmosphereDriver::set_initial_conditions ()
 }
 
 void AtmosphereDriver::
-read_fields_from_file (const std::vector<std::string>& field_names_nc,
-                       const std::vector<std::string>& field_names_eamxx,
-                       const std::shared_ptr<const AbstractGrid>& grid,
-                       const std::string& file_name,
-                       const util::TimeStamp& t0)
-{
-  EKAT_REQUIRE_MSG(field_names_nc.size()==field_names_eamxx.size(),
-                   "Error! Field name arrays must have same size.\n");
-
-  if (field_names_nc.size()==0) {
-    return;
-  }
-
-  // NOTE: we cannot pass the field_mgr and m_grids_mgr, since the input
-  //       grid may not be in the grids_manager and may not be the grid
-  //       of the field mgr. This sounds weird, but there is a precise
-  //       use case: when grid is a shallow clone of the fm grid, where
-  //       we changed the name of some field tags (e.g., we set the name
-  //       of COL to ncol_d). This is used when reading the topography,
-  //       since the topo file *always* uses ncol_d for GLL points data,
-  //       while a non-PG2 run would have the tag name be "ncol".
-  const auto& field_mgr = m_field_mgrs.at(grid->name());
-  std::vector<Field> fields;
-  for (size_t i=0; i<field_names_nc.size(); ++i) {
-    const auto& eamxx_name = field_names_eamxx[i];
-    const auto& nc_name    = field_names_nc[i];
-    fields.push_back(field_mgr->get_field(eamxx_name).alias(nc_name));
-  }
-
-  AtmosphereInput ic_reader(file_name,grid,fields);
-  ic_reader.set_logger(m_atm_logger);
-  ic_reader.read_variables();
-  ic_reader.finalize();
-
-  for (auto& f : fields) {
-    // Set the initial time stamp
-    // NOTE: f is an alias of the field from field_mgr, so it shares all
-    //       pointers to the metadata (except for the FieldIdentifier),
-    //       so changing its timestamp will also change the timestamp
-    //       of the field in field_mgr
-    f.get_header().get_tracking().update_time_stamp(t0);
-  }
-}
-
-void AtmosphereDriver::
-read_fields_from_file (const std::vector<std::string>& field_names,
+read_fields_from_file (const std::vector<Field>& fields,
                        const std::shared_ptr<const AbstractGrid>& grid,
-                       const std::string& file_name,
-                       const util::TimeStamp& t0)
+                       const std::string& file_name)
 {
-  if (field_names.size()==0) {
+  if (fields.size()==0) {
     return;
   }
 
-  // NOTE: we cannot pass the field_mgr and m_grids_mgr, since the input
-  //       grid may not be in the grids_manager and may not be the grid
-  //       of the field mgr. This sounds weird, but there is a precise
-  //       use case: when grid is a shallow clone of the fm grid, where
-  //       we changed the name of some field tags (e.g., we set the name
-  //       of COL to ncol_d). This is used when reading the topography,
-  //       since the topo file *always* uses ncol_d for GLL points data,
-  //       while a non-PG2 run would have the tag name be "ncol".
-  const auto& field_mgr = m_field_mgrs.at(grid->name());
-  std::vector<Field> fields;
-  for (const auto& fn : field_names) {
-    fields.push_back(field_mgr->get_field(fn));
-  }
-
   AtmosphereInput ic_reader(file_name,grid,fields);
   ic_reader.set_logger(m_atm_logger);
   ic_reader.read_variables();
-  ic_reader.finalize();
-
-  for (auto& f : fields) {
-    // Set the initial time stamp
-    f.get_header().get_tracking().update_time_stamp(t0);
-  }
 }
 
 void AtmosphereDriver::

components/eamxx/src/control/atmosphere_driver.hpp

@@ -185,21 +185,10 @@ class AtmosphereDriver
   void set_initial_conditions ();
   void restart_model ();
 
-  // Read fields from a file when the names of the fields in
-  // EAMxx do not match exactly with the .nc file. Example is
-  // for topography data files, where GLL and PG2 grid have
-  // different naming conventions for phis.
-  void read_fields_from_file (const std::vector<std::string>& field_names_nc,
-                              const std::vector<std::string>& field_names_eamxx,
+  // Read fields from a file
+  void read_fields_from_file (const std::vector<Field>& fields,
                               const std::shared_ptr<const AbstractGrid>& grid,
-                              const std::string& file_name,
-                              const util::TimeStamp& t0);
-  // Read fields from a file when the names of the fields in
-  // EAMxx match with the .nc file.
-  void read_fields_from_file (const std::vector<std::string>& field_names,
-                              const std::shared_ptr<const AbstractGrid>& grid,
-                              const std::string& file_name,
-                              const util::TimeStamp& t0);
+                              const std::string& file_name);
   void register_groups ();
 
   std::map<std::string,field_mgr_ptr>       m_field_mgrs;

components/eamxx/src/physics/spa/eamxx_spa_process_interface.cpp

@@ -60,15 +60,12 @@ void SPA::set_grids(const std::shared_ptr<const GridsManager> grids_manager)
   auto spa_data_file = m_params.get<std::string>("spa_data_file");
   auto spa_map_file  = m_params.get<std::string>("spa_remap_file","");
 
-  // IOP cases cannot have a remap file. IOP file reader is itself a remaper,
-  // where a single column of data corresponding to the closest lat/lon pair to
-  // the IOP lat/lon parameters is read from file, and that column data is mapped
-  // to all columns of the IdentityRemapper source fields.
+  // IOP cases cannot have a remap file. We will create a IOPRemapper as the horiz remapper
   EKAT_REQUIRE_MSG(spa_map_file == "" or spa_map_file == "None" or not m_iop_data_manager,
     "Error! Cannot define spa_remap_file for cases with an Intensive Observation Period defined. "
     "The IOP class defines it's own remap from file data -> model data.\n");
 
-  SPAHorizInterp = SPAFunc::create_horiz_remapper (m_grid,spa_data_file,spa_map_file, m_iop_data_manager!=nullptr);
+  SPAHorizInterp = SPAFunc::create_horiz_remapper (m_grid,spa_data_file,spa_map_file, m_iop_data_manager);
 
   // Grab a sw and lw field from the horiz interp, and check sw/lw dim against what we hardcoded in this class
   auto nswbands_data = SPAHorizInterp->get_src_field(4).get_header().get_identifier().get_layout().dim("swband");
@@ -124,15 +121,8 @@ void SPA::set_grids(const std::shared_ptr<const GridsManager> grids_manager)
     Kokkos::deep_copy(spa_hybm,spa_hybm_h);
   }
 
-  // 4. Create reader for spa data. The reader is either an
-  //    AtmosphereInput object (for reading into standard
-  //    grids) or a SpaFunctions::IOPReader (for reading into
-  //    an IOP grid).
-  if (m_iop_data_manager) {
-    SPAIOPDataReader = SPAFunc::create_spa_data_reader(m_iop_data_manager,SPAHorizInterp,spa_data_file);
-  } else {
-    SPADataReader = SPAFunc::create_spa_data_reader(SPAHorizInterp,spa_data_file);
-  }
+  // 4. Create reader for spa data.
+  SPADataReader = SPAFunc::create_spa_data_reader(SPAHorizInterp,spa_data_file);
 }
 
 // =========================================================================================
@@ -224,7 +214,7 @@ void SPA::initialize_impl (const RunType /* run_type */)
   // Note: At the first time step, the data will be moved into spa_beg,
   //       and spa_end will be reloaded from file with the new month.
   const int curr_month = timestamp().get_month()-1; // 0-based
-  SPAFunc::update_spa_data_from_file(SPADataReader,SPAIOPDataReader,timestamp(),curr_month,*SPAHorizInterp,SPAData_end);
+  SPAFunc::update_spa_data_from_file(*SPADataReader,curr_month,*SPAHorizInterp,SPAData_end);
 
   // 6. Set property checks for fields in this process
   using Interval = FieldWithinIntervalCheck;
@@ -246,7 +236,7 @@ void SPA::run_impl (const double dt)
   /* Update the SPATimeState to reflect the current time, note the addition of dt */
   SPATimeState.t_now = ts.frac_of_year_in_days();
   /* Update time state and if the month has changed, update the data.*/
-    SPAFunc::update_spa_timestate(SPADataReader,SPAIOPDataReader,ts,*SPAHorizInterp,SPATimeState,SPAData_start,SPAData_end);
+    SPAFunc::update_spa_timestate(*SPADataReader,ts,*SPAHorizInterp,SPATimeState,SPAData_start,SPAData_end);
 
   // Call the main SPA routine to get interpolated aerosol forcings.
   const auto& pmid_tgt = get_field_in("p_mid").get_view<const Spack**>();

components/eamxx/src/physics/spa/eamxx_spa_process_interface.hpp

@@ -79,8 +79,6 @@ class SPA : public AtmosphereProcess
 
   // IO structure to read in data for standard grids (keep it around to avoid re-creating PIO decomps)
   std::shared_ptr<AtmosphereInput>   SPADataReader;
-  // Similar to above, but stores info to read data for IOP grid
-  std::shared_ptr<SPAFunc::IOPReader>  SPAIOPDataReader;
 
   // Structures to store the data used for interpolation
   std::shared_ptr<AbstractRemapper>  SPAHorizInterp;

components/eamxx/src/physics/spa/spa_functions.hpp

@@ -30,8 +30,6 @@ struct SPAFunctions
 
   using gid_type = AbstractGrid::gid_type;
 
-  using iop_data_ptr_type = std::shared_ptr<control::IOPDataManager>;
-
   template <typename S>
   using view_1d = typename KT::template view_1d<S>;
   template <typename S>
@@ -127,33 +125,6 @@ struct SPAFunctions
     SPAData         data;         // All spa fields
   }; // SPAInput
 
-  struct IOPReader {
-    IOPReader (iop_data_ptr_type& iop_,
-               const std::string file_name_,
-               const std::vector<Field>& io_fields_,
-               const std::shared_ptr<const AbstractGrid>& io_grid_)
-      : iop_data_manager(iop_), file_name(file_name_)
-    {
-      field_mgr = std::make_shared<FieldManager>(io_grid_);
-      for (auto& f : io_fields_) {
-        field_mgr->add_field(f);
-        field_names.push_back(f.name());
-      }
-
-      // Set IO info for this grid and file in IOP object
-      iop_data_manager->setup_io_info(file_name, io_grid_);
-    }
-
-    void read_variables(const int time_index, const util::TimeStamp& ts) {
-      iop_data_manager->read_fields_from_file_for_iop(file_name, field_names, ts, field_mgr, time_index);
-    }
-
-    iop_data_ptr_type iop_data_manager;
-    std::string file_name;
-    std::vector<std::string> field_names;
-    std::shared_ptr<FieldManager> field_mgr;
-  };
-
   // The output is really just SPAData, but for clarity it might
   // help to see a SPAOutput along a SPAInput in functions signatures
   using SPAOutput = SPAData;
@@ -166,19 +137,13 @@ struct SPAFunctions
       const std::shared_ptr<const AbstractGrid>& model_grid,
       const std::string& spa_data_file,
       const std::string& map_file,
-      const bool use_iop = false);
+      const std::shared_ptr<control::IOPDataManager>& iop_data_mgr = nullptr);
 
   static std::shared_ptr<AtmosphereInput>
   create_spa_data_reader (
       const std::shared_ptr<AbstractRemapper>& horiz_remapper,
       const std::string& spa_data_file);
 
-  static std::shared_ptr<IOPReader>
-  create_spa_data_reader (
-      iop_data_ptr_type& iop_data_manager,
-      const std::shared_ptr<AbstractRemapper>& horiz_remapper,
-      const std::string& spa_data_file);
-
   static void spa_main(
     const SPATimeState& time_state,
     const view_2d<const Spack>& p_tgt,
@@ -189,21 +154,18 @@ struct SPAFunctions
     const SPAOutput&  data_out);
 
   static void update_spa_data_from_file(
-    std::shared_ptr<AtmosphereInput>& scorpio_reader,
-    std::shared_ptr<IOPReader>&       iop_reader,
-    const util::TimeStamp&            ts,
-    const int                         time_index, // zero-based
-    AbstractRemapper&                 spa_horiz_interp,
-    SPAInput&                         spa_input);
+    AtmosphereInput&    scorpio_reader,
+    const int           time_index, // zero-based
+    AbstractRemapper&   spa_horiz_interp,
+    SPAInput&           spa_input);
 
   static void update_spa_timestate(
-    std::shared_ptr<AtmosphereInput>& scorpio_reader,
-    std::shared_ptr<IOPReader>&       iop_reader,
-    const util::TimeStamp&            ts,
-    AbstractRemapper&                 spa_horiz_interp,
-    SPATimeState&                     time_state,
-    SPAInput&                         spa_beg,
-    SPAInput&                         spa_end);
+    AtmosphereInput&          scorpio_reader,
+    const util::TimeStamp&    ts,
+    AbstractRemapper&         spa_horiz_interp,
+    SPATimeState&             time_state,
+    SPAInput&                 spa_beg,
+    SPAInput&                 spa_end);
 
   // The following three are called during spa_main
   static void perform_time_interpolation (