Adds contents of constructor and initialize methods to marine LETKF c…

…lass (#2635) Adds contents of constructor and initialize methods to marine LETKF class Partially addresses NOAA-EMC/GDASApp#1091 --------- Co-authored-by: Rahul Mahajan <[email protected]> Co-authored-by: Cory Martin <[email protected]>
NOAA-EMC · Jul 11, 2024 · 5ef4db7 · 5ef4db7
1 parent 8998ec7
commit 5ef4db7
Show file tree

Hide file tree

Showing 9 changed files with 126 additions and 42 deletions.
diff --git a/env/HERA.env b/env/HERA.env
@@ -140,13 +140,13 @@ elif [[ "${step}" = "ocnanalecen" ]]; then
     [[ ${NTHREADS_OCNANALECEN} -gt ${nth_max} ]] && export NTHREADS_OCNANALECEN=${nth_max}
     export APRUN_OCNANALECEN="${launcher} -n ${npe_ocnanalecen} --cpus-per-task=${NTHREADS_OCNANALECEN}"
 
-elif [[ "${step}" = "ocnanalletkf" ]]; then
+elif [[ "${step}" = "marineanalletkf" ]]; then
 
-    nth_max=$((npe_node_max / npe_node_ocnanalletkf))
+    nth_max=$((npe_node_max / npe_node_marineanalletkf))
 
-    export NTHREADS_OCNANALLETKF=${nth_ocnanalletkf:-${nth_max}}
-    [[ ${NTHREADS_OCNANALLETKF} -gt ${nth_max} ]] && export NTHREADS_OCNANALLETKF=${nth_max}
-    export APRUN_OCNANALLETKF="${launcher} -n ${npe_ocnanalletkf} --cpus-per-task=${NTHREADS_OCNANALLETKF}"
+    export NTHREADS_MARINEANALLETKF=${nth_marineanalletkf:-${nth_max}}
+    [[ ${NTHREADS_MARINEANALLETKF} -gt ${nth_max} ]] && export NTHREADS_MARINEANALLETKF=${nth_max}
+    export APRUN_MARINEANALLETKF="${launcher} -n ${npe_marineanalletkf} --cpus-per-task=${NTHREADS_MARINEANALLETKF}"
 
 elif [[ "${step}" = "anal" ]] || [[ "${step}" = "analcalc" ]]; then
 

diff --git a/env/ORION.env b/env/ORION.env
@@ -148,13 +148,13 @@ elif [[ "${step}" = "ocnanalecen" ]]; then
     [[ ${NTHREADS_OCNANALECEN} -gt ${nth_max} ]] && export NTHREADS_OCNANALECEN=${nth_max}
     export APRUN_OCNANALECEN="${launcher} -n ${npe_ocnanalecen} --cpus-per-task=${NTHREADS_OCNANALECEN}"
 
-elif [[ "${step}" = "ocnanalletkf" ]]; then
+elif [[ "${step}" = "marineanalletkf" ]]; then
 
-    nth_max=$((npe_node_max / npe_node_ocnanalletkf))
+    nth_max=$((npe_node_max / npe_node_marineanalletkf))
 
-    export NTHREADS_OCNANALLETKF=${nth_ocnanalletkf:-${nth_max}}
-    [[ ${NTHREADS_OCNANALLETKF} -gt ${nth_max} ]] && export NTHREADS_OCNANALLETKF=${nth_max}
-    export APRUN_OCNANALLETKF="${launcher} -n ${npe_ocnanalletkf} --cpus-per-task=${NTHREADS_OCNANALLETKF}"
+    export NTHREADS_MARINEANALLETKF=${nth_marineanalletkf:-${nth_max}}
+    [[ ${NTHREADS_MARINEANALLETKF} -gt ${nth_max} ]] && export NTHREADS_MARINEANALLETKF=${nth_max}
+    export APRUN_MARINEANALLETKF="${launcher} -n ${npe_marineanalletkf} --cpus-per-task=${NTHREADS_MARINEANALLETKF}"
 
 elif [[ "${step}" = "anal" ]] || [[ "${step}" = "analcalc" ]]; then
 

diff --git a/jobs/JGDAS_GLOBAL_OCEAN_ANALYSIS_LETKF → jobs/JGLOBAL_MARINE_ANALYSIS_LETKF b/jobs/JGDAS_GLOBAL_OCEAN_ANALYSIS_LETKF → jobs/JGLOBAL_MARINE_ANALYSIS_LETKF
@@ -1,6 +1,6 @@
 #!/bin/bash
 source "${HOMEgfs}/ush/preamble.sh"
-source "${HOMEgfs}/ush/jjob_header.sh" -e "ocnanalletkf" -c "base ocnanal ocnanalletkf"
+source "${HOMEgfs}/ush/jjob_header.sh" -e "marineanalletkf" -c "base ocnanal marineanalletkf"
 
 ##############################################
 # Set variables used in the script
@@ -13,8 +13,10 @@ gPDY=${GDATE:0:8}
 gcyc=${GDATE:8:2}
 
 YMD=${gPDY} HH=${gcyc} declare_from_tmpl -rx \
-   COM_OCEAN_HISTORY_PREV:COM_OCEAN_HISTORY_TMPL \
-   COM_ICE_HISTORY_PREV:COM_ICE_HISTORY_TMPL
+   COMIN_OCEAN_HISTORY_PREV:COM_OCEAN_HISTORY_TMPL \
+   COMIN_ICE_HISTORY_PREV:COM_ICE_HISTORY_TMPL
+
+YMD=${PDY} HH=${cyc} declare_from_tmpl -rx COMIN_OBS:COM_OBS_TMPL 
 
 ##############################################
 # Begin JOB SPECIFIC work

diff --git a/jobs/rocoto/ocnanalletkf.sh → jobs/rocoto/marineanalletkf.sh b/jobs/rocoto/ocnanalletkf.sh → jobs/rocoto/marineanalletkf.sh
@@ -8,7 +8,7 @@ source "${HOMEgfs}/ush/preamble.sh"
 status=$?
 [[ ${status} -ne 0 ]] && exit "${status}"
 
-export job="ocnanalletkf"
+export job="marineanalletkf"
 export jobid="${job}.$$"
 
 ###############################################################
@@ -18,6 +18,6 @@ export PYTHONPATH
 
 ###############################################################
 # Execute the JJOB
-"${HOMEgfs}/jobs/JGDAS_GLOBAL_OCEAN_ANALYSIS_LETKF"
+"${HOMEgfs}/jobs/JGLOBAL_MARINE_ANALYSIS_LETKF"
 status=$?
 exit "${status}"
diff --git a/parm/config/gfs/config.marineanalletkf b/parm/config/gfs/config.marineanalletkf
@@ -0,0 +1,18 @@
+#!/bin/bash
+
+########## config.marineanalletkf ##########
+# Ocn Analysis specific
+
+echo "BEGIN: config.marineanalletkf"
+
+# Get task specific resources
+. "${EXPDIR}/config.resources" marineanalletkf
+
+export MARINE_LETKF_EXEC="${JEDI_BIN}/gdas.x"
+export MARINE_LETKF_YAML_TMPL="${PARMgfs}/gdas/soca/letkf/letkf.yaml.j2"
+export MARINE_LETKF_STAGE_YAML_TMPL="${PARMgfs}/gdas/soca/letkf/letkf_stage.yaml.j2"
+
+export GRIDGEN_EXEC="${JEDI_BIN}/gdas_soca_gridgen.x"
+export GRIDGEN_YAML="${PARMgfs}/gdas/soca/gridgen/gridgen.yaml"
+
+echo "END: config.marineanalletkf"
diff --git a/parm/config/gfs/config.ocnanal b/parm/config/gfs/config.ocnanal
@@ -16,8 +16,8 @@ export SOCA_NINNER=@SOCA_NINNER@
 export CASE_ANL=@CASE_ANL@
 export DOMAIN_STACK_SIZE=116640000  #TODO: Make the stack size resolution dependent
 export JEDI_BIN=${HOMEgfs}/sorc/gdas.cd/build/bin
-
-export COMIN_OBS=@COMIN_OBS@
+export SOCA_FIX_STAGE_YAML_TMPL="${PARMgfs}/gdas/soca/soca_fix_stage.yaml.j2"
+export SOCA_ENS_BKG_STAGE_YAML_TMPL="${PARMgfs}/gdas/soca/soca_ens_bkg_stage.yaml.j2"
 
 # NICAS
 export NICAS_RESOL=@NICAS_RESOL@

diff --git a/parm/config/gfs/config.ocnanalletkf b/parm/config/gfs/config.ocnanalletkf
diff --git a/parm/config/gfs/config.resources b/parm/config/gfs/config.resources
@@ -25,7 +25,7 @@ if (( $# != 1 )); then
     echo "waveinit waveprep wavepostsbs wavepostbndpnt wavepostbndpntbll wavepostpnt"
     echo "wavegempak waveawipsbulls waveawipsgridded"
     echo "postsnd awips gempak npoess"
-    echo "ocnanalprep prepoceanobs ocnanalbmat ocnanalrun ocnanalecen ocnanalletkf ocnanalchkpt ocnanalpost ocnanalvrfy"
+    echo "ocnanalprep prepoceanobs ocnanalbmat ocnanalrun ocnanalecen marineanalletkf ocnanalchkpt ocnanalpost ocnanalvrfy"
     exit 1
 
 fi
@@ -557,32 +557,32 @@ case ${step} in
     export memory_ocnanalecen
     ;;
 
-  "ocnanalletkf")
+  "marineanalletkf")
     npes=16
     case ${OCNRES} in
       "025")
         npes=480
