Merge pull request #2636 from E3SM-Project/bartgol/active-gases-pg2-a…

…nd-restart-fixes

Some fixes related to radiation interface and restarts

components/eamxx/cime_config/eamxx_buildnml.py

@@ -15,7 +15,7 @@
 
 # SCREAM imports
 from eamxx_buildnml_impl import get_valid_selectors, get_child, refine_type, \
-        resolve_all_inheritances, gen_atm_proc_group, check_all_values
+        resolve_all_inheritances, gen_atm_proc_group, check_all_values, find_node
 from atm_manip import apply_atm_procs_list_changes_from_buffer, apply_non_atm_procs_list_changes_from_buffer
 
 from utils import ensure_yaml # pylint: disable=no-name-in-module
@@ -104,6 +104,107 @@ def do_cime_vars(entry, case, refine=False, extra=None):
 
     return entry
 
+###############################################################################
+def perform_consistency_checks(case, xml):
+###############################################################################
+    """
+    There may be separate parts of the xml that must satisfy some consistency
+    Here, we run any such check, so we can catch errors before submit time
+
+    >>> from eamxx_buildnml_impl import MockCase
+    >>> xml_str = '''
+    ... <params>
+    ...   <rrtmgp>
+    ...     <rad_frequency type="integer">3</rad_frequency>
+    ...   </rrtmgp>
+    ... </params>
+    ... '''
+    >>> import xml.etree.ElementTree as ET
+    >>> xml = ET.fromstring(xml_str)
+    >>> case = MockCase({'ATM_NCPL':'24', 'REST_N':24, 'REST_OPTION':'nsteps'})
+    >>> perform_consistency_checks(case,xml)
+    >>> case = MockCase({'ATM_NCPL':'24', 'REST_N':2, 'REST_OPTION':'nsteps'})
+    >>> perform_consistency_checks(case,xml)
+    Traceback (most recent call last):
+    CIME.utils.CIMEError: ERROR: rrtmgp::rad_frequency incompatible with restart frequency.
+     Please, ensure restart happens on a step when rad is ON
+    >>> case = MockCase({'ATM_NCPL':'24', 'REST_N':10800, 'REST_OPTION':'nseconds'})
+    >>> perform_consistency_checks(case,xml)
+    >>> case = MockCase({'ATM_NCPL':'24', 'REST_N':7200, 'REST_OPTION':'nseconds'})
+    >>> perform_consistency_checks(case,xml)
+    Traceback (most recent call last):
+    CIME.utils.CIMEError: ERROR: rrtmgp::rad_frequency incompatible with restart frequency.
+     Please, ensure restart happens on a step when rad is ON
+      rest_tstep: 7200
+      rad_testep: 10800.0
+    >>> case = MockCase({'ATM_NCPL':'24', 'REST_N':180, 'REST_OPTION':'nminutes'})
+    >>> perform_consistency_checks(case,xml)
+    >>> case = MockCase({'ATM_NCPL':'24', 'REST_N':120, 'REST_OPTION':'nminutes'})
+    >>> perform_consistency_checks(case,xml)
+    Traceback (most recent call last):
+    CIME.utils.CIMEError: ERROR: rrtmgp::rad_frequency incompatible with restart frequency.
+     Please, ensure restart happens on a step when rad is ON
+      rest_tstep: 7200
+      rad_testep: 10800.0
+    >>> case = MockCase({'ATM_NCPL':'24', 'REST_N':6, 'REST_OPTION':'nhours'})
+    >>> perform_consistency_checks(case,xml)
+    >>> case = MockCase({'ATM_NCPL':'24', 'REST_N':8, 'REST_OPTION':'nhours'})
+    >>> perform_consistency_checks(case,xml)
+    Traceback (most recent call last):
+    CIME.utils.CIMEError: ERROR: rrtmgp::rad_frequency incompatible with restart frequency.
+     Please, ensure restart happens on a step when rad is ON
+      rest_tstep: 28800
+      rad_testep: 10800.0
+    >>> case = MockCase({'ATM_NCPL':'12', 'REST_N':2, 'REST_OPTION':'ndays'})
+    >>> perform_consistency_checks(case,xml)
+    >>> case = MockCase({'ATM_NCPL':'10', 'REST_N':2, 'REST_OPTION':'ndays'})
+    >>> perform_consistency_checks(case,xml)
+    Traceback (most recent call last):
+    CIME.utils.CIMEError: ERROR: rrtmgp::rad_frequency incompatible with restart frequency.
+     Please, ensure restart happens on a step when rad is ON
+     For daily (or less frequent) restart, rad_frequency must divide ATM_NCPL
+    """
+
+    # RRTMGP can be supercycled. Restarts cannot fall in the middle
+    # of a rad superstep
+    rrtmgp = find_node(xml,"rrtmgp")
+    rest_opt = case.get_value("REST_OPTION")
+    if rrtmgp is not None and rest_opt is not None and rest_opt not in ["never","none"]:
+        rest_n = int(case.get_value("REST_N"))
+        rad_freq = int(find_node(rrtmgp,"rad_frequency").text)
+        atm_ncpl = int(case.get_value("ATM_NCPL"))
+        atm_tstep = 86400 / atm_ncpl
+        rad_tstep = atm_tstep * rad_freq
+
+
+        if rad_freq==1:
+            pass
+        elif rest_opt in ["nsteps", "nstep"]:
+            expect (rest_n % rad_freq == 0,
+                    "rrtmgp::rad_frequency incompatible with restart frequency.\n"
+                    " Please, ensure restart happens on a step when rad is ON")
+        elif rest_opt in ["nseconds", "nsecond", "nminutes", "nminute", "nhours", "nhour"]:
+            if rest_opt in ["nseconds", "nsecond"]:
+                factor = 1
+            elif rest_opt in ["nminutes", "nminute"]:
+                factor = 60
+            else:
+                factor = 3600
+
+            rest_tstep = factor*rest_n
+            expect (rest_tstep % rad_tstep == 0,
+                    "rrtmgp::rad_frequency incompatible with restart frequency.\n"
+                    " Please, ensure restart happens on a step when rad is ON\n"
+                    f"  rest_tstep: {rest_tstep}\n"
+                    f"  rad_testep: {rad_tstep}")
+
+        else:
+            # for "very infrequent" restarts, we request rad_freq to divide atm_ncpl
+            expect (atm_ncpl % rad_freq ==0,
+                    "rrtmgp::rad_frequency incompatible with restart frequency.\n"
+                    " Please, ensure restart happens on a step when rad is ON\n"
+                    " For daily (or less frequent) restart, rad_frequency must divide ATM_NCPL")
+
 ###############################################################################
 def ordered_dump(data, item, Dumper=yaml.SafeDumper, **kwds):
 ###############################################################################
