Add capability to create aerosol AOD analysis products in GRIB2 format via UPP

dev/parm/config/gcafs/config.aeroanl.j2

@@ -29,7 +29,16 @@
 export STAGE_YAML="${PARMgfs}/gdas/aero/aero_det_stage.yaml.j2"
 export SAVE_YAML="${PARMgfs}/gdas/aero/aero_det_save.yaml.j2"
 
+<<<<<<< HEAD
+export JEDI_CONFIG_YAML="${PARMgfs}/gdas/aeroanl_jedi_config.yaml.j2"
+export UPP_CONFIG_YAML="${PARMgfs}/gdas/aeroanl_upp_config.yaml.j2"
+
+export AERO_STAGE_VARIATIONAL_TMPL="${PARMgfs}/gdas/aero_stage_variational.yaml.j2"
+export AERO_FINALIZE_VARIATIONAL_TMPL="${PARMgfs}/gdas/aero_finalize_variational.yaml.j2"
+export AERO_BMATRIX_RESCALE_TMPL="aero_gen_bmatrix_rescale_default.yaml.j2"
+=======
 export AERO_BMATRIX_RESCALE_YAML="aero_gen_bmatrix_rescale_default.yaml.j2"
+>>>>>>> upstream/develop
 
 export io_layout_x="{{ IO_LAYOUT_X }}"
 export io_layout_y="{{ IO_LAYOUT_Y }}"

dev/parm/config/gcafs/config.resources

@@ -520,10 +520,10 @@ case ${step} in
         echo "FATAL ERROR: Resources not defined for job ${step} at resolution ${CASE}"
         exit 4
     esac
-    ntasks=1
+    ntasks=120
     threads_per_task=1
     tasks_per_node=$(( max_tasks_per_node / threads_per_task ))
-    memory="13072M"
+    memory=${mem_node_max}
     ;;
 
   "marineanlinit")

dev/parm/config/gfs/config.com

@@ -110,5 +110,6 @@ declare -rx COM_CHEM_HISTORY_TMPL=${COM_BASE}'/model/chem/history'
 declare -rx COM_CHEM_ANALYSIS_TMPL=${COM_BASE}'/analysis/chem'
 declare -rx COM_CHEM_BMAT_TMPL=${COM_CHEM_ANALYSIS_TMPL}'/bmatrix'
 declare -rx COM_CHEM_ANLMON_TMPL=${COM_BASE}'/products/chem/anlmon'
+declare -rx COM_CHEM_GRIB_TMPL=${COM_BASE}'/products/chem/grib2'
 
 declare -rx COM_MED_RESTART_TMPL=${COM_BASE}'/model/med/restart'

env/GAEAC5.env

@@ -97,6 +97,13 @@ case ${step} in
     export NTHREADS_AEROANL=${NTHREADSmax}
     export APRUN_AEROANL="${APRUN_default} --cpus-per-task=${NTHREADS_AEROANL}"
  ;;
+ "aeroanlfinal")
+
+    export APRUNCFP="${launcher} -n \$ncmd ${mpmd_opt}"
+
+    export NTHREADS_AEROANLFINAL=${NTHREADSmax}
+    export APRUN_AEROANLFINAL="${APRUN_default} --cpus-per-task=${NTHREADS_AEROANLFINAL}"
+ ;;
  "aeroanlgenb")
 
     export NTHREADS_AEROANLGENB=${NTHREADSmax}

env/GAEAC6.env

@@ -97,6 +97,13 @@ case ${step} in
     export NTHREADS_AEROANL=${NTHREADSmax}
     export APRUN_AEROANL="${APRUN_default} --cpus-per-task=${NTHREADS_AEROANL}"
  ;;
+ "aeroanlfinal")
+
+    export APRUNCFP="${launcher} -n \$ncmd ${mpmd_opt}"
+
+    export NTHREADS_AEROANLFINAL=${NTHREADSmax}
+    export APRUN_AEROANLFINAL="${APRUN_default} --cpus-per-task=${NTHREADS_AEROANLFINAL}"
+ ;;
  "aeroanlgenb")
 
     export NTHREADS_AEROANLGENB=${NTHREADSmax}

env/HERA.env

@@ -98,6 +98,13 @@ elif [[ "${step}" = "aeroanlvar" ]]; then
     export NTHREADS_AEROANL=${NTHREADSmax}
     export APRUN_AEROANL="${APRUN_default} --cpus-per-task=${NTHREADS_AEROANL}"
 