-        memory_ocnanalletkf="96GB"
+        memory_marineanalletkf="96GB"
         ;;
       "050")
         npes=16
-        memory_ocnanalletkf="96GB"
+        memory_marineanalletkf="96GB"
         ;;
       "500")
         npes=16
-        memory_ocnanalletkf="24GB"
+        memory_marineanalletkf="24GB"
         ;;
       *)
         echo "FATAL ERROR: Resources not defined for job ${step} at resolution ${OCNRES}"
         exit 4
     esac
 
-    export wtime_ocnanalletkf="00:10:00"
-    export npe_ocnanalletkf=${npes}
-    export nth_ocnanalletkf=1
+    export wtime_marineanalletkf="00:10:00"
+    export npe_marineanalletkf=${npes}
+    export nth_marineanalletkf=1
     export is_exclusive=True
-    export npe_node_ocnanalletkf=$(( npe_node_max / nth_ocnanalletkf ))
-    export memory_ocnanalletkf
+    export npe_node_marineanalletkf=$(( npe_node_max / nth_marineanalletkf ))
+    export memory_marineanalletkf
     ;;
 
 

diff --git a/ush/python/pygfs/task/marine_letkf.py b/ush/python/pygfs/task/marine_letkf.py
@@ -1,11 +1,16 @@
 #!/usr/bin/env python3
 
