Merge pull request #398 from xylar/fix-omega-pr-baseline-comparison

Translate variable names for Omega before baseline validation

Author: xylar

docs/developers_guide/ocean/api.md

@@ -7,6 +7,12 @@
    :toctree: generated/
 
    Ocean
+   Ocean.map_to_native_model_vars
+   Ocean.map_var_list_to_native_model
+   Ocean.write_model_dataset
+   Ocean.map_from_native_model_vars
+   Ocean.map_var_list_from_native_model
+   Ocean.open_model_dataset
 
    add_tasks.add_ocean_tasks
 ```
@@ -314,10 +320,10 @@
 
    rpe.analysis.Analysis
    rpe.analysis.Analysis.run
-   
+
    viz.Viz
    viz.Viz.run
-   
+
 ```
 
 ### single_column

polaris/ocean/model/ocean_io_step.py

@@ -1,68 +1,17 @@
-import importlib.resources as imp_res
-from typing import Dict, Union
-
-import xarray as xr
-from mpas_tools.io import write_netcdf
-from ruamel.yaml import YAML
-
 from polaris import Step
+from polaris.tasks.ocean import Ocean
 
 
 class OceanIOStep(Step):
     """
     A step that writes input and/or output files for Omega or MPAS-Ocean
-
-    Attributes
-    ----------
-    mpaso_to_omega_dim_map : dict
-        A map from MPAS-Ocean dimension names to their Omega equivalents
-
-    omega_to_mpaso_dim_map : dict
-        A map from Omega dimension names to their MPAS-Ocean equivalents, the
-        inverse of ``mpaso_to_omega_dim_map``
-
-    mpaso_to_omega_var_map : dict
-        A map from MPAS-Ocean variable names to their Omega equivalents
-
-    omega_to_mpaso_var_map : dict
-        A map from Omega variable names to their MPAS-Ocean equivalents, the
-        inverse of ``mpaso_to_omega_var_map``
     """
 
-    def __init__(self, component, name, **kwargs):
-        """
-        Create a new step
-
-        Parameters
-        ----------
-        component : polaris.Component
-            The component the step belongs to
-
-        name : str
-            the name of the task
-
-        kwargs
-            keyword arguments passed to `polaris.Step()`
-        """
-        super().__init__(component=component, name=name, **kwargs)
-
-        self.mpaso_to_omega_dim_map: Union[None, Dict[str, str]] = None
-        self.omega_to_mpaso_dim_map: Union[None, Dict[str, str]] = None
-        self.mpaso_to_omega_var_map: Union[None, Dict[str, str]] = None
-        self.omega_to_mpaso_var_map: Union[None, Dict[str, str]] = None
+    # make sure component is of type Ocean
+    component: Ocean
 
-    def setup(self):
-        """
-        If the ocean model is Omega, set up maps between Omega and MPAS-Ocean
-        variable names
-        """
-        config = self.config
-        model = config.get('ocean', 'model')
-        if model == 'omega':
-            self._read_var_map()
-        elif model != 'mpas-ocean':
-            raise ValueError(f'Unexpected ocean model: {model}')
-        super().setup()
+    def __init__(self, component: Ocean, **kwargs):
+        super().__init__(component=component, **kwargs)
 
     def map_to_native_model_vars(self, ds):
         """
@@ -81,22 +30,7 @@ def map_to_native_model_vars(self, ds):
             The same dataset with variables renamed as appropriate for the
             ocean model being run
         """
-        config = self.config
-        model = config.get('ocean', 'model')
-        if model == 'omega':
-            assert self.mpaso_to_omega_dim_map is not None
-            rename = {
-                k: v
-                for k, v in self.mpaso_to_omega_dim_map.items()
-                if k in ds.dims
-            }
-            assert self.mpaso_to_omega_var_map is not None
-            rename_vars = {
-                k: v for k, v in self.mpaso_to_omega_var_map.items() if k in ds
-            }
-            rename.update(rename_vars)
-            ds = ds.rename(rename)
-        return ds
+        return self.component.map_to_native_model_vars(ds)
 
     def write_model_dataset(self, ds, filename):
         """
@@ -111,8 +45,7 @@ def write_model_dataset(self, ds, filename):
         filename : str
             The path for the NetCDF file to write
         """