+elif [[ "${step}" = "aeroanlfinal" ]]; then
+
+    export APRUNCFP="${launcher} -n \$ncmd ${mpmd_opt}"
+
+    export NTHREADS_AEROANLFINAL=${NTHREADSmax}
+    export APRUN_AEROANLFINAL="${APRUN_default} --cpus-per-task=${NTHREADS_AEROANLFINAL}"
+
 elif [[ "${step}" = "aeroanlgenb" ]]; then
 
     export NTHREADS_AEROANLGENB=${NTHREADSmax}

env/HERCULES.env

@@ -102,6 +102,13 @@ case ${step} in
 
     export NTHREADS_AEROANL=${NTHREADSmax}
     export APRUN_AEROANL="${APRUN_default} --cpus-per-task=${NTHREADS_AEROANL}"
+ ;;
+ "aeroanlfinal")
+
+    export APRUNCFP="${launcher} -n \$ncmd ${mpmd_opt}"
+
+    export NTHREADS_AEROANLFINAL=${NTHREADSmax}
+    export APRUN_AEROANLFINAL="${APRUN_default} --cpus-per-task=${NTHREADS_AEROANLFINAL}"
  ;;
   "aeroanlgenb")
 

env/ORION.env

@@ -94,6 +94,13 @@ elif [[ "${step}" = "aeroanlvar" ]]; then
     export NTHREADS_AEROANL=${NTHREADSmax}
     export APRUN_AEROANL="${APRUN_default} --cpus-per-task=${NTHREADS_AEROANL}"
 
+elif [[ "${step}" = "aeroanlfinal" ]]; then
+
+    export APRUNCFP="${launcher} -n \$ncmd ${mpmd_opt}"
+
+    export NTHREADS_AEROANLFINAL=${NTHREADSmax}
+    export APRUN_AEROANLFINAL="${APRUN_default} --cpus-per-task=${NTHREADS_AEROANLFINAL}"
+
 elif [[ "${step}" = "aeroanlgenb" ]]; then
 
     export NTHREADS_AEROANLGENB=${NTHREADSmax}

env/WCOSS2.env

@@ -80,6 +80,13 @@ elif [[ "${step}" = "aeroanlvar" ]]; then
     export NTHREADS_AEROANL=${NTHREADSmax}
     export APRUN_AEROANL="${APRUN_default}"
 
+elif [[ "${step}" = "aeroanlfinal" ]]; then
+
+    export APRUNCFP="${launcher} -np \$ncmd ${mpmd_opt}"
+
+    export NTHREADS_AEROANLFINAL=${NTHREADSmax}
+    export APRUN_AEROANLFINAL="${APRUN_default}"
+
 elif [[ "${step}" = "aeroanlgenb" ]]; then
 
     export NTHREADS_AEROANLGENB=${NTHREADSmax}

jobs/JGLOBAL_AERO_ANALYSIS_FINALIZE

@@ -16,7 +16,8 @@ source "${HOMEgfs}/ush/jjob_header.sh" -e "aeroanlfinal" -c "base aeroanl aeroan
 YMD=${PDY} HH=${cyc} declare_from_tmpl -rx \
     COMOUT_CHEM_ANALYSIS:COM_CHEM_ANALYSIS_TMPL \
     COMOUT_CONF:COM_CONF_TMPL \
-    COMOUT_ATMOS_RESTART:COM_ATMOS_RESTART_TMPL
+    COMOUT_ATMOS_RESTART:COM_ATMOS_RESTART_TMPL \
+    COMOUT_CHEM_GRIB:COM_CHEM_GRIB_TMPL
 
 ###############################################################
 # Run relevant script

jobs/JGLOBAL_AERO_ANALYSIS_INITIALIZE

@@ -24,10 +24,12 @@ GDUMP="${GDUMP/gcafs/gcdas}"
 # Generate COM variables from templates
 YMD=${PDY} HH=${cyc} declare_from_tmpl -rx \
     COMIN_OBS:COM_OBS_TMPL \
+    COMIN_ATMOS_ANALYSIS:COM_ATMOS_ANALYSIS_TMPL \
     COMOUT_CHEM_ANALYSIS:COM_CHEM_ANALYSIS_TMPL
 
 RUN=${GDUMP} YMD=${gPDY} HH=${gcyc} declare_from_tmpl -rx \
     COMIN_ATMOS_RESTART_PREV:COM_ATMOS_RESTART_TMPL \
+    COMIN_ATMOS_HISTORY_PREV:COM_ATMOS_HISTORY_TMPL \
     COMIN_CHEM_BMAT_PREV:COM_CHEM_BMAT_TMPL \
     COMIN_CHEM_ANALYSIS_PREV:COM_CHEM_ANALYSIS_TMPL
 

parm/post/itag.jinja

@@ -12,5 +12,6 @@
   kpo = {{ po | length }},
   po = {{ po | join(', ') }},
   rdaod = {{ rdaod | to_f90bool }}
+  nasa_on = {{ nasa_on | to_f90bool }}
 /
 

parm/post/upp.yaml

@@ -4,6 +4,7 @@ upp:
     ioform: "netcdfpara"
     po: [1000.,975.,950.,925.,900.,875.,850.,825.,800.,775.,750.,725.,700.,675.,650.,625.,600.,575.,550.,525.,500.,475.,450.,425.,400.,375.,350.,325.,300.,275.,250.,225.,200.,175.,150.,125.,100.,70.,50.,40.,30.,20.,15.,10.,7.,5.,3.,2.,1.,0.7,0.4,0.2,0.1,0.07,0.04,0.02,0.01]
     rdaod: False
+    nasa_on: False
   fix_data:
     mkdir:
       - "{{ DATA }}"

scripts/exglobal_aero_analysis_finalize.py

@@ -8,6 +8,7 @@
 
 from wxflow import Logger, cast_strdict_as_dtypedict
 from pygfs.task.aero_analysis import AerosolAnalysis
+from pygfs.task.upp import UPP
 
 
 # Initialize root logger
@@ -22,5 +23,8 @@
     # Instantiate the aerosol analysis task
     AeroAnl = AerosolAnalysis(config)
 
+    # Process aerosal products through UPP
+    AeroAnl.upp_anlproc()
+
     # Finalize JEDI aerosol variational analysis
     AeroAnl.finalize()

ush/python/pygfs/task/aero_analysis.py

@@ -11,13 +11,14 @@
 
 from wxflow import (AttrDict,
                     FileHandler,
-                    add_to_datetime, to_fv3time, to_timedelta,
+                    add_to_datetime, to_timedelta,
                     to_fv3time,
-                    Task,
+                    Task, Jinja,
                     YAMLFile, parse_j2yaml,
                     logit)
 from pygfs.jedi import Jedi
 import numpy as np
+from pygfs.task.upp import UPP
 
 logger = getLogger(__name__.split('.')[-1])
 
@@ -262,3 +263,133 @@ def add_fv3_increments(self, inc_file_YAML: str, bkg_file_YAML: str, incvars: Li
                         rstfile.variables[vname].delncattr('checksum')  # remove the checksum so fv3 does not complain
                     except (AttributeError, RuntimeError):
                         pass  # checksum is missing, move on
+
+    @logit(logger)
+    def upp_anlproc(self) -> None:
+        """Process aerosol analysis to GRIB2
+
+        This method processes aerosol analysis products from tracer fields using UPP.
+        This includes:
+        - Creating a UPP object
+        - Staging UPP fix files
+        - Creating the 'upp_dict' for UPP object
+        - Generating the upp namelist
+        - Adding atmos and aerosol increments to the background
+        - Execute upp.x
+        """
+
+        local_dict = AttrDict(
+            {
+                'UPP_RUN': "analysis",
+                'FORECAST_HOUR': 0
+            }
+        )
+        self.task_config = AttrDict(**self.task_config, **local_dict)
+        self.task_config.UPP_CONFIG = self.task_config.UPP_CONFIG_YAML
+        upp = UPP(self.task_config)
+
+        upp_yaml = upp.task_config.upp_yaml
+        upp.initialize(upp_yaml)
+
+        upp_dict = AttrDict()
+        keys = ['APRUN_AEROANLFINAL', 'forecast_hour',
+                'atmos_filename', 'flux_filename']
+
+        upp_dict = AttrDict()
+        for key in keys:
+            upp_dict[key] = upp.task_config[key]
+
+        upp_dict['NET'] = 'gfs'   # set to 'gfs' so upp can recognize
+        upp_dict['valid_datetime'] = self.task_config.current_cycle
+        upp_dict['DATA'] = os.path.join(self.task_config.DATA, 'upp')
+        upp_dict.update(upp_yaml['upp']['config'])
+
+        # Configure the namelist and write to file
+        logger.info("Creating namelist for upp.x")
+        nml_template = os.path.join(upp_dict.DATA, "itag.jinja")
+        nml_data = Jinja(nml_template, upp_dict).render
+        logger.debug(f"itag:\n{nml_data}")
+        nml_file = os.path.join(upp_dict.DATA, 'itag')
+        with open(nml_file, "w") as fho:
+            fho.write(nml_data)
+
+        # ---- add aero increments to atmf000 files
+        logger.info('Adding aero increments to RESTART files')
+        bkg_file = os.path.join(upp_dict.DATA, f"{upp_dict.atmos_filename}")
+        inc_filename = f"aeroinc_gauss.{self.task_config.current_cycle.strftime('%Y-%m-%dT%H:%M:%S')}Z.gaussian.modelLevels.nc"
+        inc_file = os.path.join(self.task_config.DATA, 'anl', inc_filename)
+        allvars = upp_yaml['aeroincvars'][:]
+        bkgvars = [var[0] for var in allvars]
+        incvars = [var[1] for var in allvars]
+        self.add_aero_gaussian_increments(inc_file, bkg_file, incvars, bkgvars)
+
+        # ---- add atmo increments to atmf000 files
+        logger.info('Adding atmo increments to RESTART files')
+        inc_file = os.path.join(upp_dict.DATA, f"{self.task_config.APREFIX}atminc.nc")
+        allvars = upp_yaml['atmincvars'][:]
+        bkgvars = [var[0] for var in allvars]
+        incvars = [var[1] for var in allvars]
+        self.add_atm_gaussian_increments(inc_file, bkg_file, incvars, bkgvars)
+
+        # reset time to 0 (analysis time)
+        flux_file = os.path.join(upp_dict.DATA, f"{upp_dict.flux_filename}")
+        with Dataset(flux_file, mode='a') as rstfile:
+            time = rstfile.variables['time']
+            time[:] = 0.0
+            time.setncattr("units", f"hours since {self.task_config.current_cycle.strftime('%Y-%m-%d %H:%M:%S')}")
+
+        upp.execute(upp_dict.DATA, upp_dict.APRUN_AEROANLFINAL, upp_dict.forecast_hour)
+
+    @logit(logger)
+    def add_aero_gaussian_increments(self, inc_file: str, bkg_file: str, incvars: List, bkgvars: List) -> None:
+        """Add aero gaussian increments to gaussian backgrounds
+
+        Parameters
+        ----------
+        inc_file : str
+           increment file
+        bkg_file : str
+           background file
+        incvars : List
+           List of increment variables to add to the background
+        bkgvars : List
+           List of background variables to which the increment variables will be added.
+        """
+        with Dataset(inc_file, mode='r') as incfile, Dataset(bkg_file, mode='a') as rstfile:
+            for incname, bkgname in zip(incvars, bkgvars):
+                increment = incfile.variables[incname][:]
+                # reordering the dimensions of increment to macth background
+                increment_reshape = np.transpose(increment, (2, 0, 1))
+
+                bkg = rstfile.variables[bkgname][:]
+                anl = bkg + increment_reshape[np.newaxis, :, :, :]
+                rstfile.variables[bkgname][:] = anl[:]
+            time = rstfile.variables['time']
+            time[:] = 0.0
+            time.setncattr("units", f"hours since {self.task_config.current_cycle.strftime('%Y-%m-%d %H:%M:%S')}")
+
+    @logit(logger)
+    def add_atm_gaussian_increments(self, inc_file: str, bkg_file: str, incvars: List, bkgvars: List) -> None:
+        """Add atm gaussian increments to gaussian backgrounds
+
+        Parameters
+        ----------
+        inc_file : str
+           increment file
+        bkg_file : str
+           background file
+        incvars : List
+           List of increment variables to add to the background
+        bkgvars : List
+           List of background variables to which the increment variables will be added.
+        """
+        with Dataset(inc_file, mode='r') as incfile, Dataset(bkg_file, mode='a') as rstfile:
+            for incname, bkgname in zip(incvars, bkgvars):
+                increment = incfile.variables[incname][:]
+                # handel latitude inversion in atminc
+                lat_axis_index = 1
+                increment_lat_inversion = np.flip(increment, axis=lat_axis_index)
+
+                bkg = rstfile.variables[bkgname][:]
+                anl = bkg + increment_lat_inversion[np.newaxis, :, :, :]
+                rstfile.variables[bkgname][:] = anl[:]