Merge branch 'eagles-project:mjs271/eamxx/DAG-refactor' (PR #6860)

Finishes PR #3101 originally submitted to the scream repo the atmospheric DAG generator (atmosphere_process_dag.xpp) and atmosphere driver to get things working again

Author: bartgol

components/eamxx/src/control/atmosphere_driver.cpp

@@ -4,6 +4,7 @@
 
 #include "physics/share/physics_constants.hpp"
 
+#include "share/scream_config.hpp"
 #include "share/atm_process/atmosphere_process_group.hpp"
 #include "share/atm_process/atmosphere_process_dag.hpp"
 #include "share/field/field_utils.hpp"
@@ -685,10 +686,29 @@ void AtmosphereDriver::create_fields()
     fm->add_to_group(fid.name(),"RESTART");
   }
 
+  auto& driver_options_pl = m_atm_params.sublist("driver_options");
+  const int verb_lvl = driver_options_pl.get<int>("atmosphere_dag_verbosity_level",-1);
+  if (verb_lvl>0) {
+    // now that we've got fields, generate a DAG with fields and dependencies
+    // NOTE: at this point, fields provided by initial conditions may (will)
+    // appear as unmet dependencies
+    AtmProcDAG dag;
+    // First, add all atm processes
+    dag.create_dag(*m_atm_process_group);
+    // Write a dot file for visualizing the DAG
+    if (m_atm_comm.am_i_root()) {
+      std::string filename = "scream_atm_createField_dag";
+      if (is_scream_standalone()) {
+        filename += ".np" + std::to_string(m_atm_comm.size());
+      }
+      filename += ".dot";
+      dag.write_dag(filename, verb_lvl);
+    }
+  }
+
   m_ad_status |= s_fields_created;
 
   // If the user requested it, we can save a dictionary of the FM fields to file
-  auto& driver_options_pl = m_atm_params.sublist("driver_options");
   if (driver_options_pl.get("save_field_manager_content",false)) {
     auto pg = m_grids_manager->get_grid("Physics");
     const auto& fm = m_field_mgrs.at(pg->name());
@@ -902,28 +922,6 @@ initialize_fields ()
     TraceGasesWorkaround::singleton().run_type = m_run_type;
   }
 
-  // See if we need to print a DAG. We do this first, cause if any input
-  // field is missing from the initial condition file, an error will be thrown.
-  // By printing the DAG first, we give the user the possibility of seeing
-  // what fields are inputs to the atm time step, so he/she can fix the i.c. file.
-  // TODO: would be nice to do the IC input first, and mark the fields in the
-  //       DAG node "Begin of atm time step" in red if there's no initialization
-  //       mechanism set for them. That is, allow field XYZ to not be found in
-  //       the IC file, and throw an error when the dag is created.
-
-  auto& driver_options_pl = m_atm_params.sublist("driver_options");
-  const int verb_lvl = driver_options_pl.get<int>("atmosphere_dag_verbosity_level",-1);
-  if (verb_lvl>0) {
-    // Check the atm DAG for missing stuff
-    AtmProcDAG dag;
-
-    // First, add all atm processes
-    dag.create_dag(*m_atm_process_group);
-
-    // Write a dot file for visualization
-    dag.write_dag("scream_atm_dag.dot",std::max(verb_lvl,0));
-  }
-
   // Initialize fields
   if (m_run_type==RunType::Restart) {
     restart_model ();
@@ -1101,7 +1099,6 @@ void AtmosphereDriver::set_initial_conditions ()
   // Check which fields need to have an initial condition.
   std::map<std::string,std::vector<std::string>> ic_fields_names;
   std::vector<FieldIdentifier> ic_fields_to_copy;
-  std::map<std::string,std::vector<std::string>> fields_inited;
 
   // Check which fields should be loaded from the topography file
   std::map<std::string,std::vector<std::string>> topography_file_fields_names;
@@ -1130,7 +1127,7 @@ void AtmosphereDriver::set_initial_conditions ()
         EKAT_ERROR_MSG ("ERROR: invalid assignment for variable " + fname + ", only scalar "
                         "double or string, or vector double arguments are allowed");
       }
-      fields_inited[grid_name].push_back(fname);
+      m_fields_inited[grid_name].push_back(fname);
     } else if (fname == "phis" or fname == "sgh30") {
       // Both phis and sgh30 need to be loaded from the topography file
       auto& this_grid_topo_file_fnames = topography_file_fields_names[grid_name];
@@ -1146,7 +1143,7 @@ void AtmosphereDriver::set_initial_conditions ()
                    grid_name == "Point Grid") {
           this_grid_topo_file_fnames.push_back("PHIS_d");
           this_grid_topo_eamxx_fnames.push_back(fname);
-          fields_inited[grid_name].push_back(fname);
+          m_fields_inited[grid_name].push_back(fname);
         } else {
           EKAT_ERROR_MSG ("Error! Requesting phis on an unknown grid: " + grid_name + ".\n");
         }
