use jcb for soca without actually using it

parm/atm/jcb-base.yaml.j2

@@ -10,64 +10,56 @@ app_path_observation_chronicle: "{{PARMgfs}}/gdas/jcb-gdas/observation_chronicle
 # Places where we deviate from the generic file name of a yaml
 # ------------------------------------------------------------
 #final_increment_file: final_increment_gaussian
-final_increment_file: final_increment_cubed_sphere
-output_ensemble_increments_file: output_ensemble_increments_cubed_sphere
-model_file: model_pseudo
-initial_condition_file: background  # Initial conditions for 4D apps is background
-
+final_increment_file: atmosphere_final_increment_cubed_sphere
+output_ensemble_increments_file: atmosphere_output_ensemble_increments_cubed_sphere
+model_file: atmosphere_model_pseudo
+initial_condition_file: atmosphere_background  # Initial conditions for 4D apps is background
+background_error_file: "{{BERROR_YAML}}"
 
-# Assimilation window
-# -------------------
+# Assimilation standard things (not prepended with model choice)
+# ----------------------------
 window_begin: "{{ ATM_WINDOW_BEGIN | to_isotime }}"
 window_length: "{{ ATM_WINDOW_LENGTH }}"
 bound_to_include: begin
-
-# Default background time is for 3D applications
-atm_background_time_iso: "{{ current_cycle | to_isotime }}"
-
-
-# Analysis standard things
-# ------------------------
-analysis_variables: [ua,va,t,ps,sphum,ice_wat,liq_wat,o3mr]
 minimizer: DRPCG
 final_diagnostics_departures: anlmob
-
-cost_type: 3D-Var
 number_of_outer_loops: 2
+analysis_variables: [ua,va,t,ps,sphum,ice_wat,liq_wat,o3mr]
 
 # Model things
 # ------------
 # Geometry
-layout_x: {{ layout_x | default(1, true) }}
-layout_y: {{ layout_y | default(1, true) }}
-npx_ges: {{ npx_ges | default(49, true) }}
-npy_ges: {{ npy_ges | default(49, true) }}
-npz_ges: {{ npz_ges | default(127, true) }}
-npx_anl: {{ npx_anl | default(49, true) }}
-npy_anl: {{ npy_anl | default(49, true) }}
-npz_anl: {{ npz_anl | default(127, true) }}
+atmosphere_layout_x: {{ layout_x | default(1, true) }}
+atmosphere_layout_y: {{ layout_y | default(1, true) }}
+atmosphere_npx_ges: {{ npx_ges | default(49, true) }}
+atmosphere_npy_ges: {{ npy_ges | default(49, true) }}
+atmosphere_npz_ges: {{ npz_ges | default(127, true) }}
+atmosphere_npx_anl: {{ npx_anl | default(49, true) }}
+atmosphere_npy_anl: {{ npy_anl | default(49, true) }}
+atmosphere_npz_anl: {{ npz_anl | default(127, true) }}
 
-fv3jedi_files_path: ./fv3jedi  # Ideally this would be {{DATA}}/fv3jedi but FMS
+atmosphere_fv3jedi_files_path: ./fv3jedi  # Ideally this would be {{DATA}}/fv3jedi but FMS
 
 # Background
-background_path: ./bkg
-background_ensemble_path: ./ens/mem%mem%
+atmosphere_background_path: ./bkg
+atmosphere_background_ensemble_path: ./ens/mem%mem%
+
+# Default background time is for 3D applications
+atmosphere_background_time_iso: "{{ current_cycle | to_isotime }}"
 
 # Background error
-bump_data_directory: "{{ DATA }}/berror"
-gsibec_path: "{{ DATA }}/berror"
-number_ensemble_members: {{ NMEM_ENS | default(1, true) }}
-layout_gsib_x: {{ layout_gsib_x | default(1, true) }}
-layout_gsib_y: {{ layout_gsib_y | default(1, true) }}
+atmosphere_bump_data_directory: "{{ DATA }}/berror"
+atmosphere_gsibec_path: "{{ DATA }}/berror"
+atmosphere_number_ensemble_members: {{ NMEM_ENS | default(1, true) }}
+atmosphere_layout_gsib_x: {{ layout_gsib_x | default(1, true) }}
+atmosphere_layout_gsib_y: {{ layout_gsib_y | default(1, true) }}
 
-background_error_file: "{{BERROR_YAML}}"
 
 # Forecasting
-forecast_length: "{{ ATM_WINDOW_LENGTH }}"
-forecast_timestep: "{{ BKG_TSTEP }}"
+atmosphere_forecast_timestep: "{{ BKG_TSTEP }}"
 
 # Write final increment on Guassian grid in variational
-final_increment_prefix: "./anl/atminc."
+atmosphere_final_increment_prefix: "./anl/atminc."
 
 
 # Observation things
