Stand alone script to run a single atm 3DVar analysis

modulefiles/GDAS/orion.lua

@@ -53,6 +53,9 @@ local mpinproc = '-n'
 setenv('MPIEXEC_EXEC', mpiexec)
 setenv('MPIEXEC_NPROC', mpinproc)
 
+setenv('R2D2_CONFIG', '/work2/noaa/da/cmartin/GDASApp/R2D2_SHARED/config_orion.yaml')
+
+execute{cmd="ulimit -s unlimited",modeA={"load"}}
 
 whatis("Name: ".. pkgName)
 whatis("Version: " .. pkgVersion)

parm/atm/berror/staticb_bump.yaml

@@ -20,11 +20,11 @@ saber blocks:
       variables: [surface_pressure]
     universe radius:
       filetype: fms restart
-      datetime: 2021-12-26T00:00:00Z
+      datetime: 2021-12-22T00:00:00Z
       set datetime on read: true
       psinfile: true
       datapath: *staticb_dir
-      prefix: cor/20211226.000000
+      prefix: cor/20211222.000000
       filename_core: cor_rh.fv_core.res.nc
       filename_trcr: cor_rh.fv_tracer.res.nc
       filename_sfcd: cor_rh.sfc_data.nc
@@ -36,11 +36,11 @@ saber blocks:
   active variables: *active_vars
   file:
     filetype: fms restart
-    datetime: 2021-12-26T00:00:00Z
+    datetime: 2021-12-22T00:00:00Z
     set datetime on read: true
     psinfile: true
     datapath: *staticb_dir
-    prefix: stddev/20211226.000000
+    prefix: stddev/20211222.000000
     filename_core: stddev.fv_core.res.nc
     filename_trcr: stddev.fv_tracer.res.nc
     filename_sfcd: stddev.sfc_data.nc
@@ -57,7 +57,7 @@ saber blocks:
     universe_rad: 2000.0e3
     load_vbal: true
     load_samp_local: true
-    fname_samp: vbal/vbal_2021122600_gfs_sampling
+    fname_samp: vbal/vbal_2021122200_gfs_sampling
     vbal_block: [true, true,false, true,false,false]
 - saber block name: BUMP_PsiChiToUV
   input variables: *control_vars

parm/atm/berror/staticb_identity.yaml

@@ -0,0 +1,2 @@
+covariance model: FV3JEDI-ID
+date: $(window_begin)

parm/atm/obs/config/amsua_n19.yaml

@@ -14,7 +14,7 @@ obs operator:
   obs options:
     Sensor_ID: amsua_n19
     EndianType: little_endian
-    CoefficientPath: $(CRTM_COEFF_DIR)
+    CoefficientPath: $(CRTM_COEFF_DIR)/
 obs bias:
   input file: $(BIAS_DIR)/$(BIAS_PREFIX).amsua_n19.satbias.$(BIAS_DATE).nc4
   output file: $(BIAS_DIR)/$(OBS_PREFIX).amsua_n19.satbias.$(OBS_DATE).nc4
@@ -44,7 +44,7 @@ obs bias:
       inflation:
         ratio: 1.1
         ratio for small dataset: 2.0
-    output file: $(BIAS_DIR)/$(OBS_PREFIX).amsua_n19.satbias.$(OBS_DATE).nc4
+    output file: $(BIAS_DIR)/$(OBS_PREFIX).amsua_n19.satbias_cov.$(OBS_DATE).nc4
 obs filters:
 - filter: BlackList
   filter variables:

parm/atm/obs/config/sondes.yaml

@@ -3,7 +3,7 @@ obs space:
   obsdatain:
     obsfile: $(OBS_DIR)/$(OBS_PREFIX).sondes.$(OBS_DATE).nc4
     obsgrouping:
-      group variables: ["station_id", "LaunchTime"]
+      group variables: ["station_id"]
       sort variable: "air_pressure"
       sort order: "descending"
   obsdataout:

test/genYAML_prep.sh

@@ -6,7 +6,8 @@ cat > testoutput/genYAML_example.yaml << EOF
 template: ${srcdir}/parm/atm/variational/3dvar_dripcg.yaml
 output: ${bindir}/test/testoutput/genYAML_output_3dvar.yaml
 config:
-  paths: $<< ${srcdir}/parm/atm/common/paths.yaml
+  BERROR_YAML: \${PARMgfs}/atm/berror/staticb_bump.yaml
+  OBS_YAML_DIR: \${PARMgfs}/atm/obs/config
   atm: true
   layout_x: '1'
   layout_y: '2'

test/test_generate_yaml.py

@@ -10,7 +10,8 @@ def test_yaml_gen_yaml(parm_dir):
     output_file = os.path.join(os.getcwd(), 'testoutput', 'test_yaml_gen.yaml')
     template = os.path.join(parm_dir, 'atm/variational/3dvar_dripcg.yaml')
     config_dict = {
-        'paths': f'$<< {parm_dir}/atm/common/paths.yaml',
+        'BERROR_YAML': '${PARMgfs}/atm/berror/staticb_bump.yaml',
+        'OBS_YAML_DIR': '${PARMgfs}/atm/obs/config',
         'atm': True,
         'layout_x': '1',
         'layout_y': '2',

ush/run_single_atm_var_analysis.py

@@ -0,0 +1,103 @@
+#!/usr/bin/env python3
+import argparse
+import datetime as dt
+import logging
+import os
+import re
+import subprocess
+import sys
+import yaml
+
+
+def run_atm_var_analysis(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:
+        with open(yamlconfig, 'r') as yamlconfig_opened:
+            all_config_dict = yaml.safe_load(yamlconfig_opened)
+        logging.info(f'Loading configuration from {yamlconfig}')
+    except Exception as e:
+        logging.error(f'Error occurred when attempting to load: {yamlconfig}, error: {e}')
+    # create working directory
+    workdir = all_config_dict['working directory']
+    try:
+        os.makedirs(workdir, exist_ok=True)
+        logging.info(f'Created working directory {workdir}')
+    except OSError as error:
+        logging.error(f'Error creating {workdir}: {error}')
+    # add your ufsda python package to path
+    ufsda_path = os.path.join(all_config_dict['GDASApp home'], 'ush')
+    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]
+    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']),
+        '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_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_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']),
+    }
+    template = analysis_subconfig['var_yaml']
+    output_file = os.path.join(workdir, 'fv3jedi_var.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}')
+    # 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']
+    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'))
+    # create output directories
+    ufsda.disk_utils.mkdir(os.path.join(workdir, 'diags'))
+    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'),
+                                                   output_file)
+    # submit job to queue
+    ufsda.misc_utils.submit_batch_job(all_config_dict['job options'], workdir, job_script)
+
+
+if __name__ == "__main__":
+    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)

ush/ufsda/__init__.py

@@ -1,7 +1,7 @@
-from .disk_utils import *
+from .disk_utils import mkdir, symlink
 from .ufs_yaml import gen_yaml, parse_config
 import ufsda.stage
 import ufsda.r2d2
 import ufsda.post
 import ufsda.yamltools
-from .misc_utils import isTrue
+from .misc_utils import isTrue, create_batch_job, submit_batch_job

ush/ufsda/disk_utils.py

@@ -1,3 +1,7 @@
+import os
+import logging
+
+
 def mkdir(dirpath):
     """
     mkdir(dirpath)
@@ -7,6 +11,17 @@ def mkdir(dirpath):
     import os
     try:
         os.makedirs(dirpath, exist_ok=True)
-        print(f"{dirpath} created successfully")
+        logging.info(f"{dirpath} created successfully")
     except OSError as error:
-        print(f"{dirpath} could not be created")
+        logging.info(f"{dirpath} could not be created")
+
+
+def symlink(src, dest):
+    try:
+        os.symlink(src, dest)
+        logging.info(f"Symbolically linked {src} to {dest}")
+    except FileExistsError:
+        os.remove(dest)
+        logging.info(f"{dest} exists, removing...")
+        os.symlink(src, dest)
+        logging.info(f"Symbolically linked {src} to {dest}")

ush/ufsda/misc_utils.py

@@ -1,3 +1,8 @@
+import logging
+import os
+import subprocess
+
+
 def isTrue(str_in):
     """ isTrue(str_in)
     - function to translate shell variables to python logical variables
@@ -10,3 +15,38 @@ def isTrue(str_in):
     else:
         status = False
     return status
+
+
+def create_batch_job(job_config, working_dir, exe_path, yaml_path):
+    """
+    create a batch job submission shell script
+    """
+    # open up a file for writing
+    batch_script = os.path.join(working_dir, 'submit_job.sh')
+    with open(batch_script, 'w') as f:
+        f.write('#!/bin/bash\n')
+        if job_config['machine'] in ['orion', 'hera']:
+            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")
+        if job_config['machine'] in ['orion', 'hera']:
+            f.write(f"srun -n $SLURM_NTASKS {exe_path} {yaml_path}\n")
+    logging.info(f"Wrote batch submission script to {batch_script}")
+    return batch_script
+
+
+def submit_batch_job(job_config, working_dir, job_script):
+    """
+    submit a batch job
+    """
+    if job_config['machine'] in ['orion', 'hera']:
+        subprocess.Popen(f"sbatch {job_script}", cwd=working_dir, shell=True)

ush/ufsda/r2d2.py

@@ -19,7 +19,7 @@ def store(config):
     dump = config.get('dump', 'gdas')
     source_dir = config['source_dir']
     source_file_fmt = config['source_file_fmt']
-    obs_types = config['obs_types']
+    obs_types = config.get('obs_types', None)
 
     for time in times:
         year = Hour(time).format('%Y')
@@ -39,3 +39,35 @@ def store(config):
             kwargs['step'] = config['forecast_steps']
             kwargs['source_file'] = eval(f"f'{source_file_fmt}'"),
             r2d2.store(**kwargs)
+
+
+def fetch(config):
+    kwargs = {}
+    kwargs['ignore_missing'] = False
+    for arg in config.keys():
+        if arg in possible_args:
+            kwargs[arg] = config[arg]
+    type = config.type
+    times = date_sequence(config.start, config.end, config.step)
+    dump = config.get('dump', 'gdas')
+    obs_types = config.get('obs_types', None)
+    target_dir = config['target_dir']
+    target_file_fmt = config['target_file_fmt']
+    for time in times:
+        year = Hour(time).format('%Y')
+        month = Hour(time).format('%m')
+        day = Hour(time).format('%d')
+        hour = Hour(time).format('%H')
+        kwargs['date'] = time
+        if type in ['bc', 'ob']:
+            if type == 'ob':
+                kwargs['time_window'] = config['step']
+            for obs_type in obs_types:
+                kwargs['target_file'] = eval(f"f'{target_file_fmt}'"),
+                kwargs['obs_type'] = obs_type
+                r2d2.fetch(**kwargs)
+        else:
+            kwargs['file_type'] = config.file_type_list
+            kwargs['step'] = config['forecast_steps']
+            kwargs['target_file'] = eval(f"f'{target_file_fmt}'"),
+            r2d2.fetch(**kwargs)