Stage variational and ensemble DA job files with Jinja2-templated YAMLs

parm/config/gfs/config.atmanl

@@ -23,6 +23,9 @@ fi
 
 export CRTM_FIX_YAML="${PARMgfs}/gdas/atm_crtm_coeff.yaml.j2"
 export JEDI_FIX_YAML="${PARMgfs}/gdas/atm_jedi_fix.yaml.j2"
+export VAR_BKG_STAGING_YAML="${PARMgfs}/gdas/staging/atm_var_bkg.yaml.j2"
+export BERROR_STAGING_YAML="${PARMgfs}/gdas/staging/atm_berror_${STATICB_TYPE}.yaml.j2"
+export FV3ENS_STAGING_YAML="${PARMgfs}/gdas/staging/atm_var_fv3ens.yaml.j2"
 
 export layout_x_atmanl=@LAYOUT_X_ATMANL@
 export layout_y_atmanl=@LAYOUT_Y_ATMANL@

parm/config/gfs/config.atmensanl

@@ -12,6 +12,7 @@ export INTERP_METHOD='barycentric'
 
 export CRTM_FIX_YAML="${PARMgfs}/gdas/atm_crtm_coeff.yaml.j2"
 export JEDI_FIX_YAML="${PARMgfs}/gdas/atm_jedi_fix.yaml.j2"
+export LGETKF_BKG_STAGING_YAML="${PARMgfs}/gdas/staging/atm_lgetkf_bkg.yaml.j2"
 
 export layout_x_atmensanl=@LAYOUT_X_ATMENSANL@
 export layout_y_atmensanl=@LAYOUT_Y_ATMENSANL@

parm/gdas/staging/atm_berror_gsibec.yaml.j2

@@ -0,0 +1,8 @@
+{% set fname_list = ['gfs_gsi_global.nml', 'gsi-coeffs-gfs-global.nc4'] %}
+
+mkdir:
+- '{{ DATA }}/berror'
+copy:
+{% for fname in fname_list %}
+- ['{{ HOMEgfs }}/fix/gdas/gsibec/{{ CASE_ANL }}/{{ fname }}', '{{ DATA }}/berror']
+{% endfor %}

parm/gdas/staging/atm_lgetkf_bkg.yaml.j2

@@ -0,0 +1,32 @@
+{% set ftype_list = ['fv_core.res', 'fv_srf_wnd.res', 'fv_tracer.res', 'phy_data', 'sfc_data'] %}
+{% set time_list = [current_cycle] %}
+
+mkdir:
+{% for imem in range(1,NMEM_ENS+1) %}
+   {% set memchar = 'mem%03d' | format(imem) %}
+   {% set tmpl_dict = ({ '${ROTDIR}': ROTDIR,
+                         '${RUN}': RUN, 
+                         '${YMD}': current_cycle | to_YMD,	
+                         '${HH}': current_cycle | strftime('%H'),
+                         '${MEMDIR}': memchar }) %}
+- '{{ DATA }}/bkg/{{ memchar }}'
+- '{{ DATA }}/anl/{{ memchar }}'
+- '{{ COM_ATMOS_ANALYSIS_TMPL | replace_tmpl(tmpl_dict) }}'
+{% endfor %}
+copy:
+{% for time in time_list %}
+   {% for imem in range(1,NMEM_ENS+1) %}
+      {% set memchar = 'mem%03d' | format(imem) %}
+      {% set tmpl_dict = ({ '${ROTDIR}': ROTDIR,
+      	      	      	    '${RUN}': 'enkfgdas',
+                            '${YMD}': previous_cycle | to_YMD,	
+                            '${HH}': previous_cycle | strftime('%H'),
+                            '${MEMDIR}': memchar }) %}
+- ['{{ COM_ATMOS_RESTART_TMPL | replace_tmpl(tmpl_dict) }}/{{ time | to_fv3time }}.coupler.res', '{{ DATA }}/bkg/{{ memchar }}/']
+      {% for ftype in ftype_list %}
+         {% for itile in range(1,7) %}
+- ['{{ COM_ATMOS_RESTART_TMPL | replace_tmpl(tmpl_dict) }}/{{ time | to_fv3time }}.{{ ftype }}.tile{{ itile }}.nc', '{{ DATA }}/bkg/{{ memchar }}/']
+         {% endfor %}
+      {% endfor %}
+   {% endfor %}
+{% endfor %}

