diff --git a/ush/python/pygfs/task/marine_letkf.py b/ush/python/pygfs/task/marine_letkf.py index 02f4e2a7a45..965bc18f241 100644 --- a/ush/python/pygfs/task/marine_letkf.py +++ b/ush/python/pygfs/task/marine_letkf.py @@ -1,18 +1,26 @@ #!/usr/bin/env python3 +import copy import f90nml -import pygfs.utils.marine_da_utils as mdau +from jcb import render from logging import getLogger import os +from pygfs.jedi import Jedi from pygfs.task.analysis import Analysis +import pygfs.utils.marine_da_utils as mdau from typing import Dict -from wxflow import (AttrDict, +from wxflow import (add_to_datetime, + AttrDict, Executable, FileHandler, logit, parse_j2yaml, + save_as_yaml, + Template, + TemplateConstants, to_timedelta, - to_YMDH) + to_YMDH, + YAMLFile) logger = getLogger(__name__.split('.')[-1]) @@ -37,6 +45,7 @@ def __init__(self, config: Dict) -> None: logger.info("init") super().__init__(config) + _enspert_relpath = os.path.relpath(self.task_config.DATAens, self.task_config.DATA) _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, @@ -44,16 +53,34 @@ def __init__(self, config: Dict) -> None: 'localensembleda', _letkf_yaml_file] # compute the relative path from self.task_config.DATA to self.task_config.DATAenspert - _enspert_relpath = os.path.relpath(self.task_config.DATAens, self.task_config.DATA) + _window_begin = add_to_datetime(self.task_config.current_cycle, -to_timedelta(f"{self.task_config.assim_freq}H") / 2) + _window_end = add_to_datetime(self.task_config.current_cycle, to_timedelta(f"{self.task_config.assim_freq}H") / 2) 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.MARINE_WINDOW_MIDDLE = self.task_config.current_cycle + self.task_config.WINDOW_BEGIN = _window_begin + self.task_config.WINDOW_END = _window_end + self.task_config.MARINE_WINDOW_BEGIN = _window_begin + self.task_config.MARINE_WINDOW_END = _window_end 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') self.task_config.ENSPERT_RELPATH = _enspert_relpath + self.task_config.PARMsoca = os.path.join(self.task_config.PARMgfs, 'gdas', 'soca') + + # Create dictionary of JEDI objects + expected_keys = ['marineanlletkf'] + self.jedi_dict = Jedi.get_jedi_dict(self.task_config.JEDI_CONFIG_YAML, self.task_config, expected_keys) + self.jedi_dict.PARMgfs = self.task_config.PARMgfs + #self.jedi_dict.MARINE_WINDOW_BEGIN = _window_begin.strftime('%Y-%m-%dT%H:%M:%SZ') + self.jedi_dict.MARINE_WINDOW_BEGIN = _window_begin + self.jedi_dict.MARINE_WINDOW_END = _window_end + #self.jedi_dict.MARINE_WINDOW_LENGTH = _window_end - _window_begin + self.jedi_dict.MARINE_WINDOW_LENGTH = f"PT{self.task_config['assim_freq']}H" + print("jedi_dict: ",self.jedi_dict) + @logit(logger) def initialize(self): @@ -66,43 +93,17 @@ def initialize(self): None """ + # NOTE: This task requires the ensemble background to be staged, which is done in the + # BMAT task (which is a dependency of this task) + logger.info("initialize") # make directories and stage ensemble background files soca_fix_stage_list = parse_j2yaml(self.task_config.SOCA_FIX_YAML_TMPL, self.task_config) FileHandler(soca_fix_stage_list).sync() - stageconf = AttrDict() - keys = ['current_cycle', - 'previous_cycle', - 'COM_ICE_LETKF_TMPL', - 'COM_OCEAN_LETKF_TMPL', - 'COM_ICE_HISTORY_TMPL', - 'COM_OCEAN_HISTORY_TMPL', - 'COMIN_OCEAN_HISTORY_PREV', - 'COMIN_ICE_HISTORY_PREV', - 'COMOUT_ICE_LETKF', - 'COMOUT_OCEAN_LETKF', - 'DATA', - 'ENSPERT_RELPATH', - 'GDUMP_ENS', - 'NMEM_ENS', - 'OPREFIX', - 'PARMgfs', - 'ROTDIR', - 'RUN', - 'WINDOW_BEGIN', - 'WINDOW_MIDDLE'] - for key in keys: - stageconf[key] = self.task_config[key] - - # stage ensemble background files - soca_ens_bkg_stage_list = parse_j2yaml(self.task_config.MARINE_ENSDA_STAGE_BKG_YAML_TMPL, stageconf) - FileHandler(soca_ens_bkg_stage_list).sync() - - # stage letkf-specific files - letkf_stage_list = parse_j2yaml(self.task_config.MARINE_LETKF_STAGE_YAML_TMPL, stageconf) + # stage letkf-specific files and directories + 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.MARINE_OBS_LIST_YAML, self.task_config) # get the list of observations @@ -127,13 +128,54 @@ def initialize(self): else: logger.warning(f"{obs_file} is not available in {self.task_config.COMIN_OBS}") + # set aside for LETKF configuration + observers = {'observers': obs_to_use} + # stage the desired obs files FileHandler({'copy': obs_files_to_copy}).sync() - # make the letkf.yaml - letkf_yaml = parse_j2yaml(self.task_config.MARINE_LETKF_YAML_TMPL, stageconf) - letkf_yaml.observations.observers = obs_to_use - letkf_yaml.save(self.task_config.letkf_yaml_file) +#################################################################################################### + + + # initialize JEDI LETKF observer application + #self.task_config['jcb_algo'] = 'marine_letkf' + logger.info(f"Initializing JEDI LETKF observer application") + #self.jedi_dict['marineanlletkf'].initialize(self.task_config) + self.jedi_dict['marineanlletkf'].initialize(self.jedi_dict) + + envconfig_jcb = copy.deepcopy(self.task_config) + envconfig_jcb['cyc'] = int(os.getenv('cyc')) + envconfig_jcb['PDY'] = self.task_config.current_cycle.strftime('%Y%m%d') + envconfig_jcb['window_length'] = f"PT{self.task_config['assim_freq']}H" + + # Render the JCB configuration files + jcb_base_yaml = os.path.join(self.task_config.PARMsoca, 'marine-jcb-base.yaml') + jcb_algo_yaml = os.path.join(self.task_config.PARMsoca, 'letkf/marine-jcb-lektf.yaml.j2') + + jcb_base_config = parse_j2yaml(path=jcb_base_yaml, data=envconfig_jcb) + 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 = parse_j2yaml(path=jcb_algo_yaml, data=envconfig_jcb) + 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} + + # convert datetime to string + jcb_config['window_begin'] = self.task_config.MARINE_WINDOW_BEGIN.strftime('%Y-%m-%dT%H:%M:%SZ') + jcb_config['window_middle'] = self.task_config.MARINE_WINDOW_MIDDLE.strftime('%Y-%m-%dT%H:%M:%SZ') + jcb_config['window_length'] = f"PT{self.task_config['assim_freq']}H" + + # Render the full JEDI configuration file using JCB + jedi_config = render(jcb_config) + jedi_config['observations'] = observers + + # Save the JEDI configuration file + letkf_yaml_jcb = 'letkf.yaml' + save_as_yaml(jedi_config, letkf_yaml_jcb) + + ###################################### # swap date and stack size in mom_input.nml domain_stack_size = self.task_config.DOMAIN_STACK_SIZE