Updates to python workflow due to PR ufs-community#776

tests/WE2E/get_WE2Etest_names_subdirs_descs.sh

@@ -438,9 +438,6 @@ function get_WE2Etest_names_subdirs_descs() {
 #
   ushdir=$( readlink -f "$WE2Edir/../../ush" )
   . $ushdir/constants.sh
-  . $ushdir/set_predef_grid_params.sh
-  . $ushdir/set_gridparams_GFDLgrid.sh
-  . $ushdir/set_gridparams_ESGgrid.sh
 #
 #-----------------------------------------------------------------------
 #
@@ -1220,70 +1217,16 @@ ${test_desc}${stripped_line} "
 # parameters.  The way the latter are obtained depends on whether or not
 # a predefined grid is being used.
 #
-# If using a predefined grid, call the set_predef_grid_params() function
-# to get the grid parameters.
-#
-      if [ ! -z "${PREDEF_GRID_NAME}" ]; then
-#
-# Note:
-# Can set "quilting" to "FALSE" in the following argument list because
-# the write-component parameters are not needed below; only those of the
-# native grid are needed.
-#
-        set_predef_grid_params \
-          predef_grid_name="${PREDEF_GRID_NAME}" \
-          quilting="FALSE" \
-          outvarname_grid_gen_method="grid_gen_method" \
-          outvarname_esggrid_lon_ctr="esggrid_lon_ctr" \
-          outvarname_esggrid_lat_ctr="esggrid_lat_ctr" \
-          outvarname_esggrid_delx="esggrid_delx" \
-          outvarname_esggrid_dely="esggrid_dely" \
-          outvarname_esggrid_nx="esggrid_nx" \
-          outvarname_esggrid_ny="esggrid_ny" \
-          outvarname_esggrid_pazi="esggrid_pazi" \
-          outvarname_esggrid_wide_halo_width="esggrid_wide_halo_width" \
-          outvarname_gfdlgrid_lon_t6_ctr="gfdlgrid_lon_t6_ctr" \
-          outvarname_gfdlgrid_lat_t6_ctr="gfdlgrid_lat_t6_ctr" \
-          outvarname_gfdlgrid_stretch_fac="gfdlgrid_stretch_fac" \
-          outvarname_gfdlgrid_res="gfdlgrid_res" \
-          outvarname_gfdlgrid_refine_ratio="gfdlgrid_refine_ratio" \
-          outvarname_gfdlgrid_istart_of_rgnl_dom_on_t6g="gfdlgrid_istart_of_rgnl_dom_on_t6g" \
-          outvarname_gfdlgrid_iend_of_rgnl_dom_on_t6g="gfdlgrid_iend_of_rgnl_dom_on_t6g" \
-          outvarname_gfdlgrid_jstart_of_rgnl_dom_on_t6g="gfdlgrid_jstart_of_rgnl_dom_on_t6g" \
-          outvarname_gfdlgrid_jend_of_rgnl_dom_on_t6g="gfdlgrid_jend_of_rgnl_dom_on_t6g" \
-          outvarname_dt_atmos="dta"
-#
-# If using a custom grid, the test's configuration file should contain 
-# the grid parameters.  Source this file and set the values of the grid
-# parameters it contains to local variables.
-#
-      else
-
-        . ./${config_fn}
-        grid_gen_method="${GRID_GEN_METHOD}"
-        if [ "${grid_gen_method}" = "GFDLgrid" ]; then
-          gfdlgrid_lon_t6_ctr="${GFDLgrid_LON_T6_CTR}"
-          gfdlgrid_lat_t6_ctr="${GFDLgrid_LAT_T6_CTR}"
-          gfdlgrid_res="${GFDLgrid_RES}"
-          gfdlgrid_stretch_fac="${GFDLgrid_STRETCH_FAC}"
-          gfdlgrid_refine_ratio="${GFDLgrid_REFINE_RATIO}"
-          gfdlgrid_istart_of_rgnl_dom_on_t6g="${GFDLgrid_ISTART_OF_RGNL_DOM_ON_T6G}"
-          gfdlgrid_iend_of_rgnl_dom_on_t6g="${GFDLgrid_IEND_OF_RGNL_DOM_ON_T6G}"
-          gfdlgrid_jstart_of_rgnl_dom_on_t6g="${GFDLgrid_JSTART_OF_RGNL_DOM_ON_T6G}"
-          gfdlgrid_jend_of_rgnl_dom_on_t6g="${GFDLgrid_JEND_OF_RGNL_DOM_ON_T6G}"
-        elif [ "${grid_gen_method}" = "ESGgrid" ]; then
-          esggrid_lon_ctr="${ESGgrid_LON_CTR}"
-          esggrid_lat_ctr="${ESGgrid_LAT_CTR}"
-          esggrid_delx="${ESGgrid_DELX}"
-          esggrid_dely="${ESGgrid_DELY}"
-          esggrid_nx="${ESGgrid_NX}"
-          esggrid_ny="${ESGgrid_NY}"
-          esggrid_pazi="${ESGgrid_PAZI}"
-          esggrid_wide_halo_width="${ESGgrid_WIDE_HALO_WIDTH}"
-        fi
-        dta="${DT_ATMOS}"
+params=$(\
+  PREDEF_GRID_NAME="${PREDEF_GRID_NAME}" \
+  QUILTING="FALSE" \
+  RADIUS_EARTH=${RADIUS_EARTH} \
+  DEGS_PER_RADIAN=${DEGS_PER_RADIAN} \
+  NH4=${NH4} \
+    $ushdir/calculate_cost.py -c "${test_configs_basedir}/$subdir/${config_fn}")
+
+read dta nxny dta_r nxny_r <<< "${params}"
 
-      fi
 #
 # Save the value of dta (which is just dt_atmos) in an array.  The single 
 # quote at the beginning forces Google Sheets to interpret this quantity 
@@ -1292,49 +1235,9 @@ ${test_desc}${stripped_line} "
 #
       prim_test_dt_atmos+=( "'${dta}" )
 #
-# The way the number of grid points in the horizontal directions (nx and
-# ny) are calculated depends on the method used to generate the grid as 
-# well as the grid parameters for that method.
-#
-      if [ "${grid_gen_method}" = "GFDLgrid" ]; then
-#
-# Note:
-# The workflow generation mode (run_envir) can be set to "community" here
-# since this does not affect the values of nx and ny.
-#
-        set_gridparams_GFDLgrid \
-          lon_of_t6_ctr="${gfdlgrid_lon_t6_ctr}" \
-          lat_of_t6_ctr="${gfdlgrid_lat_t6_ctr}" \
-          res_of_t6g="${gfdlgrid_res}" \
-          stretch_factor="${gfdlgrid_stretch_fac}" \
-          refine_ratio_t6g_to_t7g="${gfdlgrid_refine_ratio}" \
-          istart_of_t7_on_t6g="${gfdlgrid_istart_of_rgnl_dom_on_t6g}" \
-          iend_of_t7_on_t6g="${gfdlgrid_iend_of_rgnl_dom_on_t6g}" \
-          jstart_of_t7_on_t6g="${gfdlgrid_jstart_of_rgnl_dom_on_t6g}" \
-          jend_of_t7_on_t6g="${gfdlgrid_jend_of_rgnl_dom_on_t6g}" \
-          verbose="$verbose" \
-          outvarname_nx_of_t7_on_t7g="nx" \
-          outvarname_ny_of_t7_on_t7g="ny"
-
-      elif [ "${grid_gen_method}" = "ESGgrid" ]; then
-
-        set_gridparams_ESGgrid \
-          lon_ctr="${esggrid_lon_ctr}" \
-          lat_ctr="${esggrid_lat_ctr}" \
-          nx="${esggrid_nx}" \
-          ny="${esggrid_ny}" \
-          pazi="${esggrid_pazi}" \
-          halo_width="${esggrid_wide_halo_width}" \
-          delx="${esggrid_delx}" \
-          dely="${esggrid_dely}" \
-          outvarname_nx="nx" \
-          outvarname_ny="ny"
-
-      fi
-#
 # Calculate the total number of horizontal grid points.
 #
-      num_grid_pts=$(( nx*ny ))
+      num_grid_pts=$nxny
 #
 # Calculate the number of time steps for the test.  Note that FCST_LEN_HRS 
 # is in units of hours while dta is in units of seconds.  Also, the factor
@@ -1348,30 +1251,6 @@ ${test_desc}${stripped_line} "
 #
       ac=$(( num_grid_pts*num_time_steps*nf ))
 #
-# Unset all grid paramters so that they are not accidentally reused for
-# the next test.
-#
-      unset gfdlgrid_lon_t6_ctr \
-            gfdlgrid_lat_t6_ctr \
-            gfdlgrid_res \
-            gfdlgrid_stretch_fac \
-            gfdlgrid_refine_ratio \
-            gfdlgrid_istart_of_rgnl_dom_on_t6g \
-            gfdlgrid_iend_of_rgnl_dom_on_t6g \
-            gfdlgrid_jstart_of_rgnl_dom_on_t6g \
-            gfdlgrid_jend_of_rgnl_dom_on_t6g \
-            esggrid_lon_ctr \
-            esggrid_lat_ctr \
-            esggrid_nx \
-            esggrid_ny \
-            esggrid_pazi \
-            esggrid_wide_halo_width \
-            esggrid_delx \
-            esggrid_dely \
-            dta \
-            nx \
-            ny
-#
 # Save the absolute cost for this test in the array that will eventually
 # contain the relative cost.  The values in this array will be divided
 # by abs_cost_ref later below to obtain relative costs.
@@ -1397,16 +1276,9 @@ ${test_desc}${stripped_line} "
 #
 #-----------------------------------------------------------------------
 #
-    set_predef_grid_params \
-      predef_grid_name="RRFS_CONUS_25km" \
-      quilting="FALSE" \
-      outvarname_esggrid_nx="nx" \
-      outvarname_esggrid_ny="ny" \
-      outvarname_dt_atmos="dta"
-
-    num_grid_pts=$(( nx*ny ))
+    num_grid_pts=$nxny_r
     fcst_len_hrs="6"
-    num_time_steps=$(( (fcst_len_hrs*3600 + dta - 1)/dta ))
+    num_time_steps=$(( (fcst_len_hrs*3600 + dta_r - 1)/dta_r ))
     abs_cost_ref=$(( num_grid_pts*num_time_steps ))
 
     for (( i=0; i<=$((num_prim_tests-1)); i++ )); do

ush/calculate_cost.py

@@ -0,0 +1,106 @@
+#!/usr/bin/env python3
+
+import os
+import unittest
+import argparse
+
+from python_utils import set_env_var, import_vars, export_vars, load_config_file, flatten_dict
+
+from set_predef_grid_params import set_predef_grid_params
+from set_gridparams_ESGgrid import set_gridparams_ESGgrid
+from set_gridparams_GFDLgrid import set_gridparams_GFDLgrid
+
+def calculate_cost(config_fn):
+    global PREDEF_GRID_NAME, QUILTING, GRID_GEN_METHOD
+
+    #import all environment variables
+    import_vars()
+
+    #get grid config parameters (predefined or custom)
+    if PREDEF_GRID_NAME:
+        set_env_var('QUILTING',False)
+        set_predef_grid_params()
+        import_vars()
+    else:
+        cfg_u = load_config_file(config_fn)
+        cfg_u = flatten_dict(cfg_u)
+        import_vars(dictionary=cfg_u)
+
+    #number of gridpoints (nx*ny) depends on grid generation method
+    if GRID_GEN_METHOD == "GFDLgrid":
+      (\
+      LON_CTR,LAT_CTR,NX,NY,NHW,STRETCH_FAC,
+      ISTART_OF_RGNL_DOM_WITH_WIDE_HALO_ON_T6SG,
+      IEND_OF_RGNL_DOM_WITH_WIDE_HALO_ON_T6SG,
+      JSTART_OF_RGNL_DOM_WITH_WIDE_HALO_ON_T6SG,
+      JEND_OF_RGNL_DOM_WITH_WIDE_HALO_ON_T6SG \
+      ) = \
+        set_gridparams_GFDLgrid( \
+        lon_of_t6_ctr=GFDLgrid_LON_T6_CTR, \
+        lat_of_t6_ctr=GFDLgrid_LAT_T6_CTR, \
+        res_of_t6g=GFDLgrid_RES, \
+        stretch_factor=GFDLgrid_STRETCH_FAC, \
+        refine_ratio_t6g_to_t7g=GFDLgrid_REFINE_RATIO, \
+        istart_of_t7_on_t6g=GFDLgrid_ISTART_OF_RGNL_DOM_ON_T6G, \
+        iend_of_t7_on_t6g=GFDLgrid_IEND_OF_RGNL_DOM_ON_T6G, \
+        jstart_of_t7_on_t6g=GFDLgrid_JSTART_OF_RGNL_DOM_ON_T6G, \
+        jend_of_t7_on_t6g=GFDLgrid_JEND_OF_RGNL_DOM_ON_T6G)
+
+    elif GRID_GEN_METHOD == "ESGgrid":
+      (\
+      LON_CTR,LAT_CTR,NX,NY,PAZI,
+      NHW,STRETCH_FAC,DEL_ANGLE_X_SG,DEL_ANGLE_Y_SG,
+      NEG_NX_OF_DOM_WITH_WIDE_HALO,
+      NEG_NY_OF_DOM_WITH_WIDE_HALO \
+      ) = \
+        set_gridparams_ESGgrid( \
+        lon_ctr=ESGgrid_LON_CTR, \
+        lat_ctr=ESGgrid_LAT_CTR, \
+        nx=ESGgrid_NX, \
+        ny=ESGgrid_NY, \
+        pazi=ESGgrid_PAZI, \
+        halo_width=ESGgrid_WIDE_HALO_WIDTH, \
+        delx=ESGgrid_DELX, \
+        dely=ESGgrid_DELY)
+
+    cost = [DT_ATMOS, NX*NY]
+
+    #reference grid (6-hour forecast on RRFS_CONUS_25km)
+    PREDEF_GRID_NAME="RRFS_CONUS_25km"
+
+    export_vars()
+    set_predef_grid_params()
+    import_vars()
+    cost.extend([DT_ATMOS, ESGgrid_NX*ESGgrid_NY])
+
+    return cost
+
+#interface
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description=\
+             'Calculates parameters needed for calculating cost.')
+    parser.add_argument('--cfg','-c',dest='cfg',required=True,
+             help='config file containing grip params')
+    args = parser.parse_args()
+
+    params = calculate_cost(args.cfg)
+    print(' '.join(map(str,params)))
+
+class Testing(unittest.TestCase):
+    def test_calculate_cost(self):
+        USHDIR = os.path.dirname(os.path.abspath(__file__))
+        config_fn = os.path.join(USHDIR, "config.community.sh")
+        params = calculate_cost(config_fn)
+        self.assertCountEqual(params, [36, 1987440, 36, 28689])
+
+    def setUp(self):
+        set_env_var('DEBUG',False)
+        set_env_var('PREDEF_GRID_NAME',"RRFS_CONUS_3km")
+        set_env_var('DT_ATMOS',36)
+        set_env_var('LAYOUT_X',18)
+        set_env_var('LAYOUT_Y',36)
+        set_env_var('BLOCKSIZE',28)
+        set_env_var('QUILTING',False)
+        set_env_var('RADIUS_EARTH',6371200.0)
+        set_env_var('DEGS_PER_RADIAN',57.2957795131)
+