@@ -77,28 +69,28 @@ observations: all_observations
 crtm_coefficient_path: "{{ DATA }}/crtm/"
 
 # Naming conventions for observational files
-atm_obsdatain_path: "{{atm_obsdatain_path}}"
-atm_obsdatain_prefix: "{{OPREFIX}}"
-atm_obsdatain_suffix: ".tm00.nc"
+atmosphere_obsdatain_path: "{{atm_obsdatain_path}}"
+atmosphere_obsdatain_prefix: "{{OPREFIX}}"
+atmosphere_obsdatain_suffix: ".tm00.nc"
 
-atm_obsdataout_path: "{{atm_obsdataout_path}}"
-atm_obsdataout_prefix: diag_
-atm_obsdataout_suffix: "_{{ current_cycle | to_YMDH }}.nc"
+atmosphere_obsdataout_path: "{{atm_obsdataout_path}}"
+atmosphere_obsdataout_prefix: diag_
+atmosphere_obsdataout_suffix: "_{{ current_cycle | to_YMDH }}.nc"
 
 # Naming conventions for bias correction files
-atm_obsbiasin_path: "{{DATA}}/obs/"
-atm_obsbiasin_prefix: "{{GPREFIX}}"
-atm_obsbiasin_suffix: ".satbias.nc"
-atm_obstlapsein_prefix: "{{GPREFIX}}"
-atm_obstlapsein_suffix: ".tlapse.txt"
-atm_obsbiascovin_prefix: "{{GPREFIX}}"
-atm_obsbiascovin_suffix: ".satbias_cov.nc"
+atmosphere_obsbiasin_path: "{{DATA}}/obs/"
+atmosphere_obsbiasin_prefix: "{{GPREFIX}}"
+atmosphere_obsbiasin_suffix: ".satbias.nc"
+atmosphere_obstlapsein_prefix: "{{GPREFIX}}"
+atmosphere_obstlapsein_suffix: ".tlapse.txt"
+atmosphere_obsbiascovin_prefix: "{{GPREFIX}}"
+atmosphere_obsbiascovin_suffix: ".satbias_cov.nc"
 
-atm_obsbiasout_path: "{{DATA}}/bc/"
-atm_obsbiasout_prefix: "{{APREFIX}}"
-atm_obsbiasout_suffix: ".satbias.nc"
-atm_obsbiascovout_prefix: "{{APREFIX}}"
-atm_obsbiascovout_suffix: ".satbias_cov.nc"
+atmosphere_obsbiasout_path: "{{DATA}}/bc/"
+atmosphere_obsbiasout_prefix: "{{APREFIX}}"
+atmosphere_obsbiasout_suffix: ".satbias.nc"
+atmosphere_obsbiascovout_prefix: "{{APREFIX}}"
+atmosphere_obsbiascovout_suffix: ".satbias_cov.nc"
 
 
 # Local Ensemble DA (LETKF)
@@ -124,5 +116,5 @@ driver_save_posterior_mean_increment: false
 driver_save_posterior_ensemble_increments: true
 
 # Diagnostics
-ensemble_increment_prefix: "./anl/mem%{member}%/atminc."
-posterior_output_gaussian: "./mem%{member}%/atmanl."
+atmosphere_ensemble_increment_prefix: "./anl/mem%{member}%/atminc."
+atmosphere_posterior_output_gaussian: "./mem%{member}%/atmanl."

parm/atm/jcb-prototype_lgetkf.yaml.j2

@@ -9,7 +9,7 @@ algorithm: local_ensemble_da
 
 # Model things
 # ------------
-background_ensemble_path: ./bkg/mem%mem%
+atmosphere_background_ensemble_path: ./bkg/mem%mem%
 
 # Observation things
 # ------------------

parm/soca/marine-jcb-3dfgat.yaml.j2

@@ -0,0 +1,5 @@
+# JEDI algorithm being invoked
+algorithm: 3dfgat
+
+# Observations
+observations: !INC ${OBS_LIST_SHORT}

parm/soca/marine-jcb-base.yaml