@@ -536,6 +637,8 @@ def _create_raw_xml_file_impl(case, xml):
     #    which atm processes are used
     apply_non_atm_procs_list_changes_from_buffer (case,xml)
 
+    perform_consistency_checks (case, xml)
+
     return xml
 
 ###############################################################################

components/eamxx/src/control/atmosphere_driver.cpp

@@ -127,13 +127,6 @@ set_params(const ekat::ParameterList& atm_params)
   create_logger ();
 
   m_ad_status |= s_params_set;
-
-  const auto pg_type = "PG2";
-  fvphyshack = m_atm_params.sublist("grids_manager").get<std::string>("physics_grid_type", "None") == pg_type;
-  if (fvphyshack) {
-    // See the [rrtmgp active gases] note in share/util/eamxx_fv_phys_rrtmgp_active_gases_workaround.hpp
-    fv_phys_rrtmgp_active_gases_init(m_atm_params);
-  }
 }
 
 void AtmosphereDriver::

components/eamxx/src/diagnostics/longwave_cloud_forcing.cpp

@@ -1,5 +1,7 @@
 #include "diagnostics/longwave_cloud_forcing.hpp"
 
+#include <ekat/kokkos/ekat_kokkos_utils.hpp>
+
 namespace scream
 {
 
@@ -25,11 +27,10 @@ void LongwaveCloudForcingDiagnostic::set_grids(const std::shared_ptr<const Grids
 
   FieldLayout scalar2d_layout_col{ {COL}, {m_num_cols} };
   FieldLayout scalar3d_layout_mid { {COL,LEV}, {m_num_cols,m_num_levs} };
-  constexpr int ps = Pack::n;
 
   // The fields required for this diagnostic to be computed
-  add_field<Required>("LW_flux_up",          scalar3d_layout_mid, W/m2,  grid_name, ps);
-  add_field<Required>("LW_clrsky_flux_up",          scalar3d_layout_mid, W/m2,  grid_name, ps);
+  add_field<Required>("LW_flux_up",        scalar3d_layout_mid, W/m2,  grid_name);
+  add_field<Required>("LW_clrsky_flux_up", scalar3d_layout_mid, W/m2,  grid_name);
 
   // Construct and allocate the diagnostic field
   FieldIdentifier fid (name(), scalar2d_layout_col, W/m2, grid_name);
@@ -41,19 +42,23 @@ void LongwaveCloudForcingDiagnostic::set_grids(const std::shared_ptr<const Grids
 // =========================================================================================
 void LongwaveCloudForcingDiagnostic::compute_diagnostic_impl()
 {
-  const auto default_policy = ekat::ExeSpaceUtils<KT::ExeSpace>::get_default_team_policy(m_num_cols,1);
+  using KT         = KokkosTypes<DefaultDevice>;
+  using ESU        = ekat::ExeSpaceUtils<KT::ExeSpace>;
+  using MemberType = typename KT::MemberType;
+
+  const auto default_policy = ESU::get_default_team_policy(m_num_cols,1);
 
   const auto& LWCF              = m_diagnostic_output.get_view<Real*>();
-  const auto& LW_flux_up        = get_field_in("LW_flux_up").get_view<const Pack**>();
-  const auto& LW_clrsky_flux_up = get_field_in("LW_clrsky_flux_up").get_view<const Pack**>();
+  const auto& LW_flux_up        = get_field_in("LW_flux_up").get_view<const Real**>();
+  const auto& LW_clrsky_flux_up = get_field_in("LW_clrsky_flux_up").get_view<const Real**>();
 
   Kokkos::parallel_for("LongwaveCloudForcingDiagnostic",
                        default_policy,
                        KOKKOS_LAMBDA(const MemberType& team) {
     const int icol = team.league_rank();
-    LWCF(icol) =  LW_clrsky_flux_up(icol,0)[0] -  LW_flux_up(icol,0)[0] ;
+    LWCF(icol) =  LW_clrsky_flux_up(icol,0) -  LW_flux_up(icol,0) ;
   });
   Kokkos::fence();
 }
-// =========================================================================================
+
 } //namespace scream

components/eamxx/src/diagnostics/longwave_cloud_forcing.hpp

@@ -2,8 +2,6 @@
 #define EAMXX_LONGWAVE_CLOUD_FORCING_DIAGNOSTIC_HPP
 
 #include "share/atm_process/atmosphere_diagnostic.hpp"
-#include "share/util/scream_common_physics_functions.hpp"
-#include "ekat/kokkos/ekat_subview_utils.hpp"
 
 namespace scream
 {
@@ -15,13 +13,6 @@ namespace scream
 class LongwaveCloudForcingDiagnostic : public AtmosphereDiagnostic
 {
 public:
-  using Pack          = ekat::Pack<Real,SCREAM_PACK_SIZE>;
-  using PF            = scream::PhysicsFunctions<DefaultDevice>;
-
-  using KT            = KokkosTypes<DefaultDevice>;
-  using MemberType    = typename KT::MemberType;
-  using view_1d       = typename KT::template view_1d<Pack>;
-
   // Constructors
   LongwaveCloudForcingDiagnostic (const ekat::Comm& comm, const ekat::ParameterList& params);
 

components/eamxx/src/diagnostics/shortwave_cloud_forcing.cpp

@@ -1,5 +1,7 @@
 #include "diagnostics/shortwave_cloud_forcing.hpp"
 
+#include <ekat/kokkos/ekat_kokkos_utils.hpp>
+
 namespace scream
 {
 
@@ -24,13 +26,12 @@ void ShortwaveCloudForcingDiagnostic::set_grids(const std::shared_ptr<const Grid
 
   FieldLayout scalar2d_layout_col{ {COL}, {m_num_cols} };
   FieldLayout scalar3d_layout_mid { {COL,LEV}, {m_num_cols,m_num_levs} };
-  constexpr int ps = Pack::n;
 
   // The fields required for this diagnostic to be computed
-  add_field<Required>("SW_flux_dn",          scalar3d_layout_mid, W/m2,  grid_name, ps);
-  add_field<Required>("SW_flux_up",          scalar3d_layout_mid, W/m2,  grid_name, ps);
-  add_field<Required>("SW_clrsky_flux_dn",          scalar3d_layout_mid, W/m2,  grid_name, ps);
-  add_field<Required>("SW_clrsky_flux_up",          scalar3d_layout_mid, W/m2,  grid_name, ps);
+  add_field<Required>("SW_flux_dn",        scalar3d_layout_mid, W/m2, grid_name);
+  add_field<Required>("SW_flux_up",        scalar3d_layout_mid, W/m2, grid_name);
+  add_field<Required>("SW_clrsky_flux_dn", scalar3d_layout_mid, W/m2, grid_name);
+  add_field<Required>("SW_clrsky_flux_up", scalar3d_layout_mid, W/m2, grid_name);
 
   // Construct and allocate the diagnostic field
   FieldIdentifier fid (name(), scalar2d_layout_col, W/m2, grid_name);
@@ -42,19 +43,23 @@ void ShortwaveCloudForcingDiagnostic::set_grids(const std::shared_ptr<const Grid
 
 void ShortwaveCloudForcingDiagnostic::compute_diagnostic_impl()
 {
-  const auto default_policy = ekat::ExeSpaceUtils<KT::ExeSpace>::get_default_team_policy(m_num_cols,1);
+  using KT         = KokkosTypes<DefaultDevice>;
+  using ESU        = ekat::ExeSpaceUtils<KT::ExeSpace>;
+  using MemberType = typename KT::MemberType;
+
+  const auto default_policy = ESU::get_default_team_policy(m_num_cols,1);
 
   const auto& SWCF              = m_diagnostic_output.get_view<Real*>();
-  const auto& SW_flux_dn        = get_field_in("SW_flux_dn").get_view<const Pack**>();
-  const auto& SW_flux_up        = get_field_in("SW_flux_up").get_view<const Pack**>();
-  const auto& SW_clrsky_flux_dn = get_field_in("SW_clrsky_flux_dn").get_view<const Pack**>();
-  const auto& SW_clrsky_flux_up = get_field_in("SW_clrsky_flux_up").get_view<const Pack**>();
+  const auto& SW_flux_dn        = get_field_in("SW_flux_dn").get_view<const Real**>();
+  const auto& SW_flux_up        = get_field_in("SW_flux_up").get_view<const Real**>();
+  const auto& SW_clrsky_flux_dn = get_field_in("SW_clrsky_flux_dn").get_view<const Real**>();
+  const auto& SW_clrsky_flux_up = get_field_in("SW_clrsky_flux_up").get_view<const Real**>();
 
   Kokkos::parallel_for("ShortwaveCloudForcingDiagnostic",
                        default_policy,
                        KOKKOS_LAMBDA(const MemberType& team) {
     const int icol = team.league_rank();
-    SWCF(icol) = (SW_flux_dn(icol,0)[0] - SW_flux_up(icol,0)[0]) - (SW_clrsky_flux_dn(icol,0)[0] - SW_clrsky_flux_up(icol,0)[0]);
+    SWCF(icol) = (SW_flux_dn(icol,0) - SW_flux_up(icol,0)) - (SW_clrsky_flux_dn(icol,0) - SW_clrsky_flux_up(icol,0));
   });
   Kokkos::fence();
 }

components/eamxx/src/diagnostics/shortwave_cloud_forcing.hpp

@@ -2,8 +2,6 @@
 #define EAMXX_SHORTWAVE_CLOUD_FORCING_DIAGNOSTIC_HPP
 
 #include "share/atm_process/atmosphere_diagnostic.hpp"
-#include "share/util/scream_common_physics_functions.hpp"
-#include "ekat/kokkos/ekat_subview_utils.hpp"
 
 namespace scream
 {
@@ -15,13 +13,6 @@ namespace scream
 class ShortwaveCloudForcingDiagnostic : public AtmosphereDiagnostic
 {
 public:
-  using Pack          = ekat::Pack<Real,SCREAM_PACK_SIZE>;
-  using PF            = scream::PhysicsFunctions<DefaultDevice>;
-
-  using KT            = KokkosTypes<DefaultDevice>;
-  using MemberType    = typename KT::MemberType;
-  using view_1d       = typename KT::template view_1d<Pack>;
-
   // Constructors
   ShortwaveCloudForcingDiagnostic (const ekat::Comm& comm, const ekat::ParameterList& params);
 

components/eamxx/src/dynamics/homme/homme_grids_manager.cpp

@@ -3,6 +3,8 @@
 #include "dynamics/homme/physics_dynamics_remapper.hpp"
 #include "dynamics/homme/homme_dynamics_helpers.hpp"
 
+#include "share/util/eamxx_fv_phys_rrtmgp_active_gases_workaround.hpp"
+
 #ifndef NDEBUG
 #include "share/property_checks/field_nan_check.hpp"
 #include "share/property_checks/field_lower_bound_check.hpp"
@@ -197,6 +199,10 @@ build_physics_grid (const ci_string& type, const ci_string& rebalance) {
     return;
   }
 
+  if (type=="PG2") {
+    fvphyshack = true;
+  }
+
   // Get the grid pg_type
   const int pg_code = m_pg_codes.at(type).at(rebalance);