@@ -1158,7 +1155,7 @@ void AtmosphereDriver::set_initial_conditions ()
                         " topo file only has sgh30 for Physics PG2.\n");
         topography_file_fields_names[grid_name].push_back("SGH30");
         topography_eamxx_fields_names[grid_name].push_back(fname);
-	fields_inited[grid_name].push_back(fname);
+        m_fields_inited[grid_name].push_back(fname);
       }
     } else if (not (fvphyshack and grid_name == "Physics PG2")) {
       // The IC file is written for the GLL grid, so we only load
@@ -1170,7 +1167,7 @@ void AtmosphereDriver::set_initial_conditions ()
         // If this field is the parent of other subfields, we only read from file the subfields.
         if (not ekat::contains(this_grid_ic_fnames,fname)) {
           this_grid_ic_fnames.push_back(fname);
-	  fields_inited[grid_name].push_back(fname);
+          m_fields_inited[grid_name].push_back(fname);
         }
       } else if (fvphyshack and grid_name == "Physics GLL") {
         // [CGLL ICs in pg2] I tried doing something like this in
@@ -1187,7 +1184,7 @@ void AtmosphereDriver::set_initial_conditions ()
           } else {
             this_grid_ic_fnames.push_back(fname);
           }
-	  fields_inited[grid_name].push_back(fname);
+          m_fields_inited[grid_name].push_back(fname);
         }
       }
     }
@@ -1233,7 +1230,7 @@ void AtmosphereDriver::set_initial_conditions ()
         auto p = f.get_header().get_parent().lock();
         if (p) {
           const auto& pname = p->get_identifier().name();
-          if (ekat::contains(fields_inited[grid_name],pname)) {
+          if (ekat::contains(m_fields_inited[grid_name],pname)) {
             // The parent is already inited. No need to init this field as well.
             names.erase(it2);
             run_again = true;
@@ -1456,7 +1453,7 @@ void AtmosphereDriver::set_initial_conditions ()
     // Loop through fields and apply perturbation.
     for (size_t f=0; f<perturbed_fields.size(); ++f) {
       const auto fname = perturbed_fields[f];
-      EKAT_REQUIRE_MSG(ekat::contains(fields_inited[fm->get_grid()->name()], fname),
+      EKAT_REQUIRE_MSG(ekat::contains(m_fields_inited[fm->get_grid()->name()], fname),
                        "Error! Attempting to apply perturbation to field not in initial_conditions.\n"
                        "  - Field: "+fname+"\n"
                        "  - Grid:  "+fm->get_grid()->name()+"\n");
@@ -1655,6 +1652,28 @@ void AtmosphereDriver::initialize_atm_procs ()
   m_atm_logger->info("[EAMxx] initialize_atm_procs ... done!");
 
   report_res_dep_memory_footprint ();
+
+  auto& driver_options_pl = m_atm_params.sublist("driver_options");
+  const int verb_lvl = driver_options_pl.get<int>("atmosphere_dag_verbosity_level",-1);
+  if (verb_lvl>0) {
+    // now that we've got fields, generate a DAG with fields and dependencies
+    // NOTE: at this point, fields provided by initial conditions may (will)
+    // appear as unmet dependencies
+    AtmProcDAG dag;
+    // First, add all atm processes
+    dag.create_dag(*m_atm_process_group);
+    // process the initial conditions to maybe fulfill unmet dependencies
+    dag.process_initial_conditions(m_fields_inited);
+    // Write a dot file for visualizing the DAG
+    if (m_atm_comm.am_i_root()) {
+      std::string filename = "scream_atm_initProc_dag";
+      if (is_scream_standalone()) {
+        filename += ".np" + std::to_string(m_atm_comm.size());
+      }
+      filename += ".dot";
+      dag.write_dag(filename, verb_lvl);
+    }
+  }
 }
 
 void AtmosphereDriver::

components/eamxx/src/control/atmosphere_driver.hpp

@@ -264,6 +264,8 @@ class AtmosphereDriver
 
   // Current simulation casename
   std::string m_casename;
+  // maps grid name to a vector of its initialized fields
+  std::map<std::string, std::vector<std::string>> m_fields_inited;
 };
 
 }  // namespace control