@@ -0,0 +1,72 @@
+# Search path for model and obs for JCB
+# -------------------------------------
+algorithm_path: '{{PARMgfs}}/gdas/jcb-algorithms'
+app_path_algorithm: '{{PARMgfs}}/gdas/jcb-gdas/algorithm/marine'
+app_path_model: '{{PARMgfs}}/gdas/jcb-gdas/model/marine'
+app_path_observations: '{{PARMgfs}}/gdas/jcb-gdas/observations/marine'
+app_path_observation_chronicle: '{{PARMgfs}}/gdas/jcb-gdas/observation_chronicle/marine'
+
+
+# Places where we deviate from the generic file name of a yaml
+# ------------------------------------------------------------
+model_file: marine_model_pseudo
+initial_condition_file: marine_background  # Initial conditions for 4D apps is background
+output_file: marine_output
+final_increment_file: marine_final_increment
+
+
+# Assimilation standard things (not prepended with model choice)
+# ----------------------------
+window_begin: '{{ATM_WINDOW_BEGIN}}'
+window_length: '{{ATM_WINDOW_LENGTH}}'
+bound_to_include: begin
+minimizer: RPCG
+final_diagnostics_departures: oman
+number_of_outer_loops: 1
+analysis_variables: [cicen, hicen, hsnon, socn, tocn, uocn, vocn, ssh]
+
+
+# Model things
+# ------------
+marine_window_begin: '{{ATM_WINDOW_BEGIN}}'
+marine_window_middle: '{{ATM_WINDOW_MIDDLE}}'
+
+# Geometry
+marine_soca_files_path: .
+
+# Background
+marine_background_path: ./INPUT/
+marine_background_time: '{{ATM_WINDOW_BEGIN}}'
+
+# Pseudo model
+marine_forecast_timestep: PT3H
+marine_pseudo_model_states: !INC ${BKG_LIST}
+
+# Background error model
+background_error_file: '{{berror_model}}'
+marine_number_ensemble_members: '{{nmem_ens}}'
+marine_stddev_time: '{{ATM_WINDOW_MIDDLE}}'
+
+# Observations
+observations: all_observations
+
+marine_obsdatain_path: '{{DATA}}/obs'
+marine_obsdatain_prefix: '{{OPREFIX}}'
+marine_obsdatain_suffix: '.{{PDY}}{{cyc}}.nc4'
+marine_obsdataout_path: '{{DATA}}/diags'
+marine_obsdataout_prefix: ''
+marine_obsdataout_suffix: '.{{PDY}}{{cyc}}.nc4'
+
+# Outer loops
+marine_ninner_1: '{{SOCA_NINNER}}'
+marine_grad_red_1: 1e-7
+
+# Output
+marine_output_dir: Data
+marine_output_exp: 3dvarfgat_pseudo
+marine_output_freq: PT3H   # can only writes out the analysis at the middle of the window
+                           # when using 3D-FGAT as cost function
+
+# Final increment
+marine_final_increment_dir: Data
+marine_final_increment_exp: 3dvarfgat_pseudo

scripts/exgdas_global_marine_analysis_prep.py

@@ -3,13 +3,18 @@
 # Script description:  Stages files and generates YAML for UFS Global Marine Analysis
 
 # import os to add ush to path
+import copy
 import os
 import glob
 import dateutil.parser as dparser
 import f90nml
 from soca import bkg_utils
 from datetime import datetime, timedelta
 import pytz
+import re
+import yaml
+
+from jcb import render
 from wxflow import (Logger, Template, TemplateConstants,
                     YAMLFile, FileHandler, AttrDict, parse_j2yaml)
 
@@ -59,9 +64,25 @@ def find_clim_ens(input_date):
     return nearest_date(dirs, input_date)
 
 
+def parse_obs_list_file():
+    # Get the list of observation types from the obs_list.yaml
+    obs_list_path = os.path.join(gdas_home, 'parm', 'soca', 'obs', 'obs_list.yaml')
+    obs_types = []
+    with open(obs_list_path, 'r') as file:
+        for line in file:
+            # Remove leading/trailing whitespace and check if the line is uncommented
+            line = line.strip()
+            if line.startswith('- !INC') and not line.startswith('#'):
+                # Extract the type using regex
+                match = re.search(r'\$\{OBS_YAML_DIR\}/(.+)\.yaml', line)
+                if match:
+                    obs_types.append(str(match.group(1)))
+    return obs_types
+
 ################################################################################
 # runtime environment variables, create directories
 
+
 logger.info(f"---------------- Setup runtime environement")
 
 comin_obs = os.getenv('COMIN_OBS')
@@ -298,6 +319,52 @@ def find_clim_ens(input_date):
 # Remove empty obs spaces in var_yaml
 ufsda.yamltools.save_check(varconfig, target=var_yaml, app='var')
 
+# Produce JEDI YAML file using JCB (for demonstration purposes)
+# -------------------------------------------------------------
+
+# Make a copy of the env config before modifying to avoid breaking something else
+envconfig_jcb = copy.deepcopy(envconfig)
+
+# Add the things to the envconfig in order to template JCB files
+envconfig_jcb['PARMgfs'] = os.getenv('PARMgfs')
+envconfig_jcb['nmem_ens'] = nmem_ens
+envconfig_jcb['berror_model'] = 'marine_background_error_static_diffusion'
+if nmem_ens > 3:
+    envconfig_jcb['berror_model'] = 'marine_background_error_hybrid_diffusion_diffusion'
+envconfig_jcb['DATA'] = os.getenv('DATA')
+envconfig_jcb['OPREFIX'] = os.getenv('OPREFIX')
+envconfig_jcb['PDY'] = os.getenv('PDY')
+envconfig_jcb['cyc'] = os.getenv('cyc')
+envconfig_jcb['SOCA_NINNER'] = os.getenv('SOCA_NINNER')
+envconfig_jcb['obs_list'] = ['adt_rads_all']
+
+# Write obs_list_short
+with open('obs_list_short.yaml', 'w') as file:
+    yaml.dump(parse_obs_list_file(), file, default_flow_style=False)
+os.environ['OBS_LIST_SHORT'] = 'obs_list_short.yaml'
+
+# Render the JCB configuration files
+jcb_base_yaml = os.path.join(gdas_home, 'parm', 'soca', 'marine-jcb-base.yaml')
+jcb_algo_yaml = os.path.join(gdas_home, 'parm', 'soca', 'marine-jcb-3dfgat.yaml.j2')
+
+jcb_base_config = YAMLFile(path=jcb_base_yaml)
+jcb_base_config = Template.substitute_structure(jcb_base_config, TemplateConstants.DOUBLE_CURLY_BRACES, envconfig_jcb.get)
+jcb_base_config = Template.substitute_structure(jcb_base_config, TemplateConstants.DOLLAR_PARENTHESES, envconfig_jcb.get)
+jcb_algo_config = YAMLFile(path=jcb_algo_yaml)
+jcb_algo_config = Template.substitute_structure(jcb_algo_config, TemplateConstants.DOUBLE_CURLY_BRACES, envconfig_jcb.get)
+jcb_algo_config = Template.substitute_structure(jcb_algo_config, TemplateConstants.DOLLAR_PARENTHESES, envconfig_jcb.get)
+
+# Override base with the application specific config
+jcb_config = {**jcb_base_config, **jcb_algo_config}
+
+# Render the full JEDI configuration file using JCB
+jedi_config = render(jcb_config)
+
+# Save the JEDI configuration file
+var_yaml_jcb = os.path.join(anl_dir, 'var-jcb.yaml')
+ufsda.yamltools.save_check(jedi_config, target=var_yaml_jcb, app='var')
+
+
 ################################################################################
 # Prepare the yamls for the "checkpoint" jjob
 # prepare yaml and CICE restart for soca to cice change of variable

test/atm/global-workflow/config.atmanl

@@ -15,10 +15,10 @@ export INTERP_METHOD='barycentric'
 if [[ ${DOHYBVAR} = "YES" ]]; then
     # shellcheck disable=SC2153
     export CASE_ANL=${CASE_ENS}
-    export BERROR_YAML="background_error_hybrid_${STATICB_TYPE}_${LOCALIZATION_TYPE}"
+    export BERROR_YAML="atmosphere_background_error_hybrid_${STATICB_TYPE}_${LOCALIZATION_TYPE}"
 else
     export CASE_ANL=${CASE}
-    export BERROR_YAML="background_error_static_${STATICB_TYPE}"
+    export BERROR_YAML="atmosphere_background_error_static_${STATICB_TYPE}"
 fi
 
 export CRTM_FIX_YAML="${PARMgfs}/gdas/atm_crtm_coeff.yaml.j2"

test/atm/global-workflow/jcb-prototype_3dvar.yaml.j2

@@ -4,7 +4,7 @@ algorithm: 3dvar
 
 # Model things
 # ------------
-background_ensemble_path: ./bkg/mem%mem%
+atmosphere_background_ensemble_path: ./bkg/mem%mem%
 
 # Observation things
 # ------------------
@@ -13,4 +13,4 @@ observations:
   - sondes
 
 # The observation files in the testing are appended using the yyymmddhh similar to JEDI tests
-atm_obsdatain_suffix: ".{{ current_cycle | to_YMDH }}.nc"
+atmosphere_obsdatain_suffix: ".{{ current_cycle | to_YMDH }}.nc"

test/atm/global-workflow/jcb-prototype_lgetkf.yaml.j2

@@ -9,7 +9,7 @@ algorithm: local_ensemble_da
 
 # Model things
 # ------------
-background_ensemble_path: ./bkg/mem%mem%
+atmosphere_background_ensemble_path: ./bkg/mem%mem%
 
 # Observation things
 # ------------------
@@ -18,4 +18,4 @@ observations:
   - sondes
 
 # The observation files in the testing are appended using the yyymmddhh similar to JEDI tests
-atm_obsdatain_suffix: ".{{ current_cycle | to_YMDH }}.nc"
+atmosphere_obsdatain_suffix: ".{{ current_cycle | to_YMDH }}.nc"