diff --git a/src/swell/configuration/jedi/interfaces/geos_marine/model/r2d2.py b/src/swell/configuration/jedi/interfaces/geos_marine/model/r2d2.py index 53b104454..256089658 100644 --- a/src/swell/configuration/jedi/interfaces/geos_marine/model/r2d2.py +++ b/src/swell/configuration/jedi/interfaces/geos_marine/model/r2d2.py @@ -28,14 +28,15 @@ def r2d2(template_dict: Mapping) -> Mapping: 'file_extension': 'tar.gz', 'filename': f'{cycle_dir}/cice.res.{local_background_time}.nc'}) - an_list = [{'file_type': 'ocn.incr', + an_list = [{'file_type': 'ocn.an', 'r2d2_model': 'mom6', - 'filename': f'{cycle_dir}/mom6_increment.nc'}] + # 'filename': f'{cycle_dir}/mom6_increment.nc'}] + 'filename': f'{cycle_dir}/ocn.{experiment_id}.an.{analysis_time_iso}.nc'}] if 'cice6' in template_dict['marine_models']: - an_list.append({'file_type': 'ice.incr', + an_list.append({'file_type': 'ice.an', 'r2d2_model': 'cice6', - 'filename': f'{cycle_dir}/ice.{experiment_id}.incr.{analysis_time_iso}.nc'}) + 'filename': f'{cycle_dir}/ice.{experiment_id}.an.{analysis_time_iso}.nc'}) r2d2 = { 'fetch': { diff --git a/src/swell/suites/3dfgat_marine_cycle/flow.cylc b/src/swell/suites/3dfgat_marine_cycle/flow.cylc index a36d54a43..15be3f5fb 100644 --- a/src/swell/suites/3dfgat_marine_cycle/flow.cylc +++ b/src/swell/suites/3dfgat_marine_cycle/flow.cylc @@ -103,14 +103,17 @@ {% endif %} # Move restart to next cycle and then erase current forecast folder - SaveRestart-{{model_component}} => MoveDaRestart-{{model_component}} => CleanCycle-{{model_component}} + SaveRestart-{{model_component}} => MoveDaRestart-{{model_component}} d {% if not skip_r2d2 %} - # Compress forecast on compute node, then store in R2D2 on login node - LinkCoupledGeosOutput-{{model_component}} => RunCompressForecast-{{model_component}} => SaveForecast-{{model_component}} => MoveDaRestart-{{model_component}} - RunJediFgatExecutable-{{model_component}} => SaveObsDiags-{{model_component}} => CleanCycle-{{model_component}} + # Compress forecast on compute node, then save in R2D2 on login node; (state space) forecast, analysis, and (obs space) observation outputs. + LinkCoupledGeosOutput-{{model_component}} => SaveForecast-{{model_component}} => MoveDaRestart-{{model_component}} + RunJediFgatExecutable-{{model_component}} => SaveAnalysis-{{model_component}} => MoveDaRestart-{{model_component}} + RunJediFgatExecutable-{{model_component}} => SaveObsDiags-{{model_component}} => MoveDaRestart-{{model_component}} {% endif %} + MoveDaRestart-{{model_component}} => CleanCycle-{{model_component}} + # Clean up large files EvaObservations-{{model_component}} & EvaJediLog-{{model_component}} & EvaIncrement-{{model_component}} => CleanCycle-{{model_component}} @@ -243,6 +246,9 @@ --{{key}} = {{value}} {%- endfor %} + [[SaveAnalysis-{{model_component}}]] + script = "swell task SaveAnalysis $config -d $datetime -m {{model_component}}" + [[RunCompressForecast-{{model_component}}]] script = "swell task RunCompressForecast $config -d $datetime -m {{model_component}}" platform = {{platform}} diff --git a/src/swell/suites/3dvar_marine_cycle/flow.cylc b/src/swell/suites/3dvar_marine_cycle/flow.cylc index d7619156a..26b1c5532 100644 --- a/src/swell/suites/3dvar_marine_cycle/flow.cylc +++ b/src/swell/suites/3dvar_marine_cycle/flow.cylc @@ -105,9 +105,10 @@ SaveRestart-{{model_component}} => MoveDaRestart-{{model_component}} => CleanCycle-{{model_component}} {% if not skip_r2d2 %} - # Compress forecast on compute node, then store in R2D2 on login node - LinkCoupledGeosOutput-{{model_component}} => RunCompressForecast-{{model_component}} => SaveForecast-{{model_component}} => MoveDaRestart-{{model_component}} - RunJediVariationalExecutable-{{model_component}} => SaveObsDiags-{{model_component}} => CleanCycle-{{model_component}} + # Compress forecast on compute node, then save in R2D2 on login node; (state space) forecast, analysis, and (obs space) observation outputs. + LinkCoupledGeosOutput-{{model_component}} => SaveForecast-{{model_component}} => MoveDaRestart-{{model_component}} + RunJediVariationalExecutable-{{model_component}} => SaveAnalysis-{{model_component}} => MoveDaRestart-{{model_component}} + RunJediVariationalExecutable-{{model_component}} => SaveObsDiags-{{model_component}} => MoveDaRestart-{{model_component}} {% endif %} # Clean up large files @@ -251,15 +252,18 @@ --{{key}} = {{value}} {%- endfor %} + [[SaveAnalysis-{{model_component}}]] + script = "swell task SaveAnalysis $config -d $datetime -m {{model_component}}" + [[SaveForecast-{{model_component}}]] script = "swell task SaveForecast $config -d $datetime -m {{model_component}}" - [[SaveRestart-{{model_component}}]] - script = "swell task SaveRestart $config -d $datetime -m {{model_component}}" - [[SaveObsDiags-{{model_component}}]] script = "swell task SaveObsDiags $config -d $datetime -m {{model_component}}" + [[SaveRestart-{{model_component}}]] + script = "swell task SaveRestart $config -d $datetime -m {{model_component}}" + [[PrepareAnalysis-{{model_component}}]] script = "swell task PrepareAnalysis $config -d $datetime -m {{model_component}}" diff --git a/src/swell/tasks/get_background.py b/src/swell/tasks/get_background.py index afb92f1d1..850accc27 100644 --- a/src/swell/tasks/get_background.py +++ b/src/swell/tasks/get_background.py @@ -84,7 +84,7 @@ def execute(self) -> None: # Get window parameters local_background_time = self.da_window_params.local_background_time(window_length, window_type) - analysis_time_iso = self.da_window_params.analysis_time_iso() + analysis_time_iso, _ = self.da_window_params.analysis_time_iso() # Add to jedi config rendering dictionary self.jedi_rendering.add_key('local_background_time', local_background_time) diff --git a/src/swell/tasks/run_jedi_convert_state_soca2cice_executable.py b/src/swell/tasks/run_jedi_convert_state_soca2cice_executable.py index 297f0bfb3..e1d66a041 100644 --- a/src/swell/tasks/run_jedi_convert_state_soca2cice_executable.py +++ b/src/swell/tasks/run_jedi_convert_state_soca2cice_executable.py @@ -33,10 +33,10 @@ def execute(self) -> None: window_type = self.config.window_type() window_length = self.config.window_length() - # Compute data assimilation window parameters - # -------------------------------------------- - analysis_time = self.da_window_params.analysis_time() - analysis_time_iso = self.da_window_params.analysis_time_iso() + # Compute data assimilation window parameters, some methods return dto, which may not be needed here + # ----------------------------------------------------------------------------- + analysis_time, _ = self.da_window_params.analysis_time() + analysis_time_iso, _ = self.da_window_params.analysis_time_iso() local_background_time = self.da_window_params.local_background_time(window_length, window_type) local_background_time_iso = self.da_window_params.local_background_time_iso(window_length, diff --git a/src/swell/tasks/save_analysis.py b/src/swell/tasks/save_analysis.py new file mode 100644 index 000000000..d1c269dcb --- /dev/null +++ b/src/swell/tasks/save_analysis.py @@ -0,0 +1,185 @@ +# (C) Copyright 2021- United States Government as represented by the Administrator of the +# National Aeronautics and Space Administration. All Rights Reserved. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. + + +# -------------------------------------------------------------------------------------------------- + + +from datetime import datetime as dt +import isodate +import os +from r2d2 import store + +from swell.tasks.base.task_base import taskBase +from swell.utilities.datetime_util import datetime_formats +from swell.utilities.r2d2 import load_r2d2_credentials + + +# -------------------------------------------------------------------------------------------------- + +class SaveAnalysis(taskBase): + + # ---------------------------------------------------------------------------------------------- + + def execute(self) -> None: + + """Store analysis files for a given experiment and cycle in R2D2. + + Analysis files are handled differently then forecast files in R2D2, as they are stored + at their exact valid time (the analysis time). + + Additionally, for 4D/FGAT windows, all analyses within the window are stored, each at its own valid time, + rather than just the window-begin forecast as in the SaveForecast task. + + The marine FGAT only calculates one analysis and increment file per model (i.e., mom6 and cice6), + valid at the local background time. + + Parameters + ---------- + All inputs are extracted from the JEDI experiment file configuration. + See the taskBase constructor for more information. + """ + + # Parse common configuration as instance variables + # ------------------------------------------------- + self.marine_models = self.config.marine_models(None) or [] + self.window_type = self.config.window_type() + self.window_length = self.config.window_length() + self.window_begin_iso = self.da_window_params.window_begin_iso(self.window_length) + self.background_experiment = self.experiment_id() + self.horizontal_resolution = self.config.horizontal_resolution() + + # Load R2D2 credentials + # --------------------- + load_r2d2_credentials(self.logger, self.platform()) + + self.local_background_time, self.local_background_time_dto = self.da_window_params.local_background_time( + self.window_length, + self.window_type, + dto=True) + + analysis_time_iso, analysis_time_dto = self.da_window_params.analysis_time_iso() + + # Populate jedi_rendering template dictionary before rendering + # (mirrors run_jedi_variational_executable.py) + # ------------------------------------------------------------ + local_background_time_iso = self.da_window_params.local_background_time_iso( + self.window_length, self.window_type) + + self.jedi_rendering.add_key('local_background_time', self.local_background_time) + self.jedi_rendering.add_key('local_background_time_iso', local_background_time_iso) + self.jedi_rendering.add_key('marine_models', self.marine_models) + self.jedi_rendering.add_key('horizontal_resolution', self.horizontal_resolution) + self.jedi_rendering.add_key('analysis_time_iso', analysis_time_iso) + + # Render r2d2 interface dict once, shared by all store methods + # ------------------------------------------------------------- + self.r2d2_dict = self.jedi_rendering.render_interface_model('r2d2') + + # Dispatch to model- and window-type-specific store methods + # ---------------------------------------------------------- + model_component = self.get_model() + + if model_component == 'geos_atmosphere': + if is_4d: + self.store_atmosphere_4d() + else: + self.store_atmosphere_3d() + + elif model_component == 'geos_marine': + marine_model_configs = [('mom6', 'ocn_filename', 'ocn.an')] + if 'cice6' in self.marine_models: + marine_model_configs.append(('cice6', 'ice_filename', 'ice.an')) + + for model_name, filename_key, file_type in marine_model_configs: + self._store_an_dict(model_name, analysis_time_dto) + + else: + self.logger.abort(f'Unknown model component for SaveAnalysis: {model_component}') + + # ---------------------------------------------------------------------------------------------- + + def _store_analysis(self, + model_name: str, + ana_dto: dt, + source_file: str, + file_type: str, + ) -> None: + """Call r2d2.store for a single analysis file at its exact valid time. + + Parameters + ---------- + model_name : str + R2D2 model identifier (e.g. 'geos', 'mom6', 'cice6', 'geos_cf'). + ana_dto : datetime + Valid datetime of the analysis file (used directly as the r2d2 date). + source_file : str + Absolute path to the analysis file to store. + file_type : str + R2D2 file_type label (e.g. 'MOM.res', 'cice.res', 'bkg'). + """ + + file_extension = file_type.split('.')[-1] if '.' in file_type else 'nc' + + self.logger.info(f'Storing {os.path.basename(source_file)} ' + f'at {ana_dto.strftime(datetime_formats["iso_format"])}') + + # store( + # item='analysis', + # model=model_name, + # experiment=self.background_experiment, + # resolution=self.horizontal_resolution, + # date=ana_dto.strftime('%Y-%m-%d %H:%M:%S'), + # source_file=source_file, + # file_type=file_type, + # file_extension=file_extension, + # ) + + # ---------------------------------------------------------------------------------------------- + + def _store_an_dict(self, model_name: str, ana_dto: dt) -> None: + """Store all analyses defined in r2d2_dict['store']['an'] for one datetime. + + The filename for each entry is resolved by applying strftime to ana_dto, + ensuring it works for both static (marine, already strftime-compatible) and + datetime-templated (atmosphere) filename patterns. + + Parameters + ---------- + model_name : str + R2D2 model identifier. + ana_dto : datetime + Valid datetime of the analysis — used as the r2d2 date and for filename resolution. + """ + + for an in self.r2d2_dict['store']['an']: + if an.get('r2d2_model') != model_name: + continue + file_type = an['file_type'] + source_file = ana_dto.strftime(an['filename']) + self._store_analysis(model_name, ana_dto, source_file, file_type) + + # ---------------------------------------------------------------------------------------------- + # Atmosphere store methods (not implemeted yet, will need to resolve datetime-templated filenames for each analysis within the window) + # ---------------------------------------------------------------------------------------------- + + def store_atmosphere_3d(self) -> None: + """Store a single atmospheric analysis at the middle of a 3D window.""" + + self.logger.abort('Storing the window-begin atmospheric analysis is not ready yet.') + + # ---------------------------------------------------------------------------------------------- + + def store_atmosphere_4d(self) -> None: + """Store atmospheric analyses across a 4D (or FGAT) window. + + Uses states_generator to enumerate valid analysis datetimes and stores + each file, keyed by its exact valid time. + """ + + self.logger.abort('Storing the window-begin atmospheric analysis is not ready yet.') + +# -------------------------------------------------------------------------------------------------- diff --git a/src/swell/tasks/task_questions.py b/src/swell/tasks/task_questions.py index 71fca1608..b858e3a82 100644 --- a/src/swell/tasks/task_questions.py +++ b/src/swell/tasks/task_questions.py @@ -664,6 +664,18 @@ class TaskQuestions(QuestionContainer, Enum): # -------------------------------------------------------------------------------------------------- + RunCompressForecast = QuestionList( + list_name="RunCompressForecast", + questions=[ + window_questions, + qd.r2d2_experiment_id(), + qd.background_frequency(), + qd.horizontal_resolution(), + ] + ) + + # -------------------------------------------------------------------------------------------------- + RunJediConvertStateSoca2ciceExecutable = QuestionList( list_name="RunJediConvertStateSoca2ciceExecutable", questions=[ @@ -851,18 +863,6 @@ class TaskQuestions(QuestionContainer, Enum): # -------------------------------------------------------------------------------------------------- - RunCompressForecast = QuestionList( - list_name="RunCompressForecast", - questions=[ - window_questions, - qd.r2d2_experiment_id(), - qd.background_frequency(), - qd.horizontal_resolution(), - ] - ) - - # -------------------------------------------------------------------------------------------------- - RunJediEdaExecutable = QuestionList( list_name="RunJediEdaExecutable", questions=[ @@ -877,6 +877,18 @@ class TaskQuestions(QuestionContainer, Enum): # -------------------------------------------------------------------------------------------------- + SaveAnalysis = QuestionList( + list_name="SaveAnalysis", + questions=[ + window_questions, + qd.r2d2_experiment_id(), + qd.background_experiment(), + qd.horizontal_resolution(), + ] + ) + + # -------------------------------------------------------------------------------------------------- + SaveForecast = QuestionList( list_name="SaveForecast", questions=[ @@ -887,6 +899,7 @@ class TaskQuestions(QuestionContainer, Enum): ] ) + # -------------------------------------------------------------------------------------------------- SaveForecastCf = QuestionList( diff --git a/src/swell/utilities/data_assimilation_window_params.py b/src/swell/utilities/data_assimilation_window_params.py index 454cbefd2..79710e364 100644 --- a/src/swell/utilities/data_assimilation_window_params.py +++ b/src/swell/utilities/data_assimilation_window_params.py @@ -101,17 +101,31 @@ def __get_analysis_time__(self) -> datetime.datetime: # ---------------------------------------------------------------------------------------------- - def analysis_time(self) -> str: + def analysis_time(self) -> Tuple[str, datetime.datetime]: + """Get the analysis time for the current cycle. + + Returns + ------- + Tuple[str, datetime.datetime] + A tuple containing the analysis time formatted as a string and the corresponding datetime object. + """ analysis_time_dto = self.__get_analysis_time__() - return analysis_time_dto.strftime(datetime_formats['directory_format']) + return analysis_time_dto.strftime(datetime_formats['directory_format']), analysis_time_dto # ---------------------------------------------------------------------------------------------- - def analysis_time_iso(self) -> str: + def analysis_time_iso(self) -> Tuple[str, datetime.datetime]: + """Get the analysis time for the current cycle in ISO format. + + Returns + ------- + Tuple[str, datetime.datetime] + A tuple containing the analysis time formatted as an ISO string and the corresponding datetime object. + """ analysis_time_dto = self.__get_analysis_time__() - return analysis_time_dto.strftime(datetime_formats['iso_format']) + return analysis_time_dto.strftime(datetime_formats['iso_format']), analysis_time_dto # ----------------------------------------------------------------------------------------------