Compress and SaveForecast with R2D2v3

pyproject.toml

@@ -18,7 +18,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "swell"
-version = "20260603"
+version = "20260728"
 authors = [{name = "NASA Global Modeling and Assimilation Office"}]
 description = "Workflow suites, tasks and configuration for coupled data assimilation"
 requires-python = ">=3.11"

src/swell/configuration/jedi/interfaces/geos_atmosphere/model/r2d2.py

@@ -31,6 +31,7 @@ def r2d2(template_dict: Mapping) -> Mapping:
         'store': {
             'fc': [
                 {'file_type': 'bkg',
+                 'filename': f'{cycle_dir}/bkg.%Y%m%dT%H%M%SZ.nc4',
                  'r2d2_model': 'geos'}
             ]
         }

src/swell/configuration/jedi/interfaces/geos_cf/model/r2d2.py

@@ -63,6 +63,7 @@ def r2d2(template_dict: Mapping) -> Mapping:
         'store': {
             'fc': [
                 {'file_type': 'bkg',
+                 'filename': f'{cycle_dir}/bkg.%Y%m%dT%H%M%SZ.nc4',
                  'r2d2_model': 'geos_cf'}
             ]
         }

src/swell/configuration/jedi/interfaces/geos_marine/model/r2d2.py

@@ -12,18 +12,20 @@
 
 
 def r2d2(template_dict: Mapping) -> Mapping:
+    analysis_time_iso = template_dict['analysis_time_iso']
     cycle_dir = template_dict['cycle_dir']
-    local_background_time = template_dict['local_background_time']
     experiment_id = template_dict['experiment_id']
-    analysis_time_iso = template_dict['analysis_time_iso']
+    local_background_time = template_dict['local_background_time']
 
     fc_list = [{'file_type': 'MOM.res',
                 'r2d2_model': 'mom6',
+                'file_extension': 'tar.gz',
                 'filename': f'{cycle_dir}/MOM6.res.{local_background_time}.nc'}]
 
     if 'cice6' in template_dict['marine_models']:
         fc_list.append({'file_type': 'cice.res',
                         'r2d2_model': 'cice6',
+                        'file_extension': 'tar.gz',
                         'filename': f'{cycle_dir}/cice.res.{local_background_time}.nc'})
 
     an_list = [{'file_type': 'ocn.incr',

src/swell/suites/3dfgat_marine_cycle/flow.cylc

@@ -106,14 +106,11 @@
             SaveRestart-{{model_component}} => MoveDaRestart-{{model_component}} => CleanCycle-{{model_component}}
 
             {% if not skip_r2d2 %}
-            # Save analysis output
-            # RunJediFgatExecutable-{{model_component}} => SaveAnalysis-{{model_component}}
+            # 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}}
             {% endif %}
 
-            # Save model output
-            # MoveBackground-{{model_component}} => StoreBackground-{{model_component}}
-
             # Clean up large files
             EvaObservations-{{model_component}} & EvaJediLog-{{model_component}} & EvaIncrement-{{model_component}} =>
             CleanCycle-{{model_component}}
@@ -246,6 +243,18 @@
             --{{key}} = {{value}}
         {%- endfor %}
 
+    [[RunCompressForecast-{{model_component}}]]
+        script = "swell task RunCompressForecast $config -d $datetime -m {{model_component}}"
+        platform = {{platform}}
+        execution time limit = {{scheduling["RunCompressForecast"]["execution_time_limit"]}}
+        [[[directives]]]
+        {%- for key, value in scheduling["RunCompressForecast"]["directives"][model_component].items() %}
+            --{{key}} = {{value}}
+        {%- endfor %}
+
+    [[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}}"
 

src/swell/suites/3dvar_marine/suite_config.py

@@ -37,7 +37,6 @@ class SuiteConfig(QuestionContainer, Enum):
             qd.horizontal_resolution("72x36"),
             qd.vertical_resolution("50"),
             qd.total_processors(6),
-            qd.obs_experiment("s2s_v1"),
             qd.observations([
                 "adt_cryosat2n",
                 "adt_jason3",
@@ -54,6 +53,7 @@ class SuiteConfig(QuestionContainer, Enum):
                 "temp_profile_xbt"
             ]),
             qd.background_time_offset("PT18H"),
+            qd.background_experiment("pr790_test"),
             qd.clean_patterns(['*.nc4', '*.txt']),
         ]
     )
@@ -68,3 +68,31 @@ class SuiteConfig(QuestionContainer, Enum):
     )
 
     # --------------------------------------------------------------------------------------------------
+
+    _3dvar_marine_tier2 = QuestionList(
+        list_name="3dvar_marine_tier2",
+        questions=[
+            _3dvar_marine,
+            qd.start_cycle_point("2023-01-02T12:00:00Z"),
+            qd.final_cycle_point("2023-01-03T12:00:00Z"),
+        ],
+        geos_marine=[
+            qd.cycle_times(['T12']),
+            qd.marine_models(['mom6', 'cice6']),
+            qd.analysis_variables([
+                "sea_water_salinity",
+                "sea_water_potential_temperature",
+                "sea_surface_height_above_geoid",
+                "sea_water_cell_thickness",
+                "sea_ice_area_fraction",
+                "sea_ice_thickness",
+                "sea_ice_snow_thickness"
+            ]),
+            qd.window_length("P1D"),
+            qd.horizontal_resolution("1440x1080"),
+            qd.vertical_resolution("75"),
+            qd.total_processors(720),
+        ]
+    )
+
+    # --------------------------------------------------------------------------------------------------