parm/gdas/staging/atm_var_bkg.yaml.j2

@@ -0,0 +1,14 @@
+{% set ftype_list = ['fv_core.res', 'fv_srf_wnd.res', 'fv_tracer.res', 'phy_data', 'sfc_data'] %}
+{% set time_list = [current_cycle] %}
+
+mkdir:
+- '{{ DATA }}/bkg'
+copy:
+{% for time in time_list %}
+- ['{{ COM_ATMOS_RESTART_PREV }}/{{ time | to_fv3time }}.coupler.res', '{{ DATA }}/bkg/']
+   {% for ftype in ftype_list %}
+      {% for itile in range(1,ntiles+1) %}
+- ['{{ COM_ATMOS_RESTART_PREV }}/{{ time | to_fv3time }}.{{ ftype }}.tile{{ itile }}.nc', '{{ DATA }}/bkg/']
+      {% endfor %}
+   {% endfor %}
+{% endfor %}

parm/gdas/staging/atm_var_fv3ens.yaml.j2

@@ -0,0 +1,24 @@
+{% set ftype_list = ['fv_core.res', 'fv_srf_wnd.res', 'fv_tracer.res', 'phy_data', 'sfc_data'] %}
+{% set time_list = [current_cycle] %}
+
+mkdir:
+{% for imem in range(1,NMEM_ENS+1) %}
+- '{{ DATA }}/ens/{{ 'mem%03d' | format(imem) }}'
+{% endfor %}
+copy:
+{% for time in time_list %}
+   {% for imem in range(1,NMEM_ENS+1) %}
+      {% set memchar = 'mem%03d' | format(imem) %}
+      {% set tmpl_dict = ({ '${ROTDIR}': ROTDIR,
+                            '${RUN}': 'enkfgdas',
+                            '${YMD}': previous_cycle | to_YMD,
+                            '${HH}': previous_cycle | strftime('%H'),
+                            '${MEMDIR}': memchar }) %}
+- ['{{ COM_ATMOS_RESTART_TMPL | replace_tmpl(tmpl_dict) }}/{{ time | to_fv3time }}.coupler.res', '{{ DATA }}/ens/{{ memchar }}/']
+      {% for ftype in ftype_list %}
+         {% for itile in range(1,ntiles+1) %}
+- ['{{ COM_ATMOS_RESTART_TMPL | replace_tmpl(tmpl_dict) }}/{{ time | to_fv3time }}.{{ ftype }}.tile{{ itile }}.nc', '{{ DATA }}/ens/{{ memchar }}/']
+         {% endfor %}
+      {% endfor %}
+   {% endfor %}
+{% endfor %}

ush/python/pygfs/task/aero_analysis.py

@@ -29,33 +29,33 @@ class AerosolAnalysis(Analysis):
     def __init__(self, config):
         super().__init__(config)
 
-        _res = int(self.config['CASE'][1:])
-        _res_anl = int(self.config['CASE_ANL'][1:])
-        _window_begin = add_to_datetime(self.runtime_config.current_cycle, -to_timedelta(f"{self.config['assim_freq']}H") / 2)
-        _jedi_yaml = os.path.join(self.runtime_config.DATA, f"{self.runtime_config.CDUMP}.t{self.runtime_config['cyc']:02d}z.aerovar.yaml")
+        _res = int(self.task_config['CASE'][1:])
+        _res_anl = int(self.task_config['CASE_ANL'][1:])
+        _window_begin = add_to_datetime(self.task_config.current_cycle, -to_timedelta(f"{self.task_config['assim_freq']}H") / 2)
+        _jedi_yaml = os.path.join(self.task_config.DATA, f"{self.task_config.RUN}.t{self.task_config['cyc']:02d}z.aerovar.yaml")
 
         # Create a local dictionary that is repeatedly used across this class
         local_dict = AttrDict(
             {
                 'npx_ges': _res + 1,
                 'npy_ges': _res + 1,
-                'npz_ges': self.config.LEVS - 1,
-                'npz': self.config.LEVS - 1,
+                'npz_ges': self.task_config.LEVS - 1,
+                'npz': self.task_config.LEVS - 1,
                 'npx_anl': _res_anl + 1,
                 'npy_anl': _res_anl + 1,
-                'npz_anl': self.config['LEVS'] - 1,
+                'npz_anl': self.task_config['LEVS'] - 1,
                 'AERO_WINDOW_BEGIN': _window_begin,
-                'AERO_WINDOW_LENGTH': f"PT{self.config['assim_freq']}H",
-                'aero_bkg_fhr': map(int, str(self.config['aero_bkg_times']).split(',')),
-                'OPREFIX': f"{self.runtime_config.CDUMP}.t{self.runtime_config.cyc:02d}z.",  # TODO: CDUMP is being replaced by RUN
-                'APREFIX': f"{self.runtime_config.CDUMP}.t{self.runtime_config.cyc:02d}z.",  # TODO: CDUMP is being replaced by RUN
-                'GPREFIX': f"gdas.t{self.runtime_config.previous_cycle.hour:02d}z.",
+                'AERO_WINDOW_LENGTH': f"PT{self.task_config['assim_freq']}H",
+                'aero_bkg_fhr': map(int, str(self.task_config['aero_bkg_times']).split(',')),
+                'OPREFIX': f"{self.task_config.RUN}.t{self.task_config.cyc:02d}z.",
+                'APREFIX': f"{self.task_config.RUN}.t{self.task_config.cyc:02d}z.",
+                'GPREFIX': f"gdas.t{self.task_config.previous_cycle.hour:02d}z.",
                 'jedi_yaml': _jedi_yaml,
             }
         )
 
