Merge branch 'bartgol/eamxx/rad-cosp-freq-adjustments' (PR #7337)

Currently, rad runs on steps 1, rad_freq+1, 2*rad_freq+1, etc.,
which is not what one would expect, and makes it hard to craft
an instant output yaml file to catch rad steps (if freq>1).

This PR makes the frequency logic of rad similar to the output,
with the only exception that we always run it on the 1st step
(regardless of frequency).

[non-BFB] for EAMxx cases

Author: bartgol

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, find_node
+        resolve_all_inheritances, gen_atm_proc_group, check_all_values
 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
@@ -29,7 +29,6 @@
 # Cime imports
 from standard_script_setup import * # pylint: disable=wildcard-import
 from CIME.utils import expect, safe_copy, SharedArea
-from CIME.test_status import TestStatus, RUN_PHASE
 
 logger = logging.getLogger(__name__) # pylint: disable=undefined-variable
 
@@ -105,121 +104,6 @@ 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 (3 steps) incompatible with restart frequency (2 steps).
-     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")
-    is_test = case.get_value("TEST")
-    caseraw = case.get_value("CASE")
-    caseroot = case.get_value("CASEROOT")
-    casebaseid  = case.get_value("CASEBASEID")
-    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
-
-        # Some tests (ERS) make late (run-phase) changes, so we cannot validate restart
-        # settings until RUN phase
-        is_test_not_yet_run = False
-        if is_test:
-            test_name = casebaseid if casebaseid is not None else caseraw
-            ts = TestStatus(test_dir=caseroot, test_name=test_name)
-            phase = ts.get_latest_phase()
-            if phase != RUN_PHASE:
-                is_test_not_yet_run = True
-
-        if rad_freq==1 or is_test_not_yet_run:
-            pass
-        elif rest_opt in ["nsteps", "nstep"]:
-            expect (rest_n % rad_freq == 0,
-                    f"rrtmgp::rad_frequency ({rad_freq} steps) incompatible with "
-                    f"restart frequency ({rest_n} steps).\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):
 ###############################################################################
@@ -743,8 +627,6 @@ def _create_raw_xml_file_impl(case, xml, filepath=None):
 
         raise e
 
-    perform_consistency_checks (case, xml)
-
     return xml
 
 ###############################################################################

components/eamxx/cime_config/namelist_defaults_eamxx.xml

@@ -551,6 +551,7 @@ be lost if SCREAM_HACK_XML is not enabled.
           true
       </do_subcol_sampling>
       <pool_size_multiplier type="real">1.0</pool_size_multiplier>
+      <force_run_after_restart type="logical" doc="Force rad to run on first step after restart, regardless of rad frequency">false</force_run_after_restart>
     </rrtmgp>
 
     <mac_aero_mic inherit="atm_proc_group">

components/eamxx/src/control/atmosphere_driver.cpp

@@ -1620,9 +1620,11 @@ void AtmosphereDriver::run (const int dt) {
   EKAT_REQUIRE_MSG (dt>0, "Error! Input time step must be positive.\n");
 
   // Print current timestamp information
+  auto end_of_step = m_current_ts + dt;
   m_atm_logger->log(ekat::logger::LogLevel::info,
-    "Atmosphere step = " + std::to_string(m_current_ts.get_num_steps()) + "\n" +
-    "  model start-of-step time = " + m_current_ts.get_date_string() + " " + m_current_ts.get_time_string() + "\n");
+    "Atmosphere step = " + std::to_string(end_of_step.get_num_steps()) + "\n" +
+    "  model beg-of-step timestamp: " + m_current_ts.get_date_string() + " " + m_current_ts.get_time_string() + "\n"
+    "  model end-of-step timestamp: " + end_of_step.get_date_string() + " " + end_of_step.get_time_string() + "\n");
 
   // Reset accum fields to 0
   // Note: at the 1st timestep this is redundant, since we did it at init,

components/eamxx/src/physics/cosp/eamxx_cosp.hpp

@@ -32,15 +32,13 @@ class Cosp : public AtmosphereProcess
   // Set the grid
   void set_grids (const std::shared_ptr<const GridsManager> grids_manager);
 
-  inline bool cosp_do(const int icosp, const int nstep) {
+  inline bool cosp_do(const int cosp_freq, const int nstep) {
       // If icosp == 0, then never do cosp;
-      // Otherwise, we always call cosp at the first step,
-      // and afterwards we do cosp if the timestep is divisible
-      // by icosp
-      if (icosp == 0) {
+      // Otherwise, do cosp if the timestep is divisible by cosp_freq
+      if (cosp_freq == 0) {
           return false;
       } else {
-          return ( (nstep == 0) || (nstep % icosp == 0) );
+          return nstep % cosp_freq == 0;
       }
   }
 

components/eamxx/src/physics/rrtmgp/eamxx_rrtmgp_process_interface.cpp

@@ -79,6 +79,9 @@ RRTMGPRadiation (const ekat::Comm& comm, const ekat::ParameterList& params)
   }
 
   m_ngas = m_gas_names.size();
+
+  // Determine rad timestep, specified as number of atm steps
+  m_rad_freq_in_steps = m_params.get<Int>("rad_frequency", 1);
 }
 
 void RRTMGPRadiation::set_grids(const std::shared_ptr<const GridsManager> grids_manager) {
@@ -198,7 +201,6 @@ void RRTMGPRadiation::set_grids(const std::shared_ptr<const GridsManager> grids_
   add_field<Computed>("LW_clrsky_flux_dn", scalar3d_int, W/m2, grid_name);
   add_field<Computed>("LW_clnsky_flux_up", scalar3d_int, W/m2, grid_name);
   add_field<Computed>("LW_clnsky_flux_dn", scalar3d_int, W/m2, grid_name);
-  add_field<Computed>("rad_heating_pdel", scalar3d_mid, Pa*K/s, grid_name);
   // Cloud properties added as computed fields for diagnostic purposes
   add_field<Computed>("cldlow"        , scalar2d, nondim, grid_name);
   add_field<Computed>("cldmed"        , scalar2d, nondim, grid_name);
@@ -219,6 +221,11 @@ void RRTMGPRadiation::set_grids(const std::shared_ptr<const GridsManager> grids_
   add_field<Computed>("eff_radius_qc_at_cldtop", scalar2d, micron, grid_name);
   add_field<Computed>("eff_radius_qi_at_cldtop", scalar2d, micron, grid_name);
 
+  // This field is needed for restart
+  Field rad_heating_pdel (FieldIdentifier("rad_heating_pdel", scalar3d_mid, Pa*K/s, grid_name));
+  rad_heating_pdel.allocate_view();
+  add_internal_field(rad_heating_pdel);
+
   // Translation of variables from EAM
   // --------------------------------------------------------------
   // EAM name | EAMXX name       | Description
@@ -230,12 +237,19 @@ void RRTMGPRadiation::set_grids(const std::shared_ptr<const GridsManager> grids_
   // netsw      sfc_flux_sw_net    net (down - up) SW flux at surface
   // flwds      sfc_flux_lw_dn     downwelling LW flux at surface
   // --------------------------------------------------------------
-  add_field<Computed>("sfc_flux_dir_nir", scalar2d, W/m2, grid_name);
-  add_field<Computed>("sfc_flux_dir_vis", scalar2d, W/m2, grid_name);
-  add_field<Computed>("sfc_flux_dif_nir", scalar2d, W/m2, grid_name);
-  add_field<Computed>("sfc_flux_dif_vis", scalar2d, W/m2, grid_name);
-  add_field<Computed>("sfc_flux_sw_net" , scalar2d, W/m2, grid_name);
-  add_field<Computed>("sfc_flux_lw_dn"  , scalar2d, W/m2, grid_name);
+
+  // We need to ensure that these fields are added to the RESTART group,
+  // since the cpl will need them at every step, and rrtmgp may not run
+  // the 1st step after restart (depending on rad freq).
+  // NOTE: technically, we know rad freq, so we *could* avoid adding them
+  //       to the rest file if rad_freq=1. But a) that is not common at high
+  //       res anyways, and b) that could prevent changing rad_freq upon restart
+  add_field<Computed>("sfc_flux_dir_nir", scalar2d, W/m2, grid_name, "RESTART");
+  add_field<Computed>("sfc_flux_dir_vis", scalar2d, W/m2, grid_name, "RESTART");
+  add_field<Computed>("sfc_flux_dif_nir", scalar2d, W/m2, grid_name, "RESTART");
+  add_field<Computed>("sfc_flux_dif_vis", scalar2d, W/m2, grid_name, "RESTART");
+  add_field<Computed>("sfc_flux_sw_net" , scalar2d, W/m2, grid_name, "RESTART");
+  add_field<Computed>("sfc_flux_lw_dn"  , scalar2d, W/m2, grid_name, "RESTART");
 
   // Boundary flux fields for energy and mass conservation checks
   if (has_column_conservation_check()) {
@@ -453,12 +467,9 @@ void RRTMGPRadiation::init_buffers(const ATMBufferManager &buffer_manager)
   EKAT_REQUIRE_MSG(used_mem==requested_buffer_size_in_bytes(), "Error! Used memory != requested memory for RRTMGPRadiation.");
 } // RRTMGPRadiation::init_buffers
 
-void RRTMGPRadiation::initialize_impl(const RunType /* run_type */) {
+void RRTMGPRadiation::initialize_impl(const RunType run_type) {
   using PC = scream::physics::Constants<Real>;
 
-  // Determine rad timestep, specified as number of atm steps
-  m_rad_freq_in_steps = m_params.get<Int>("rad_frequency", 1);
-
   // Determine orbital year. If orbital_year is negative, use current year
   // from timestamp for orbital year; if positive, use provided orbital year
   // for duration of simulation.
@@ -532,6 +543,13 @@ void RRTMGPRadiation::initialize_impl(const RunType /* run_type */) {
     auto co_vmr = get_field_out("co_volume_mix_ratio").get_view<Real**>();
     Kokkos::deep_copy(co_vmr, m_params.get<double>("covmr", 1.0e-7));
   }
+
+  // Ensure rad_heating_pdel is recognized as initialized by the driver
+  auto& rad_heating = get_internal_field("rad_heating_pdel");
+  rad_heating.get_header().get_tracking().update_time_stamp(start_of_step_ts());
+
+  m_force_run_on_next_step = run_type==RunType::Initial or
+    m_params.get("force_run_after_restart",false);
 }
 
 // =========================================================================================
@@ -597,7 +615,7 @@ void RRTMGPRadiation::run_impl (const double dt) {
   auto d_lw_clrsky_flux_dn = get_field_out("LW_clrsky_flux_dn").get_view<Real**>();
   auto d_lw_clnsky_flux_up = get_field_out("LW_clnsky_flux_up").get_view<Real**>();
   auto d_lw_clnsky_flux_dn = get_field_out("LW_clnsky_flux_dn").get_view<Real**>();
-  auto d_rad_heating_pdel = get_field_out("rad_heating_pdel").get_view<Real**>();
+  auto d_rad_heating_pdel = get_internal_field("rad_heating_pdel").get_view<Real**>();
   auto d_sfc_flux_dir_vis = get_field_out("sfc_flux_dir_vis").get_view<Real*>();
   auto d_sfc_flux_dir_nir = get_field_out("sfc_flux_dir_nir").get_view<Real*>();
   auto d_sfc_flux_dif_vis = get_field_out("sfc_flux_dif_vis").get_view<Real*>();
@@ -638,8 +656,8 @@ void RRTMGPRadiation::run_impl (const double dt) {
   const auto do_aerosol_rad = m_do_aerosol_rad;
 
   // Are we going to update fluxes and heating this step?
-  auto ts = start_of_step_ts();
-  auto update_rad = scream::rrtmgp::radiation_do(m_rad_freq_in_steps, ts.get_num_steps());
+  auto ts = end_of_step_ts();
+  auto update_rad = m_force_run_on_next_step or scream::rrtmgp::radiation_do(m_rad_freq_in_steps, ts.get_num_steps());
 
   if (update_rad) {
     // On each chunk, we internally "reset" the GasConcs object to subview the concs 3d array
@@ -1211,6 +1229,7 @@ void RRTMGPRadiation::run_impl (const double dt) {
     });
   }
 
+  m_force_run_on_next_step = false;
 }
 // =========================================================================================
 

components/eamxx/src/physics/rrtmgp/eamxx_rrtmgp_process_interface.hpp

@@ -240,6 +240,8 @@ class RRTMGPRadiation : public AtmosphereProcess {
 
   // Struct which contains local variables
   Buffer m_buffer;
+
+  bool m_force_run_on_next_step = false;
 };  // class RRTMGPRadiation
 
 }  // namespace scream

components/eamxx/src/physics/rrtmgp/rrtmgp_utils.hpp

@@ -37,15 +37,15 @@ void compute_heating_rate (
                                   ));
 }
 
-inline bool radiation_do(const int irad, const int nstep) {
-  // If irad == 0, then never do radiation;
+inline bool radiation_do(const int rad_freq, const int nstep) {
+  // If rad_freq == 0, then never do radiation;
   // Otherwise, we always call radiation at the first step,
   // and afterwards we do radiation if the timestep is divisible
-  // by irad
-  if (irad == 0) {
+  // by rad_freq
+  if (rad_freq == 0) {
     return false;
   } else {
-    return ( (nstep == 0) || (nstep % irad == 0) );
+    return nstep % rad_freq == 0;
   }
 }
 

components/eamxx/src/physics/rrtmgp/tests/rrtmgp_unit_tests.cpp

@@ -411,25 +411,14 @@ TEST_CASE("rrtmgp_test_compute_broadband_surface_flux_k") {
 }
 
 TEST_CASE("rrtmgp_test_radiation_do_k") {
-  // If we specify rad every step, radiation_do should always be true
-  REQUIRE(scream::rrtmgp::radiation_do(1, 0) == true);
-  REQUIRE(scream::rrtmgp::radiation_do(1, 1) == true);
-  REQUIRE(scream::rrtmgp::radiation_do(1, 2) == true);
-
-    // Test cases where we want rad called every other step
-  REQUIRE(scream::rrtmgp::radiation_do(2, 0) == true);
-  REQUIRE(scream::rrtmgp::radiation_do(2, 1) == false);
-  REQUIRE(scream::rrtmgp::radiation_do(2, 2) == true);
-  REQUIRE(scream::rrtmgp::radiation_do(2, 3) == false);
-
-  // Test cases where we want rad every third step
-  REQUIRE(scream::rrtmgp::radiation_do(3, 0) == true);
-  REQUIRE(scream::rrtmgp::radiation_do(3, 1) == false);
-  REQUIRE(scream::rrtmgp::radiation_do(3, 2) == false);
-  REQUIRE(scream::rrtmgp::radiation_do(3, 3) == true);
-  REQUIRE(scream::rrtmgp::radiation_do(3, 4) == false);
-  REQUIRE(scream::rrtmgp::radiation_do(3, 5) == false);
-  REQUIRE(scream::rrtmgp::radiation_do(3, 6) == true);
+  // Rad runs whenever step is a multiple of freq
+  // NOTE: the rrtmgp process class handles logic for running on 1st step
+  for (int istep : {1,2,3,4,5,6}) {
+    for (int rad_freq : {1,2,3}) {
+      bool divides = istep%rad_freq == 0;
+      REQUIRE( scream::rrtmgp::radiation_do(rad_freq,istep)== divides );
+    }
+  }
 }
 
 TEST_CASE("rrtmgp_test_check_range_k") {

components/eamxx/src/share/atm_process/atmosphere_process.cpp

@@ -130,7 +130,7 @@ void AtmosphereProcess::run (const double dt) {
       // Print hash of INPUTS before run
       print_global_state_hash(name() + "-pre-sc-" + std::to_string(m_subcycle_iter),
                               m_start_of_step_ts,
-                              true, false, false);
+                              true, false, true);
 
     // Run derived class implementation
     run_impl(dt_sub);
@@ -139,7 +139,7 @@ void AtmosphereProcess::run (const double dt) {
       // Print hash of OUTPUTS/INTERNALS after run
       print_global_state_hash(name() + "-pst-sc-" + std::to_string(m_subcycle_iter),
                               m_end_of_step_ts,
-                              true, true, true);
+                              false, true, true);
 
     if (has_column_conservation_check()) {
       // Run the column local mass and energy conservation checks