src/swell/suites/3dvar_marine_cycle/flow.cylc

@@ -105,14 +105,11 @@
             SaveRestart-{{model_component}} => MoveDaRestart-{{model_component}} => CleanCycle-{{model_component}}
 
             {% if not skip_r2d2 %}
-            # Save analysis output
-            # RunJediVariationalExecutable-{{model_component}} => SaveAnalysis-{{model_component}}
+            # 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}}
             {% endif %}
 
-            # Save model output
-            # MoveBackground-{{model_component}} => StoreBackground-{{model_component}}
-
             # Clean up large files
             EvaObservations-{{model_component}} & EvaJediLog-{{model_component}} & EvaIncrement-{{model_component}} =>
             CleanCycle-{{model_component}}
@@ -217,6 +214,15 @@
 
     {% endif %}
 
+    [[RunCompressForecast-{{model_component}}]]
+        script = "swell task RunCompressForecast $config -d $datetime -m {{model_component}}"
+        platform = {{platform}}
+        execution time limit = {{scheduling["RunCompressForecast"]["execution_time_limit"]}}
+        [[[directives]]]
+        {%- for key, value in scheduling["RunCompressForecast"]["directives"][model_component].items() %}
+            --{{key}} = {{value}}
+        {%- endfor %}
+
     [[RenderJediObservations-{{model_component}}]]
         script = "swell task RenderJediObservations $config -d $datetime -m {{model_component}}"
 
@@ -245,6 +251,9 @@
             --{{key}} = {{value}}
         {%- endfor %}
 
+    [[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}}"
 

src/swell/tasks/geos_marine/get_background_geos_marine.py

@@ -0,0 +1,111 @@
+# (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.
+
+
+# --------------------------------------------------------------------------------------------------
+
+import os
+import tarfile
+
+import r2d2
+
+from swell.tasks.base.task_base import taskBase
+from swell.utilities.r2d2 import load_r2d2_credentials
+
+# --------------------------------------------------------------------------------------------------
+
+
+def _marine_archive_path(cycle_dir: str, model_name: str, local_background_time: str) -> str:
+    """Return the canonical path for a marine background tar.gz archive.
+
+    Args:
+        cycle_dir: Directory where the cycle artifacts are written.
+        model_name: Name of the marine model, such as "mom6" or "cice6".
+        local_background_time: Background time string used in the archive name.
+
+    Returns:
+        The full path to the expected archive file.
+    """
+    return os.path.join(cycle_dir, f"{model_name}.{local_background_time}.tar.gz")
+
+
+# --------------------------------------------------------------------------------------------------
+
+
+class GetBackground(taskBase):
+    """Fetch marine forecast archives from R2D2 and unpack them into the cycle directory.
+
+    This task retrieves the pre-compressed marine background archives produced
+    by the marine forecast save task and extracts them into the current cycle
+    directory for downstream use.
+    """
+
+    def execute(self) -> None:
+        """Retrieve the archived marine forecast outputs for the current cycle.
+
+        The method loads the R2D2 credentials, resolves the active background
+        time, fetches the archive for each configured marine model, and unpacks
+        it into the cycle directory.
+        """
+
+        load_r2d2_credentials(
+            self.logger,
+            self.platform(),
+            r2d2_server=self.config.r2d2_server(default=None),
+        )
+
+        r2d2_datastore = self.config.r2d2_datastore(default=None)
+        marine_models = self.config.marine_models(None) or []
+        window_type = self.config.window_type()
+        window_length = self.config.window_length()
+        horizontal_resolution = self.config.horizontal_resolution()
+        background_experiment = self.config.background_experiment()
+
+        local_background_time, local_background_time_dto = \
+            self.da_window_params.local_background_time(window_length, window_type, dto=True)
+
+        os.makedirs(self.cycle_dir(), 0o755, exist_ok=True)
+
+        marine_model_configs = [('mom6', 'MOM.res')]
+        if 'cice6' in marine_models:
+            marine_model_configs.append(('cice6', 'cice.res'))
+
+        for model_name, file_type in marine_model_configs:
+            archive_path = _marine_archive_path(
+                self.cycle_dir(), model_name, local_background_time)
+
+            self.logger.info(
+                f'Fetching marine background archive {os.path.basename(archive_path)} '
+                f'({file_type}) for {model_name}'
+            )
+
+            fetch_kwargs = dict(
+                item='forecast',
+                target_file=archive_path,
+                model=model_name,
+                experiment=background_experiment,
+                file_extension='tar.gz',
+                resolution=horizontal_resolution,
+                step='PT00',
+                date=local_background_time_dto.strftime('%Y-%m-%d %H:%M:%S'),
+                file_type=file_type,
+            )
+            if r2d2_datastore:
+                fetch_kwargs['data_store'] = r2d2_datastore
+
+            r2d2.fetch(**fetch_kwargs)
+
+            if not os.path.exists(archive_path):
+                self.logger.abort(f'Marine archive not fetched: {archive_path}')
+
+            self.logger.info(f'Unpacking {os.path.basename(archive_path)} into {self.cycle_dir()}')
+            with tarfile.open(archive_path, 'r:gz') as tar:
+                tar.extractall(path=self.cycle_dir())
+
+            os.remove(archive_path)
+            self.logger.info(f'Unpacked {os.path.basename(archive_path)} into {self.cycle_dir()}')
+
+# --------------------------------------------------------------------------------------------------