Create JEDI class

scripts/exglobal_atm_analysis_fv3_increment.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 # exglobal_atm_analysis_fv3_increment.py
 # This script creates an AtmAnalysis object
-# and runs the init_fv3_increment and fv3_increment methods
+# and runs the initialize_fv3inc and execute methods
 # which convert the JEDI increment into an FV3 increment
 import os
 
@@ -17,7 +17,9 @@
     # Take configuration from environment and cast it as python dictionary
     config = cast_strdict_as_dtypedict(os.environ)
 
-    # Instantiate the atm analysis task
-    AtmAnl = AtmAnalysis(config)
-    AtmAnl.init_fv3_increment()
-    AtmAnl.fv3_increment()
+    # Instantiate the atm analysis object
+    AtmAnlFV3Inc = AtmAnalysis(config)
+
+    # Initialize and execute
+    AtmAnlFV3Inc.initialize_fv3inc()
+    AtmAnlFV3Inc.execute(config.APRUN_ATMANLFV3INC)

scripts/exglobal_atm_analysis_initialize.py

@@ -21,4 +21,4 @@
 
     # Instantiate the atm analysis task
     AtmAnl = AtmAnalysis(config)
-    AtmAnl.initialize()
+    AtmAnl.initialize_var()

scripts/exglobal_atm_analysis_variational.py

@@ -19,4 +19,4 @@
 
     # Instantiate the atm analysis task
     AtmAnl = AtmAnalysis(config)
-    AtmAnl.variational()
+    AtmAnl.execute(config.APRUN_ATMANLVAR, ['fv3jedi', 'variational'])

scripts/exglobal_atmens_analysis_fv3_increment.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 # exglobal_atmens_analysis_fv3_increment.py
 # This script creates an AtmEnsAnalysis object
-# and runs the init_fv3_increment and fv3_increment methods
+# and runs the initialize_fv3inc and execute methods
 # which convert the JEDI increment into an FV3 increment
 import os
 
@@ -17,7 +17,9 @@
     # Take configuration from environment and cast it as python dictionary
     config = cast_strdict_as_dtypedict(os.environ)
 
-    # Instantiate the atmens analysis task
-    AtmEnsAnl = AtmEnsAnalysis(config)
-    AtmEnsAnl.init_fv3_increment()
-    AtmEnsAnl.fv3_increment()
+    # Instantiate the atmens analysis object
+    AtmEnsAnlFV3Inc = AtmEnsAnalysis(config)
+
+    # Initialize and execute
+    AtmEnsAnlFV3Inc.initialize_fv3inc()
+    AtmEnsAnlFV3Inc.execute(config.APRUN_ATMENSANLFV3INC)

scripts/exglobal_atmens_analysis_initialize.py

@@ -21,4 +21,4 @@
 
     # Instantiate the atmens analysis task
     AtmEnsAnl = AtmEnsAnalysis(config)
-    AtmEnsAnl.initialize()
+    AtmEnsAnl.initialize_var()

scripts/exglobal_atmens_analysis_letkf.py

@@ -1,8 +1,8 @@
 #!/usr/bin/env python3
 # exglobal_atmens_analysis_letkf.py
 # This script creates an AtmEnsAnalysis object
-# and runs the letkf method
-# which executes the global atm local ensemble analysis
+# and runs the execute method which executes
+# the global atm local ensemble analysis
 import os
 
 from wxflow import Logger, cast_strdict_as_dtypedict
@@ -19,4 +19,4 @@
 
     # Instantiate the atmens analysis task
     AtmEnsAnl = AtmEnsAnalysis(config)
-    AtmEnsAnl.letkf()
+    AtmEnsAnl.execute(config.APRUN_ATMENSANLLETKF, ['fv3jedi', 'localensembleda']))

ush/python/pygfs/task/atm_analysis.py

@@ -5,24 +5,21 @@
 import gzip
 import tarfile
 from logging import getLogger
-from typing import Dict, List, Any
 
 from wxflow import (AttrDict,
                     FileHandler,
                     add_to_datetime, to_fv3time, to_timedelta, to_YMDH,
-                    chdir,
-                    parse_j2yaml, save_as_yaml,
-                    logit,
-                    Executable,
-                    WorkflowException)
-from pygfs.task.analysis import Analysis
+                    Task, 
+                    parse_j2yaml,
+                    logit)
+from pygfs.task.jedi import JEDI
 
 logger = getLogger(__name__.split('.')[-1])
 
 
-class AtmAnalysis(Analysis):
+class AtmAnalysis(Task):
     """
-    Class for global atm analysis tasks
+    Class for JEDI-based global atm analysis tasks
     """
     @logit(logger, name="AtmAnalysis")
     def __init__(self, config):
