Generalize single cycle script to also work for hofx

parm/atm/hofx/hofx_nomodel.yaml

@@ -0,0 +1,16 @@
+window begin: '$(ATM_WINDOW_BEGIN)'
+window length: $(ATM_WINDOW_LENGTH)
+geometry: $(GEOM_BKG)
+state:
+  datapath: $(BKG_DIR)
+  filetype: fms restart
+  datetime: $(BKG_ISOTIME)
+  filename_core: $(BKG_YYYYmmddHHMMSS).fv_core.res.nc
+  filename_trcr: $(BKG_YYYYmmddHHMMSS).fv_tracer.res.nc
+  filename_sfcd: $(BKG_YYYYmmddHHMMSS).sfc_data.nc
+  filename_sfcw: $(BKG_YYYYmmddHHMMSS).fv_srf_wnd.res.nc
+  filename_cplr: $(BKG_YYYYmmddHHMMSS).coupler.res
+  state variables: [ua,va,t,delp,sphum,ice_wat,liq_wat,o3mr,phis,
+                    slmsk,sheleg,tsea,vtype,stype,vfrac,stc,smc,snwdph,
+                    u_srf,v_srf,f10m]
+observations: $<< $(OBS_LIST)

parm/atm/variational/3dvar_dripcg.yaml

@@ -8,11 +8,11 @@ cost function:
     datapath: $(BKG_DIR)
     filetype: fms restart
     datetime: $(BKG_ISOTIME)
-    filename_core: $(BKG_YYYYMMDDpHHMMSS).fv_core.res.nc
-    filename_trcr: $(BKG_YYYYMMDDpHHMMSS).fv_tracer.res.nc
-    filename_sfcd: $(BKG_YYYYMMDDpHHMMSS).sfc_data.nc
-    filename_sfcw: $(BKG_YYYYMMDDpHHMMSS).fv_srf_wnd.res.nc
-    filename_cplr: $(BKG_YYYYMMDDpHHMMSS).coupler.res
+    filename_core: $(BKG_YYYYmmddHHMMSS).fv_core.res.nc
+    filename_trcr: $(BKG_YYYYmmddHHMMSS).fv_tracer.res.nc
+    filename_sfcd: $(BKG_YYYYmmddHHMMSS).sfc_data.nc
+    filename_sfcw: $(BKG_YYYYmmddHHMMSS).fv_srf_wnd.res.nc
+    filename_cplr: $(BKG_YYYYmmddHHMMSS).coupler.res
     state variables: [ua,va,t,delp,sphum,ice_wat,liq_wat,o3mr,phis,
                       slmsk,sheleg,tsea,vtype,stype,vfrac,stc,smc,snwdph,
                       u_srf,v_srf,f10m]

ush/run_single_atm_var_analysis.py → ush/run_jedi_exe.py

@@ -9,7 +9,7 @@
 import yaml
 
 
-def run_atm_var_analysis(yamlconfig):
+def run_jedi_exe(yamlconfig):
     logging.basicConfig(format='%(asctime)s:%(levelname)s:%(message)s', level=logging.INFO, datefmt='%Y-%m-%d %H:%M:%S')
     # open YAML file to get config
     try:
@@ -18,6 +18,13 @@ def run_atm_var_analysis(yamlconfig):
         logging.info(f'Loading configuration from {yamlconfig}')
     except Exception as e:
         logging.error(f'Error occurred when attempting to load: {yamlconfig}, error: {e}')
+    # check if the specified app mode is valid
+    app_mode = all_config_dict['GDASApp mode']
+    supported_app_modes = ['hofx', 'variational']
+    if app_mode not in supported_app_modes:
+        raise KeyError(f"'{app_mode}' not supported. " +
+                       "Current GDASApp modes supported are: " +
+                       f"{' | '.join(supported_app_modes)}")
     # create working directory
     workdir = all_config_dict['working directory']
     try:
@@ -30,67 +37,68 @@ def run_atm_var_analysis(yamlconfig):
     sys.path.append(ufsda_path)
     import ufsda
     # compute config for YAML for executable
-    analysis_subconfig = all_config_dict['analysis options']
-    valid_time = analysis_subconfig['valid_time']
-    h = re.findall('PT(\\d+)H', analysis_subconfig['atm_window_length'])[0]
+    executable_subconfig = all_config_dict['executable options']
+    valid_time = executable_subconfig['valid_time']
+    h = re.findall('PT(\\d+)H', executable_subconfig['atm_window_length'])[0]
     prev_cycle = valid_time - dt.timedelta(hours=int(h))
     window_begin = valid_time - dt.timedelta(hours=int(h)/2)
     cyc = valid_time.strftime("%H")
     cdate = valid_time.strftime("%Y%m%d%H")
     gcyc = prev_cycle.strftime("%H")
     gdate = prev_cycle.strftime("%Y%m%d%H")
     var_config = {
-        'BERROR_YAML': analysis_subconfig['berror_yaml'],
-        'OBS_YAML_DIR': analysis_subconfig['obs_yaml_dir'],
-        'OBS_LIST': analysis_subconfig['obs_list'],
-        'atm': analysis_subconfig['atm'],
-        'layout_x': str(analysis_subconfig['layout_x']),
-        'layout_y': str(analysis_subconfig['layout_y']),
+        'BERROR_YAML': executable_subconfig.get('berror_yaml', './'),
+        'OBS_YAML_DIR': executable_subconfig['obs_yaml_dir'],
+        'OBS_LIST': executable_subconfig['obs_list'],
+        'atm': executable_subconfig.get('atm', False),
+        'layout_x': str(executable_subconfig['layout_x']),
+        'layout_y': str(executable_subconfig['layout_y']),
         'BKG_DIR': os.path.join(workdir, 'bkg'),
         'fv3jedi_fix_dir': os.path.join(workdir, 'Data', 'fv3files'),
         'fv3jedi_fieldset_dir': os.path.join(workdir, 'Data', 'fieldsets'),
         'ANL_DIR': os.path.join(workdir, 'anl'),
         'fv3jedi_staticb_dir': os.path.join(workdir, 'berror'),
         'BIAS_DIR': os.path.join(workdir, 'obs'),
         'CRTM_COEFF_DIR': os.path.join(workdir, 'crtm'),
-        'BIAS_PREFIX': f"{analysis_subconfig['dump']}.t{gcyc}z",
+        'BIAS_PREFIX': f"{executable_subconfig['dump']}.t{gcyc}z",
         'BIAS_DATE': f"{gdate}",
         'DIAG_DIR': os.path.join(workdir, 'diags'),
         'OBS_DIR': os.path.join(workdir, 'obs'),
-        'OBS_PREFIX': f"{analysis_subconfig['dump']}.t{cyc}z",
+        'OBS_PREFIX': f"{executable_subconfig['dump']}.t{cyc}z",
         'OBS_DATE': f"{cdate}",
         'valid_time': f"{valid_time.strftime('%Y-%m-%dT%H:%M:%SZ')}",
         'window_begin': f"{window_begin.strftime('%Y-%m-%dT%H:%M:%SZ')}",
         'prev_valid_time': f"{prev_cycle.strftime('%Y-%m-%dT%H:%M:%SZ')}",
-        'atm_window_length': analysis_subconfig['atm_window_length'],
-        'CASE': analysis_subconfig['case'],
-        'CASE_ENKF': analysis_subconfig['case_enkf'],
-        'DOHYBVAR': analysis_subconfig['dohybvar'],
-        'LEVS': str(analysis_subconfig['levs']),
+        'atm_window_length': executable_subconfig['atm_window_length'],
+        'CASE': executable_subconfig['case'],
+        'CASE_ENKF': executable_subconfig.get('case_enkf', executable_subconfig['case']),
+        'DOHYBVAR': executable_subconfig.get('dohybvar', False),
+        'LEVS': str(executable_subconfig['levs']),
     }
-    template = analysis_subconfig['var_yaml']
-    output_file = os.path.join(workdir, 'fv3jedi_var.yaml')
+    template = executable_subconfig['yaml_template']
+    output_file = os.path.join(workdir, f"gdas_{app_mode}.yaml")
     # set some environment variables
     os.environ['PARMgfs'] = os.path.join(all_config_dict['GDASApp home'], 'parm')
     # generate YAML for executable based on input config
     logging.info(f'Using YAML template {template}')
     ufsda.yamltools.genYAML(var_config, template=template, output=output_file)
-    logging.info(f'Wrote Variational DA YAML file to {output_file}')
+    logging.info(f'Wrote YAML file to {output_file}')
     # use R2D2 to stage backgrounds, obs, bias correction files, etc.
     ufsda.stage.gdas_single_cycle(var_config)
     # link additional fix files needed (CRTM, fieldsets, etc.)
-    gdasfix = analysis_subconfig['gdas_fix_root']
+    gdasfix = executable_subconfig['gdas_fix_root']
     ufsda.stage.gdas_fix(gdasfix, workdir, var_config)
     # link executable
-    varexe = os.path.join(all_config_dict['GDASApp home'], 'build', 'bin', 'fv3jedi_var.x')
-    ufsda.disk_utils.symlink(varexe, os.path.join(workdir, 'fv3jedi_var.x'))
+    baseexe = os.path.basename(executable_subconfig['executable'])
+    ufsda.disk_utils.symlink(executable_subconfig['executable'], os.path.join(workdir, baseexe))
     # create output directories
     ufsda.disk_utils.mkdir(os.path.join(workdir, 'diags'))
-    ufsda.disk_utils.mkdir(os.path.join(workdir, 'anl'))
+    if app_mode in ['variational']:
+        ufsda.disk_utils.mkdir(os.path.join(workdir, 'anl'))
     # generate job submission script
     job_script = ufsda.misc_utils.create_batch_job(all_config_dict['job options'],
                                                    workdir,
-                                                   os.path.join(workdir, 'fv3jedi_var.x'),
+                                                   os.path.join(workdir, baseexe),
                                                    output_file)
     # submit job to queue
     ufsda.misc_utils.submit_batch_job(all_config_dict['job options'], workdir, job_script)
@@ -100,4 +108,4 @@ def run_atm_var_analysis(yamlconfig):
     parser = argparse.ArgumentParser()
     parser.add_argument('-c', '--config', type=str, help='Input YAML Configuration', required=True)
     args = parser.parse_args()
-    run_atm_var_analysis(args.config)
+    run_jedi_exe(args.config)

ush/ufsda/misc_utils.py

@@ -23,7 +23,7 @@ def isTrue(str_in):
     return status
 
 
-def create_batch_job(job_config, working_dir, exe_path, yaml_path):
+def create_batch_job(job_config, working_dir, executable, yaml_path):
     """
     create a batch job submission shell script
     """
@@ -32,20 +32,25 @@ def create_batch_job(job_config, working_dir, exe_path, yaml_path):
     with open(batch_script, 'w') as f:
         f.write('#!/bin/bash\n')
         if scheduler[job_config['machine']] == 'slurm':
-            f.write('#SBATCH -J GDASApp\n')
-            f.write('#SBATCH -o GDASApp.o%J\n')
-            f.write(f"#SBATCH -A {job_config['account']}\n")
-            f.write(f"#SBATCH -q {job_config['queue']}\n")
-            f.write(f"#SBATCH -p {job_config['partition']}\n")
-            f.write(f"#SBATCH --ntasks={job_config['ntasks']}\n")
-            f.write(f"#SBATCH --cpus-per-task={job_config['cpus-per-task']}\n")
-            f.write(f"#SBATCH --exclusive\n")
-            f.write(f"#SBATCH -t {job_config['walltime']}\n")
-        f.write(f"module use {job_config['modulepath']}\n")
-        f.write(f"module load GDAS/{job_config['machine']}\n")
-        f.write(f"cd {working_dir}\n")
+            sbatch = f"""#SBATCH -J GDASApp
+#SBATCH -o GDASApp.o%J
+#SBATCH -A {job_config['account']}
+#SBATCH -q {job_config['queue']}
+#SBATCH -p {job_config['partition']}
+#SBATCH --ntasks={job_config['ntasks']}
+#SBATCH --cpus-per-task={job_config['cpus-per-task']}
+#SBATCH --exclusive
+#SBATCH -t {job_config['walltime']}"""
+            f.write(sbatch)
+        commands = f"""
+module purge
+module use {job_config['modulepath']}
+module load GDAS/{job_config['machine']}
+cd {working_dir}
+"""
+        f.write(commands)
         if scheduler[job_config['machine']] == 'slurm':
-            f.write(f"srun -n $SLURM_NTASKS {exe_path} {yaml_path}\n")
+            f.write(f"srun -n $SLURM_NTASKS {executable} {yaml_path}\n")
     logging.info(f"Wrote batch submission script to {batch_script}")
     return batch_script
 

ush/ufsda/yamltools.py

@@ -48,7 +48,9 @@ def parse_config(input_config_dict, template=None, clean=True):
     if config_out.get('atm', True):
         config_out = atmanl_case(config_out)
     config_out.pop('atm', None)  # pop out boolean variable that will cause issues later
-    # do a first round of includes first
+    # do a first round of substitutions first
+    config_out = replace_vars(config_out)
+    # now do a first round of includes
     config_out = include_yaml(config_out)
     # pull common key values out to top layer
     config_out = pop_out_common(config_out)
@@ -123,10 +125,10 @@ def calc_time_vars(config):
                                                     "%Y-%m-%dT%H:%M:%SZ")
     else:
         raise KeyError("Neither $(valid_time) nor ${PDY}${cyc} defined")
-    config['YYYYMMDDpHHMMSS'] = valid_time_obj.strftime('%Y%m%d.%H%M%S')
+    config['YYYYmmddHHMMSS'] = valid_time_obj.strftime('%Y%m%d.%H%M%S')
     # for now bkg_time == valid_time, will change for fgat
     bkg_time_obj = valid_time_obj
-    config['BKG_YYYYMMDDpHHMMSS'] = bkg_time_obj.strftime('%Y%m%d.%H%M%S')
+    config['BKG_YYYYmmddHHMMSS'] = bkg_time_obj.strftime('%Y%m%d.%H%M%S')
     config['BKG_ISOTIME'] = bkg_time_obj.strftime('%Y-%m-%dT%H:%M:%SZ')
     if 'assim_freq' in os.environ:
         config['ATM_WINDOW_LENGTH'] = f"PT{os.environ['assim_freq']}H"