ush/constants.sh

@@ -7,13 +7,13 @@
 #
 
 # Pi.
-pi_geom="3.14159265358979323846264338327"
+PI_GEOM="3.14159265358979323846264338327"
 
 # Degrees per radian.
-degs_per_radian=$( bc -l <<< "360.0/(2.0*${pi_geom})" )
+DEGS_PER_RADIAN=$( bc -l <<< "360.0/(2.0*${PI_GEOM})" )
 
 # Radius of the Earth in meters.
-radius_Earth="6371200.0"
+RADIUS_EARTH="6371200.0"
 #
 #-----------------------------------------------------------------------
 #

ush/get_crontab_contents.py

@@ -8,7 +8,6 @@
 
 from python_utils import import_vars, set_env_var, print_input_args, \
                          run_command, define_macos_utilities, print_info_msg
-from constants import valid_vals_BOOLEAN
 
 def get_crontab_contents(called_from_cron):
     """

ush/python_utils/__init__.py

@@ -1,5 +1,5 @@
 try:
-  from .misc import uppercase, lowercase, find_pattern_in_str, find_pattern_in_file
+  from .misc import uppercase, lowercase, find_pattern_in_str, find_pattern_in_file, flatten_dict
   from .check_for_preexist_dir_file import check_for_preexist_dir_file
   from .check_var_valid_value import check_var_valid_value
   from .count_files import count_files

ush/python_utils/misc.py

@@ -55,3 +55,21 @@ def find_pattern_in_file(pattern, file_name):
                 return match.groups()
     return None
 
+def flatten_dict(dictionary,keys=None):
+    """ Faltten a recursive dictionary (e.g.yaml/json) to be one level deep
+    Args:
+      dictionary: the source dictionary
+      keys: list of keys on top level whose contents to flatten, if None all of them
+    Returns:
+      A one-level deep dictionary for the selected set of keys
+    """
+    flat_dict = {}
+    for k,v in dictionary.items():
+        if not keys or k in keys:
+            if isinstance(v,dict):
+                r = flatten_dict(v)
+                flat_dict.update(r)
+            else:
+                flat_dict[k] = v
+    return flat_dict
+