-        ds = self.map_to_native_model_vars(ds)
-        write_netcdf(ds=ds, fileName=filename)
+        self.component.write_model_dataset(ds, filename)
 
     def map_from_native_model_vars(self, ds):
         """
@@ -130,22 +63,7 @@ def map_from_native_model_vars(self, ds):
         ds : xarray.Dataset
             The same dataset with variables named as expected in MPAS-Ocean
         """
-        config = self.config
-        model = config.get('ocean', 'model')
-        if model == 'omega':
-            assert self.omega_to_mpaso_dim_map is not None
-            rename = {
-                k: v
-                for k, v in self.omega_to_mpaso_dim_map.items()
-                if k in ds.dims
-            }
-            assert self.omega_to_mpaso_var_map is not None
-            rename_vars = {
-                k: v for k, v in self.omega_to_mpaso_var_map.items() if k in ds
-            }
-            rename.update(rename_vars)
-            ds = ds.rename(rename)
-        return ds
+        return self.component.map_from_native_model_vars(ds)
 
     def open_model_dataset(self, filename, **kwargs):
         """
@@ -165,27 +83,4 @@ def open_model_dataset(self, filename, **kwargs):
         ds : xarray.Dataset
             The dataset with variables named as expected in MPAS-Ocean
         """
-        ds = xr.open_dataset(filename, **kwargs)
-        ds = self.map_from_native_model_vars(ds)
-        return ds
-
-    def _read_var_map(self):
-        """
-        Read the map from MPAS-Ocean to Omega dimension and variable names
-        """
-        package = 'polaris.ocean.model'
-        filename = 'mpaso_to_omega.yaml'
-        text = imp_res.files(package).joinpath(filename).read_text()
-
-        yaml_data = YAML(typ='rt')
-        nested_dict = yaml_data.load(text)
-        self.mpaso_to_omega_dim_map = nested_dict['dimensions']
-        self.mpaso_to_omega_var_map = nested_dict['variables']
-        assert self.mpaso_to_omega_dim_map is not None
-        self.omega_to_mpaso_dim_map = {
-            v: k for k, v in self.mpaso_to_omega_dim_map.items()
-        }
-        assert self.mpaso_to_omega_var_map is not None
-        self.omega_to_mpaso_var_map = {
-            v: k for k, v in self.mpaso_to_omega_var_map.items()
-        }
+        return self.component.open_model_dataset(filename, **kwargs)

polaris/ocean/model/ocean_model_step.py

@@ -4,6 +4,7 @@
 from ruamel.yaml import YAML
 
 from polaris.model_step import ModelStep
+from polaris.tasks.ocean import Ocean
 
 
 class OceanModelStep(ModelStep):
@@ -26,9 +27,12 @@ class OceanModelStep(ModelStep):
         working directory)
     """
 
+    # make sure component is of type Ocean
+    component: Ocean
+
     def __init__(
         self,
-        component,
+        component: Ocean,
         name,
         subdir=None,
         indir=None,
@@ -323,6 +327,28 @@ def update_namelist_eos(self):
 
         self.add_model_config_options(options=replacements)
 
+    def validate_baselines(self):
+        """
+        Compare variables between output files in this step and in the same
+        step from a baseline run if one was provided.
+
+        Returns
+        -------
+        compared : bool
+            Whether a baseline comparison was performed
+
+        success : bool
+            Whether the outputs were successfully validated against a baseline
+        """
+        # translate variable names to native model names
+        validate_vars = {}
+        for filename, vars in self.validate_vars.items():
+            validate_vars[filename] = (
+                self.component.map_var_list_to_native_model(vars)
+            )
+        self.validate_vars = validate_vars
+        return super().validate_baselines()
+
     def _update_ntasks(self):
         """
         Update ``ntasks`` and ``min_tasks`` for the step based on the estimated