@@ -31,7 +28,6 @@ def __init__(self, config):
         _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.atmvar.yaml")
 
         # Create a local dictionary that is repeatedly used across this class
         local_dict = AttrDict(
@@ -48,7 +44,6 @@ def __init__(self, config):
                 '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,
                 'atm_obsdatain_path': f"{self.task_config.DATA}/obs/",
                 'atm_obsdataout_path': f"{self.task_config.DATA}/diags/",
                 'BKG_TSTEP': "PT1H"  # Placeholder for 4D applications
@@ -58,21 +53,39 @@ def __init__(self, config):
         # Extend task_config with local_dict
         self.task_config = AttrDict(**self.task_config, **local_dict)
 
+        # Create JEDI object
+        self.jedi = JEDI(self.task_config)
+
     @logit(logger)
-    def initialize(self: Analysis) -> None:
+    def initialize(self) -> None:
         """Initialize a global atm analysis
 
         This method will initialize a global atm analysis using JEDI.
         This includes:
+        - staging observation files
+        - staging bias correction files
         - staging CRTM fix files
         - staging FV3-JEDI fix files
         - staging B error files
         - staging model backgrounds
-        - generating a YAML file for the JEDI executable
         - creating output directories
         """
         super().initialize()
 
+        # get JEDI variational configuration and save to YAML file
+        self.jedi.get_config(self.task_config)
+
+        # link JEDI variational executable
+        self.jedi.link_exe(self.task_config)
+
+        # stage observations
+        obs_dict = self.jedi.get_obs_dict()
+        FileHandler(obs_dict).sync()
+
+        # stage bias corrections
+        bias_dict = self.jedi.get_bias_dict()
+        FileHandler(bias_dict).sync()
+
         # stage CRTM fix files
         logger.info(f"Staging CRTM fix files from {self.task_config.CRTM_FIX_YAML}")
         crtm_fix_list = parse_j2yaml(self.task_config.CRTM_FIX_YAML, self.task_config)
@@ -102,11 +115,6 @@ def initialize(self: Analysis) -> None:
         bkg_staging_dict = parse_j2yaml(self.task_config.VAR_BKG_STAGING_YAML, self.task_config)
         FileHandler(bkg_staging_dict).sync()
 
-        # generate variational YAML file
-        logger.debug(f"Generate variational YAML file: {self.task_config.jedi_yaml}")
-        save_as_yaml(self.task_config.jedi_config, self.task_config.jedi_yaml)
-        logger.info(f"Wrote variational YAML to: {self.task_config.jedi_yaml}")
-
         # need output dir for diags and anl
         logger.debug("Create empty output [anl, diags] directories to receive output from executable")
         newdirs = [
@@ -116,64 +124,23 @@ def initialize(self: Analysis) -> None:
         FileHandler({'mkdir': newdirs}).sync()
 
     @logit(logger)
-    def variational(self: Analysis) -> None:
-
-        chdir(self.task_config.DATA)
-
-        exec_cmd = Executable(self.task_config.APRUN_ATMANLVAR)
-        exec_name = os.path.join(self.task_config.DATA, 'gdas.x')
-        exec_cmd.add_default_arg(exec_name)
-        exec_cmd.add_default_arg('fv3jedi')
-        exec_cmd.add_default_arg('variational')
-        exec_cmd.add_default_arg(self.task_config.jedi_yaml)
-
-        try:
-            logger.debug(f"Executing {exec_cmd}")
-            exec_cmd()
-        except OSError:
-            raise OSError(f"Failed to execute {exec_cmd}")
-        except Exception:
-            raise WorkflowException(f"An error occured during execution of {exec_cmd}")
-
-        pass
+    def execute(self, aprun_cmd: str, jedi_args: Optional[str] = None) -> None:
+        super().execute()
 
+        self.jedi.execute(self.task_config, aprun_cmd, jedi_args)
+
     @logit(logger)
-    def init_fv3_increment(self: Analysis) -> None:
-        # Setup JEDI YAML file
-        self.task_config.jedi_yaml = os.path.join(self.task_config.DATA,
-                                                  f"{self.task_config.JCB_ALGO}.yaml")
-        save_as_yaml(self.get_jedi_config(self.task_config.JCB_ALGO), self.task_config.jedi_yaml)
-
-        # Link JEDI executable to run directory
-        self.task_config.jedi_exe = self.link_jediexe()
-
-    @logit(logger)
-    def fv3_increment(self: Analysis) -> None:
-        # Run executable
-        exec_cmd = Executable(self.task_config.APRUN_ATMANLFV3INC)
-        exec_cmd.add_default_arg(self.task_config.jedi_exe)
-        exec_cmd.add_default_arg(self.task_config.jedi_yaml)
-
-        try:
-            logger.debug(f"Executing {exec_cmd}")
-            exec_cmd()
-        except OSError:
-            raise OSError(f"Failed to execute {exec_cmd}")
-        except Exception:
-            raise WorkflowException(f"An error occured during execution of {exec_cmd}")
-
-    @logit(logger)
-    def finalize(self: Analysis) -> None:
+    def finalize(self) -> None:
         """Finalize a global atm analysis
 
         This method will finalize a global atm analysis using JEDI.
         This includes:
         - tar output diag files and place in ROTDIR
         - copy the generated YAML file from initialize to the ROTDIR
         - copy the updated bias correction files to ROTDIR
-        - write UFS model readable atm incrment file
-
         """
+        super().finalize()
+
         # ---- tar up diags
         # path of output tar statfile
         atmstat = os.path.join(self.task_config.COM_ATMOS_ANALYSIS, f"{self.task_config.APREFIX}atmstat")
@@ -198,7 +165,7 @@ def finalize(self: Analysis) -> None:
 
         # copy full YAML from executable to ROTDIR
         logger.info(f"Copying {self.task_config.jedi_yaml} to {self.task_config.COM_ATMOS_ANALYSIS}")
-        src = os.path.join(self.task_config.DATA, f"{self.task_config.RUN}.t{self.task_config.cyc:02d}z.atmvar.yaml")
+        src = self.task_config.jedi_yaml
         dest = os.path.join(self.task_config.COM_ATMOS_ANALYSIS, f"{self.task_config.RUN}.t{self.task_config.cyc:02d}z.atmvar.yaml")
         logger.debug(f"Copying {src} to {dest}")
         yaml_copy = {