-        # task_config is everything that this task should need
-        self.task_config = AttrDict(**self.config, **self.runtime_config, **local_dict)
+        # Extend task_config with local_dict
+        self.task_config = AttrDict(**self.task_config, **local_dict)
 
     @logit(logger)
     def initialize(self: Analysis) -> None:
@@ -157,8 +157,8 @@ def finalize(self: Analysis) -> None:
                 archive.add(diaggzip, arcname=os.path.basename(diaggzip))
 
         # copy full YAML from executable to ROTDIR
-        src = os.path.join(self.task_config['DATA'], f"{self.task_config['CDUMP']}.t{self.runtime_config['cyc']:02d}z.aerovar.yaml")
-        dest = os.path.join(self.task_config.COM_CHEM_ANALYSIS, f"{self.task_config['CDUMP']}.t{self.runtime_config['cyc']:02d}z.aerovar.yaml")
+        src = os.path.join(self.task_config['DATA'], f"{self.task_config['RUN']}.t{self.task_config['cyc']:02d}z.aerovar.yaml")
+        dest = os.path.join(self.task_config.COM_CHEM_ANALYSIS, f"{self.task_config['RUN']}.t{self.task_config['cyc']:02d}z.aerovar.yaml")
         yaml_copy = {
             'mkdir': [self.task_config.COM_CHEM_ANALYSIS],
             'copy': [[src, dest]]

ush/python/pygfs/task/aero_emissions.py

@@ -42,7 +42,9 @@ def __init__(self, config: Dict[str, Any]) -> None:
         localdict = AttrDict(
             {'variable_used_repeatedly': local_variable}
         )
-        self.task_config = AttrDict(**self.config, **self.runtime_config, **localdict)
+
+        # Extend task_config with localdict
+        self.task_config = AttrDict(**self.task_config, **localdict)
 
     @staticmethod
     @logit(logger)

ush/python/pygfs/task/aero_prepobs.py

@@ -24,23 +24,23 @@ class AerosolObsPrep(Task):
     def __init__(self, config: Dict[str, Any]) -> None:
         super().__init__(config)
 
-        _window_begin = add_to_datetime(self.runtime_config.current_cycle, -to_timedelta(f"{self.config['assim_freq']}H") / 2)
-        _window_end = add_to_datetime(self.runtime_config.current_cycle, +to_timedelta(f"{self.config['assim_freq']}H") / 2)
+        _window_begin = add_to_datetime(self.task_config.current_cycle, -to_timedelta(f"{self.task_config['assim_freq']}H") / 2)
+        _window_end = add_to_datetime(self.task_config.current_cycle, +to_timedelta(f"{self.task_config['assim_freq']}H") / 2)
 
         local_dict = AttrDict(
             {
                 'window_begin': _window_begin,
                 'window_end': _window_end,
-                'sensors': str(self.config['SENSORS']).split(','),
-                'data_dir': self.config['VIIRS_DATA_DIR'],
+                'sensors': str(self.task_config['SENSORS']).split(','),
+                'data_dir': self.task_config['VIIRS_DATA_DIR'],
                 'input_files': '',
-                'OPREFIX': f"{self.runtime_config.RUN}.t{self.runtime_config.cyc:02d}z.",
-                'APREFIX': f"{self.runtime_config.RUN}.t{self.runtime_config.cyc:02d}z."
+                'OPREFIX': f"{self.task_config.RUN}.t{self.task_config.cyc:02d}z.",
+                'APREFIX': f"{self.task_config.RUN}.t{self.task_config.cyc:02d}z."
             }
         )
 
         # task_config is everything that this task should need
-        self.task_config = AttrDict(**self.config, **self.runtime_config, **local_dict)
+        self.task_config = AttrDict(**self.task_config, **local_dict)
 
     @logit(logger)
     def initialize(self) -> None:
@@ -64,8 +64,8 @@ def initialize(self) -> None:
             self.task_config.prepaero_config = self.get_obsproc_config(sensor)
 
             # generate converter YAML file
-            template = f"{self.runtime_config.CDUMP}.t{self.runtime_config['cyc']:02d}z.prepaero_viirs_{sensor}.yaml"
-            _prepaero_yaml = os.path.join(self.runtime_config.DATA, template)
+            template = f"{self.task_config.RUN}.t{self.task_config['cyc']:02d}z.prepaero_viirs_{sensor}.yaml"
+            _prepaero_yaml = os.path.join(self.task_config.DATA, template)
             self.task_config.prepaero_yaml.append(_prepaero_yaml)
             logger.debug(f"Generate PrepAeroObs YAML file: {_prepaero_yaml}")
             save_as_yaml(self.task_config.prepaero_config, _prepaero_yaml)

ush/python/pygfs/task/analysis.py

@@ -27,7 +27,7 @@ class Analysis(Task):
     def __init__(self, config: Dict[str, Any]) -> None:
         super().__init__(config)
         # Store location of GDASApp jinja2 templates
-        self.gdasapp_j2tmpl_dir = os.path.join(self.config.PARMgfs, 'gdas')
+        self.gdasapp_j2tmpl_dir = os.path.join(self.task_config.PARMgfs, 'gdas')
 
     def initialize(self) -> None:
         super().initialize()
@@ -54,7 +54,7 @@ def get_jedi_config(self, algorithm: Optional[str] = None) -> Dict[str, Any]:
         ----------
         algorithm (optional) : str
             Name of the algorithm to use in the JEDI configuration. Will override the algorithm
-            set in the self.config.JCB_<>_YAML file
+            set in the self.task_config.JCB_<>_YAML file
 
         Returns
         ----------
@@ -120,7 +120,7 @@ def get_obs_dict(self) -> Dict[str, Any]:
             basename = os.path.basename(obfile)
             copylist.append([os.path.join(self.task_config['COM_OBS'], basename), obfile])
         obs_dict = {
-            'mkdir': [os.path.join(self.runtime_config['DATA'], 'obs')],
+            'mkdir': [os.path.join(self.task_config['DATA'], 'obs')],
             'copy': copylist
         }
         return obs_dict
@@ -161,7 +161,7 @@ def get_bias_dict(self) -> Dict[str, Any]:
                     # TODO: Why is this specific to ATMOS?
 
         bias_dict = {
-            'mkdir': [os.path.join(self.runtime_config.DATA, 'bc')],
+            'mkdir': [os.path.join(self.task_config.DATA, 'bc')],
             'copy': copylist
         }
         return bias_dict
@@ -180,7 +180,7 @@ def add_fv3_increments(self, inc_file_tmpl: str, bkg_file_tmpl: str, incvars: Li
            List of increment variables to add to the background
         """
 
-        for itile in range(1, self.config.ntiles + 1):
+        for itile in range(1, self.task_config.ntiles + 1):
             inc_path = inc_file_tmpl.format(tilenum=itile)
             bkg_path = bkg_file_tmpl.format(tilenum=itile)
             with Dataset(inc_path, mode='r') as incfile, Dataset(bkg_path, mode='a') as rstfile:
@@ -194,44 +194,6 @@ def add_fv3_increments(self, inc_file_tmpl: str, bkg_file_tmpl: str, incvars: Li
                     except (AttributeError, RuntimeError):
                         pass  # checksum is missing, move on
 
-    @logit(logger)
-    def get_bkg_dict(self, task_config: Dict[str, Any]) -> Dict[str, List[str]]:
-        """Compile a dictionary of model background files to copy
-
-        This method is a placeholder for now... will be possibly made generic at a later date
-
-        Parameters
-        ----------
-        task_config: Dict
-            a dictionary containing all of the configuration needed for the task
-
-        Returns
-        ----------
-        bkg_dict: Dict
-            a dictionary containing the list of model background files to copy for FileHandler
-        """
-        bkg_dict = {'foo': 'bar'}
-        return bkg_dict
-
-    @logit(logger)
-    def get_berror_dict(self, config: Dict[str, Any]) -> Dict[str, List[str]]:
-        """Compile a dictionary of background error files to copy
-
-        This method is a placeholder for now... will be possibly made generic at a later date
-
-        Parameters
-        ----------
-        config: Dict
-            a dictionary containing all of the configuration needed
-
-        Returns
-        ----------
-        berror_dict: Dict
-            a dictionary containing the list of background error files to copy for FileHandler
-        """
-        berror_dict = {'foo': 'bar'}
-        return berror_dict
-
     @logit(logger)
     def link_jediexe(self) -> None:
         """Compile a dictionary of background error files to copy
@@ -258,68 +220,6 @@ def link_jediexe(self) -> None:
 
         return exe_dest
 
-    @staticmethod
-    @logit(logger)
-    def get_fv3ens_dict(config: Dict[str, Any]) -> Dict[str, Any]:
-        """Compile a dictionary of ensemble member restarts to copy
-
-        This method constructs a dictionary of ensemble FV3 restart files (coupler, core, tracer)
-        that are needed for global atmens DA and returns said dictionary for use by the FileHandler class.
-
-        Parameters
-        ----------
-        config: Dict
-            a dictionary containing all of the configuration needed
-
-        Returns
-        ----------
-        ens_dict: Dict
-            a dictionary containing the list of ensemble member restart files to copy for FileHandler
-        """
-        # NOTE for now this is FV3 restart files and just assumed to be fh006
-
-        # define template
-        template_res = config.COM_ATMOS_RESTART_TMPL
-        prev_cycle = config.previous_cycle
-        tmpl_res_dict = {
-            'ROTDIR': config.ROTDIR,
-            'RUN': config.RUN,
-            'YMD': to_YMD(prev_cycle),
-            'HH': prev_cycle.strftime('%H'),
-            'MEMDIR': None
-        }
-
-        # construct ensemble member file list
-        dirlist = []
-        enslist = []
-        for imem in range(1, config.NMEM_ENS + 1):
-            memchar = f"mem{imem:03d}"
-
-            # create directory path for ensemble member restart
-            dirlist.append(os.path.join(config.DATA, config.dirname, f'mem{imem:03d}'))
-
-            # get FV3 restart files, this will be a lot simpler when using history files
-            tmpl_res_dict['MEMDIR'] = memchar
-            rst_dir = Template.substitute_structure(template_res, TemplateConstants.DOLLAR_CURLY_BRACE, tmpl_res_dict.get)
-            run_dir = os.path.join(config.DATA, config.dirname, memchar)
-
-            # atmens DA needs coupler
-            basename = f'{to_fv3time(config.current_cycle)}.coupler.res'
-            enslist.append([os.path.join(rst_dir, basename), os.path.join(config.DATA, config.dirname, memchar, basename)])
-
-            # atmens DA needs core, srf_wnd, tracer, phy_data, sfc_data
-            for ftype in ['fv_core.res', 'fv_srf_wnd.res', 'fv_tracer.res', 'phy_data', 'sfc_data']:
-                template = f'{to_fv3time(config.current_cycle)}.{ftype}.tile{{tilenum}}.nc'
-                for itile in range(1, config.ntiles + 1):
-                    basename = template.format(tilenum=itile)
-                    enslist.append([os.path.join(rst_dir, basename), os.path.join(run_dir, basename)])
-
-        ens_dict = {
-            'mkdir': dirlist,
-            'copy': enslist,
-        }
-        return ens_dict
-
     @staticmethod
     @logit(logger)
     def tgz_diags(statfile: str, diagdir: str) -> None: