Switch Aerosol DA to use JCB and Jedi class (NOAA-EMC#3125)

This PR switches to using JCB for aerosol DA YAML generation + makes use of the Jedi class refactoring already done by @DavidNew-NOAA for the atmospheric JEDI-based DA. Resolves NOAA-EMC#3035 Note this PR is in draft for 2 reasons. 1) Need associated GDASApp PR to be opened and merged first and 2) requires some fix file updates. This will be coordinated with an issue and linked to this PR once @andytangborn is satisfied with a first-attempt at the fix files. --------- Co-authored-by: RussTreadon-NOAA <Russ.Treadon@noaa.gov> Co-authored-by: DavidNew-NOAA <david.new@noaa.gov>
sbanihash · Dec 14, 2024 · 220bae8 · 220bae8
1 parent 41972b6
commit 220bae8
Show file tree

Hide file tree

Showing 16 changed files with 238 additions and 357 deletions.
diff --git a/parm/config/gfs/config.aeroanl b/parm/config/gfs/config.aeroanl
@@ -7,41 +7,39 @@ echo "BEGIN: config.aeroanl"
 
 # define analysis resolution based on deterministic res
 case ${CASE} in
-  "C1152" | "C768" | "C384" | "C192")
-    CASE_ANL="C192"
+  "C1152" | "C768" | "C384")
+    CASE_ANL="C384"
     ;;
-  "C96" | "C48")
+  "C192" | "C96" | "C48")
     CASE_ANL=${CASE}
     ;;
   *)
     echo "FATAL ERROR: Aerosol DA not supported at ${CASE} resolution"
     exit 4
 esac
 export CASE_ANL
-export OBS_LIST="${PARMgfs}/gdas/aero/obs/lists/gdas_aero.yaml.j2"
+export JCB_ALGO_YAML_VAR=${PARMgfs}/gdas/aero/jcb-prototype_3dvar.yaml.j2
 export STATICB_TYPE='diffusion'
-export BERROR_YAML="${PARMgfs}/gdas/aero/berror/staticb_${STATICB_TYPE}.yaml.j2"
+export BERROR_YAML="aero_background_error_static_${STATICB_TYPE}"
 export BERROR_DATA_DIR="${FIXgfs}/gdas/aero/clim_b"
 
 export CRTM_FIX_YAML="${PARMgfs}/gdas/aero_crtm_coeff.yaml.j2"
 export JEDI_FIX_YAML="${PARMgfs}/gdas/aero_jedi_fix.yaml.j2"
 
+export JEDI_CONFIG_YAML="${PARMgfs}/gdas/aeroanl_jedi_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 io_layout_x=@IO_LAYOUT_X@
 export io_layout_y=@IO_LAYOUT_Y@
 
 export JEDIEXE="${EXECgfs}/gdas.x"
-export BMATEXE="${EXECgfs}/gdasapp_chem_diagb.x"
-export DIFFUSIONEXE="${EXECgfs}/gdas_fv3jedi_error_covariance_toolbox.x"
 
 if [[ "${DOIAU}" == "YES" ]]; then
   export aero_bkg_times="3,6,9"
-  export JEDIYAML="${PARMgfs}/gdas/aero/variational/3dvar_fgat_gfs_aero.yaml.j2"
 else
   export aero_bkg_times="6," # Trailing comma is necessary so this is treated as a list
-  export JEDIYAML="${PARMgfs}/gdas/aero/variational/3dvar_gfs_aero.yaml.j2"
 fi
 
 echo "END: config.aeroanl"
diff --git a/parm/config/gfs/config.aeroanlgenb b/parm/config/gfs/config.aeroanlgenb
@@ -8,9 +8,8 @@ echo "BEGIN: config.aeroanlgenb"
 # Get task specific resources
 source "${EXPDIR}/config.resources" aeroanlgenb
 
-export BMATYAML="${PARMgfs}/gdas/aero/berror/aero_diagb.yaml.j2"
-export DIFFUSIONYAML="${PARMgfs}/gdas/aero/berror/aero_diffusionparm.yaml.j2"
-export INTERPYAML="${PARMgfs}/gdas/aero/berror/aero_interp.yaml.j2"
+export JEDI_CONFIG_YAML="${PARMgfs}/gdas/aero_bmat_jedi_config.yaml.j2"
+export JCB_BASE_YAML="${PARMgfs}/gdas/aero/jcb-base.yaml.j2"
 export AERO_BMATRIX_STAGE_TMPL="${PARMgfs}/gdas/aero_stage_bmatrix_bkg.yaml.j2"
 export AERO_BMATRIX_FINALIZE_TMPL="${PARMgfs}/gdas/aero_finalize_bmatrix_bkg.yaml.j2"
 export aero_diffusion_iter=10

diff --git a/parm/gdas/aero_bmat_jedi_config.yaml.j2 b/parm/gdas/aero_bmat_jedi_config.yaml.j2
@@ -0,0 +1,19 @@
+aero_interpbkg:
+  rundir: '{{ DATA }}'
+  exe_src: '{{ EXECgfs }}/gdas.x'
+  jedi_args: ['fv3jedi', 'convertstate']
+  mpi_cmd: '{{ APRUN_AEROANLGENB }}'
+  jcb_base_yaml: '{{ PARMgfs }}/gdas/aero/jcb-base.yaml.j2'
+  jcb_algo: aero_convert_background
+aero_diagb:
+  rundir: '{{ DATA }}'
+  exe_src: '{{ EXECgfs }}/gdasapp_chem_diagb.x'
+  mpi_cmd: '{{ APRUN_AEROANLGENB }}'
+  jcb_base_yaml: '{{ PARMgfs }}/gdas/aero/jcb-base.yaml.j2'
+  jcb_algo: aero_gen_bmatrix_diagb
+aero_diffusion:
+  rundir: '{{ DATA }}'
+  exe_src: '{{ EXECgfs }}/gdas_fv3jedi_error_covariance_toolbox.x'
+  mpi_cmd: '{{ APRUN_AEROANLGENB }}'
+  jcb_base_yaml: '{{ PARMgfs }}/gdas/aero/jcb-base.yaml.j2'
+  jcb_algo: aero_gen_bmatrix_diffusion
diff --git a/parm/gdas/aero_finalize_bmatrix_bkg.yaml.j2 b/parm/gdas/aero_finalize_bmatrix_bkg.yaml.j2
@@ -4,9 +4,9 @@
 {% set background_time = current_cycle | add_to_datetime(offset_td) %}
 copy:
 ### copy YAMLs used
-{% set yaml_list = ['chem_diagb.yaml', 'chem_diffusion.yaml'] %}
+{% set yaml_list = ['aero_diagb.yaml', 'aero_diffusion.yaml', 'aero_interpbkg.yaml'] %}
 {% for fname in yaml_list %}
-- ["{{ DATA }}/{{ HEAD }}{{ fname }}", "{{ COMOUT_CHEM_BMAT }}/{{ HEAD }}{{ fname }}"]
+- ["{{ DATA }}/{{ fname }}", "{{ COMOUT_CHEM_BMAT }}/{{ HEAD }}{{ fname }}"]
 {% endfor %}
 ### copy stddev files to ROTDIR
 {% for tile in range(1, ntiles+1) %}

diff --git a/parm/gdas/aero_finalize_variational.yaml.j2 b/parm/gdas/aero_finalize_variational.yaml.j2
@@ -12,7 +12,7 @@ mkdir:
 - "{{ COMOUT_ATMOS_RESTART }}"
 copy:
 ## copy variational YAML to ROTDIR
-- ["{{ DATA }}/{{ APREFIX }}aerovar.yaml", "{{ COMOUT_CHEM_ANALYSIS }}/{{ APREFIX }}aerovar.yaml"]
+- ["{{ DATA }}/aeroanlvar.yaml", "{{ COMOUT_CHEM_ANALYSIS }}/{{ APREFIX }}aerovar.yaml"]
 ## copy increments
 {% for tile in range(1,ntiles+1) %}
 - ["{{ DATA }}/anl/aeroinc.{{ current_cycle | to_fv3time }}.fv_tracer.res.tile{{ tile }}.nc", "{{ COMOUT_CHEM_ANALYSIS }}/aeroinc.{{ current_cycle | to_fv3time }}.fv_tracer.res.tile{{ tile }}.nc"]

diff --git a/parm/gdas/aero_stage_bmatrix_bkg.yaml.j2 b/parm/gdas/aero_stage_bmatrix_bkg.yaml.j2
@@ -12,6 +12,7 @@ mkdir:
 - "{{ DATA }}/stddev"
 - "{{ DATA }}/clm_stddev"
 - "{{ DATA }}/diffusion"
+- "{{ DATA }}/rescale"
 copy:
 ######################################
 # copy deterministic background files
@@ -30,9 +31,15 @@ copy:
 - ["{{ COM_ATMOS_RESTART_TMPL | replace_tmpl(tmpl_dict) }}/{{ background_time | to_fv3time }}.{{ ftype }}.tile{{ tile }}.nc", "{{ DATA }}/bkg/{{ background_time | to_fv3time }}.{{ ftype }}.tile{{ tile }}.nc"]
   {% endfor %}
 {% endfor %}
+######################################
 # copy climatological stddev files
 ######################################
 {% for tile in range(1, ntiles+1) %}
 - ["{{ BERROR_DATA_DIR }}/stddev.fv_tracer.res.tile{{ tile }}.nc", "{{ DATA }}/clm_stddev/stddev.fv_tracer.res.tile{{ tile }}.nc"]
 {% endfor %}
-
+######################################
+# copy stddev rescaling factor files
+######################################
+{% for tile in range(1, ntiles+1) %}
+- ["{{ BERROR_DATA_DIR }}/rescale.fv_tracer.res.tile{{ tile }}.nc", "{{ DATA }}/rescale/rescale.fv_tracer.res.tile{{ tile }}.nc"]
+{% endfor %}
diff --git a/parm/gdas/aeroanl_jedi_config.yaml.j2 b/parm/gdas/aeroanl_jedi_config.yaml.j2
@@ -0,0 +1,7 @@
+aeroanlvar:
+  rundir: '{{ DATA }}'
+  exe_src: '{{ EXECgfs }}/gdas.x'
+  mpi_cmd: '{{ APRUN_AEROANL }}'
+  jedi_args: ['fv3jedi', 'variational']
+  jcb_base_yaml: '{{ PARMgfs }}/gdas/aero/jcb-base.yaml.j2'
+  jcb_algo_yaml: '{{ JCB_ALGO_YAML_VAR }}'
diff --git a/scripts/exgdas_aero_analysis_generate_bmatrix.py b/scripts/exgdas_aero_analysis_generate_bmatrix.py
@@ -18,10 +18,8 @@
     # Take configuration from environment and cast it as python dictionary
     config = cast_strdict_as_dtypedict(os.environ)
 
-    # Instantiate the aerosol variance and diffusion correlation tasks
-    AeroB = AerosolBMatrix(config)
-    AeroB.initialize()
-    AeroB.interpBackground()
-    AeroB.computeVariance()
-    AeroB.computeDiffusion()
-    AeroB.finalize()
+    # Create an instance of the AerosolBMatrix task
+    aeroBMat = AerosolBMatrix(config)
+    aeroBMat.initialize()
+    aeroBMat.execute()
+    aeroBMat.finalize()
diff --git a/scripts/exglobal_aero_analysis_finalize.py b/scripts/exglobal_aero_analysis_finalize.py
@@ -21,4 +21,6 @@
 
     # Instantiate the aerosol analysis task
     AeroAnl = AerosolAnalysis(config)
+
+    # Finalize JEDI aerosol variational analysis
     AeroAnl.finalize()
diff --git a/scripts/exglobal_aero_analysis_initialize.py b/scripts/exglobal_aero_analysis_initialize.py
@@ -21,4 +21,6 @@
 
     # Instantiate the aerosol analysis task
     AeroAnl = AerosolAnalysis(config)
+
+    # Initialize JEDI aerosol variational analysis
     AeroAnl.initialize()
diff --git a/scripts/exglobal_aero_analysis_variational.py b/scripts/exglobal_aero_analysis_variational.py
@@ -19,4 +19,6 @@
 
     # Instantiate the aerosol analysis task
     AeroAnl = AerosolAnalysis(config)
-    AeroAnl.variational()
+
+    # Execute JEDI variational analysis
+    AeroAnl.execute('aeroanlvar')
diff --git a/sorc/gdas.cd b/sorc/gdas.cd
diff --git a/ush/python/pygfs/__init__.py b/ush/python/pygfs/__init__.py
@@ -2,7 +2,6 @@
 import os
 
 from .task.analysis import Analysis
-from .task.bmatrix import BMatrix
 from .task.aero_emissions import AerosolEmissions
 from .task.aero_analysis import AerosolAnalysis
 from .task.aero_bmatrix import AerosolBMatrix
+19 −19		parm/aero/berror/aero_diagb.yaml.j2
+10 −7		parm/aero/jcb-base.yaml.j2
+6 −6		parm/soca/obsprep/obsprep_config.yaml