+import f90nml
 from logging import getLogger
+import os
 from pygfs.task.analysis import Analysis
 from typing import Dict
-from wxflow import (chdir,
+from wxflow import (AttrDict,
+                    FileHandler,
                     logit,
-                    Task)
+                    parse_j2yaml,
+                    to_timedelta,
+                    to_YMDH)
 
 logger = getLogger(__name__.split('.')[-1])
 
@@ -30,6 +35,21 @@ def __init__(self, config: Dict) -> None:
         logger.info("init")
         super().__init__(config)
 
+        _half_assim_freq = to_timedelta(f"{self.task_config.assim_freq}H") / 2
+        _letkf_yaml_file = 'letkf.yaml'
+        _letkf_exec_args = [self.task_config.MARINE_LETKF_EXEC,
+                            'soca',
+                            'localensembleda',
+                            _letkf_yaml_file]
+
+        self.task_config.WINDOW_MIDDLE = self.task_config.current_cycle
+        self.task_config.WINDOW_BEGIN = self.task_config.current_cycle - _half_assim_freq
+        self.task_config.letkf_exec_args = _letkf_exec_args
+        self.task_config.letkf_yaml_file = _letkf_yaml_file
+        self.task_config.mom_input_nml_tmpl = os.path.join(self.task_config.DATA, 'mom_input.nml.tmpl')
+        self.task_config.mom_input_nml = os.path.join(self.task_config.DATA, 'mom_input.nml')
+        self.task_config.obs_dir = os.path.join(self.task_config.DATA, 'obs')
+
     @logit(logger)
     def initialize(self):
         """Method initialize for ocean and sea ice LETKF task
@@ -43,6 +63,63 @@ def initialize(self):
 
         logger.info("initialize")
 
+        # make directories and stage ensemble background files
+        ensbkgconf = AttrDict()
+        keys = ['previous_cycle', 'current_cycle', 'DATA', 'NMEM_ENS',
+                'PARMgfs', 'ROTDIR', 'COM_OCEAN_HISTORY_TMPL', 'COM_ICE_HISTORY_TMPL']
+        for key in keys:
+            ensbkgconf[key] = self.task_config[key]
+        ensbkgconf.RUN = 'enkfgdas'
+        soca_ens_bkg_stage_list = parse_j2yaml(self.task_config.SOCA_ENS_BKG_STAGE_YAML_TMPL, ensbkgconf)
+        FileHandler(soca_ens_bkg_stage_list).sync()
+        soca_fix_stage_list = parse_j2yaml(self.task_config.SOCA_FIX_STAGE_YAML_TMPL, self.task_config)
+        FileHandler(soca_fix_stage_list).sync()
+        letkf_stage_list = parse_j2yaml(self.task_config.MARINE_LETKF_STAGE_YAML_TMPL, self.task_config)
+        FileHandler(letkf_stage_list).sync()
+
+        obs_list = parse_j2yaml(self.task_config.OBS_YAML, self.task_config)
+
+        # get the list of observations
+        obs_files = []
+        for ob in obs_list['observers']:
+            obs_name = ob['obs space']['name'].lower()
+            obs_filename = f"{self.task_config.RUN}.t{self.task_config.cyc}z.{obs_name}.{to_YMDH(self.task_config.current_cycle)}.nc"
+            obs_files.append((obs_filename, ob))
+
+        obs_files_to_copy = []
+        obs_to_use = []
+        # copy obs from COMIN_OBS to DATA/obs
+        for obs_file, ob in obs_files:
+            obs_src = os.path.join(self.task_config.COMIN_OBS, obs_file)
+            obs_dst = os.path.join(self.task_config.DATA, self.task_config.obs_dir, obs_file)
+            if os.path.exists(obs_src):
+                obs_files_to_copy.append([obs_src, obs_dst])
+                obs_to_use.append(ob)
+            else:
+                logger.warning(f"{obs_file} is not available in {self.task_config.COMIN_OBS}")
+
+        # stage the desired obs files
+        FileHandler({'copy': obs_files_to_copy}).sync()
+
+        # make the letkf.yaml
+        letkfconf = AttrDict()
+        keys = ['WINDOW_BEGIN', 'WINDOW_MIDDLE', 'RUN', 'gcyc', 'NMEM_ENS']
+        for key in keys:
+            letkfconf[key] = self.task_config[key]
+        letkfconf.RUN = 'enkfgdas'
+        letkf_yaml = parse_j2yaml(self.task_config.MARINE_LETKF_YAML_TMPL, letkfconf)
+        letkf_yaml.observations.observers = obs_to_use
+        letkf_yaml.save(self.task_config.letkf_yaml_file)
+
+        # swap date and stack size in mom_input.nml
+        domain_stack_size = self.task_config.DOMAIN_STACK_SIZE
+        ymdhms = [int(s) for s in self.task_config.WINDOW_BEGIN.strftime('%Y,%m,%d,%H,%M,%S').split(',')]
+        with open(self.task_config.mom_input_nml_tmpl, 'r') as nml_file:
+            nml = f90nml.read(nml_file)
+            nml['ocean_solo_nml']['date_init'] = ymdhms
+            nml['fms_nml']['domains_stack_size'] = int(domain_stack_size)
+            nml.write(self.task_config.mom_input_nml, force=True)  # force to overwrite if necessary
+
     @logit(logger)
     def run(self):
         """Method run for ocean and sea ice LETKF task
@@ -56,8 +133,6 @@ def run(self):
 
         logger.info("run")
 
-        chdir(self.runtime_config.DATA)
-
     @logit(logger)
     def finalize(self):
         """Method finalize for ocean and sea ice LETKF task