Add map between MPAS-Ocean and Omega variables

docs/developers_guide/ocean/api.md

@@ -337,6 +337,13 @@
 .. autosummary::
    :toctree: generated/
 
+   OceanIOStep
+   OceanIOStep.setup
+   OceanIOStep.map_to_model_dataset
+   OceanIOStep.write_model_dataset
+   OceanIOStep.map_from_model_dataset
+   OceanIOStep.open_model_dataset
+
    OceanModelStep
    OceanModelStep.setup
    OceanModelStep.constrain_resources

docs/developers_guide/ocean/framework.md

@@ -8,6 +8,28 @@ The `ocean` component contains an ever expanding set of shared framework code.
 
 ## Model
 
+### Input and output from an E3SM component
+
+Steps that write input files for or read output files from either Omega or
+MPAS-Ocean should descend from the {py:class}`polaris.ocean.model.OceanIOStep`
+class.  Methods in this class facilitate mapping between MPAS-Ocean variable
+names (used in Polaris) and Omega variable names for tasks that will run Omega.
+
+To map a dataset between MPAS-Ocean variable names and those appropriate for
+the model being run, use the methods
+{py:meth}`polaris.ocean.model.OceanIOStep.map_to_model_dataset()` and
+{py:meth}`polaris.ocean.model.OceanIOStep.map_from_model_dataset()`. These
+methods should be called in Polaris immediatly before writing out input files
+and immediately after opening in output files, respectively. To make opening
+and writing easier, we also provide
+{py:meth}`polaris.ocean.model.OceanIOStep.write_model_dataset()` and
+{py:meth}`polaris.ocean.model.OceanIOStep.open_model_dataset()`, which take
+care of of the mapping in addition to writing and opening a dataset,
+respectively. As new variables are added to Omega, they should be added to the
+`variables` section in the
+[mpaso_to_omega.yaml](https://github.com/E3SM-Project/polaris/blob/main/polaris/ocean/model/mpaso_to_omega.yaml)
+file.
+
 ### Running an E3SM component
 
 Steps that run either Omega or MPAS-Ocean should descend from the
@@ -80,7 +102,8 @@ names for the `ocean` config options in the methods
 {py:meth}`polaris.ocean.model.OceanModelStep.map_yaml_options()` and
 {py:meth}`polaris.ocean.model.OceanModelStep.map_yaml_configs()`.
 As new config options are added to Omega, they should be added to the
-map in the [mpaso_to_omega.yaml](https://github.com/E3SM-Project/polaris/blob/main/polaris/ocean/model/mpaso_to_omega.yaml)
+`config` section in the
+[mpaso_to_omega.yaml](https://github.com/E3SM-Project/polaris/blob/main/polaris/ocean/model/mpaso_to_omega.yaml)
 file. Note that `config_model='Omega'` must be capitalized since this is the
 convention on the model name in Omega's own YAML files.
 

polaris/ocean/convergence/analysis.py

@@ -3,15 +3,14 @@
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
-import xarray as xr
 
-from polaris import Step
 from polaris.mpas import area_for_field, time_index_from_xtime
+from polaris.ocean.model import OceanIOStep
 from polaris.ocean.resolution import resolution_to_subdir
 from polaris.viz import use_mplstyle
 
 
-class ConvergenceAnalysis(Step):
+class ConvergenceAnalysis(OceanIOStep):
     """
     A step for analyzing the output from convergence tests
 
@@ -278,7 +277,7 @@ def compute_error(self, mesh_name, variable_name, zidx=None,
             The error of the variable given by variable_name
         """
         norm_type = {'l2': None, 'inf': np.inf}
-        ds_mesh = xr.open_dataset(f'{mesh_name}_mesh.nc')
+        ds_mesh = self.open_model_dataset(f'{mesh_name}_mesh.nc')
         config = self.config
         section = config['convergence']
         eval_time = section.getfloat('convergence_eval_time')
@@ -333,7 +332,7 @@ def exact_solution(self, mesh_name, field_name, time, zidx=None):
             The exact solution as derived from the initial condition
         """
 
-        ds_init = xr.open_dataset(f'{mesh_name}_init.nc')
+        ds_init = self.open_model_dataset(f'{mesh_name}_init.nc')
         ds_init = ds_init.isel(Time=0)
         if zidx is not None:
             ds_init = ds_init.isel(nVertLevels=zidx)
@@ -363,7 +362,7 @@ def get_output_field(self, mesh_name, field_name, time, zidx=None):
         field_mpas : xarray.DataArray
             model output field
         """
-        ds_out = xr.open_dataset(f'{mesh_name}_output.nc')
+        ds_out = self.open_model_dataset(f'{mesh_name}_output.nc')
 
         tidx = time_index_from_xtime(ds_out.xtime.values, time)
         ds_out = ds_out.isel(Time=tidx)

polaris/ocean/model/__init__.py

@@ -1,2 +1,3 @@
+from polaris.ocean.model.ocean_io_step import OceanIOStep
 from polaris.ocean.model.ocean_model_step import OceanModelStep
 from polaris.ocean.model.time import get_time_interval_string

polaris/ocean/model/mpaso_to_omega.yaml

@@ -1,4 +1,11 @@
-map:
+variables:
+  temperature: Temp
+  salinity: Salt
+  tracer1: Debug1
+  tracer2: Debug2
+  tracer3: Debug3
+
+config:
 - section:
     time_management: TimeManagement
   options:

polaris/ocean/model/ocean_io_step.py

@@ -0,0 +1,154 @@
+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
+
+
+class OceanIOStep(Step):
+    """
+    A step that writes input and/or output files for Omega or MPAS-Ocean
+
+    Attributes
+    ----------
+    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_var_map: Union[None, Dict[str, str]] = None
+        self.omega_to_mpaso_var_map: Union[None, Dict[str, str]] = None
+
+    def setup(self):
+        """
+        Determine if we will make yaml files or namelists and streams files,
+        then, determine the number of MPI tasks to use based on the estimated
+        mesh size
+        """
+        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 map_to_model_dataset(self, ds):
+        """
+        If the model is Omega, rename variables in a dataset from their
+        MPAS-Ocean names to the Omega equivalent (appropriate for input
+        datasets like an initial condition)
+
+        Parameters
+        ----------
+        ds : xarray.Dataset
+            A dataset containing MPAS-Ocean variable names
+
+        Returns
+        -------
+        ds : xarray.Dataset
+            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_var_map is not None
+            ds = ds.rename(self.mpaso_to_omega_var_map)
+        return ds
+
+    def write_model_dataset(self, ds, filename):
+        """
+        Write out the given dataset, mapping variable names from MPAS-Ocean
+        to Omega names if appropriate
+
+        Parameters
+        ----------
+        ds : xarray.Dataset
+            A dataset containing MPAS-Ocean variable names
+
+        filename : str
+            The path for the NetCDF file to write
+        """
+        ds = self.map_to_model_dataset(ds)
+        write_netcdf(ds=ds, fileName=filename)
+
+    def map_from_model_dataset(self, ds):
+        """
+        If the model is Omega, rename variables in a dataset from their
+        Omega names to the MPAS-Ocean equivalent (appropriate for datasets
+        that are output from the model)
+
+        Parameters
+        ----------
+        ds : xarray.Dataset
+            A dataset containing variable names native to either ocean model
+
+        Returns
+        -------
+        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_var_map is not None
+            ds = ds.rename(self.omega_to_mpaso_var_map)
+        return ds
+
+    def open_model_dataset(self, filename):
+        """
+        Open the given dataset, mapping variable names from Omega to MPAS-Ocean
+        names if appropriate
+
+        Parameters
+        ----------
+        filename : str
+            The path for the NetCDF file to open
+
+        Returns
+        -------
+        ds : xarray.Dataset
+            The dataset with variables named as expected in MPAS-Ocean
+        """
+        ds = xr.open_dataset(filename)
+        ds = self.map_from_model_dataset(ds)
+        return ds
+
+    def _read_var_map(self):
+        """
+        Read the map from MPAS-Ocean to Omega config options
+        """
+        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_var_map = nested_dict['variables']
+        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()}

polaris/ocean/model/ocean_model_step.py

@@ -17,7 +17,7 @@ class OceanModelStep(ModelStep):
         ``min_tasks``) are computed dynamically from the number of cells
         in the mesh
 
-    map : dict
+    config_map : dict
         A nested dictionary that maps from MPAS-Ocean to Omega model config
         options
 
@@ -104,7 +104,7 @@ def __init__(self, component, name, subdir=None, indir=None, ntasks=None,
 
         self.dynamic_ntasks = (ntasks is None and min_tasks is None)
 
-        self.map: Union[None, List[Dict[str, Dict[str, str]]]] = None
+        self.config_map: Union[None, List[Dict[str, Dict[str, str]]]] = None
         self.graph_target = graph_target
 
     def setup(self):
@@ -119,7 +119,7 @@ def setup(self):
             self.make_yaml = True
             self.config_models = ['ocean', 'Omega']
             self.yaml = 'omega.yml'
-            self._read_map()
+            self._read_config_map()
             self.partition_graph = False
         elif model == 'mpas-ocean':
             self.config_models = ['ocean', 'mpas-ocean']
@@ -272,7 +272,7 @@ def _update_ntasks(self):
         self.min_tasks = max(1,
                              4 * round(cell_count / (4 * max_cells_per_core)))
 
-    def _read_map(self):
+    def _read_config_map(self):
         """
         Read the map from MPAS-Ocean to Omega config options
         """
@@ -282,7 +282,7 @@ def _read_map(self):
 
         yaml_data = YAML(typ='rt')
         nested_dict = yaml_data.load(text)
-        self.map = nested_dict['map']
+        self.config_map = nested_dict['config']
 
     def _map_mpaso_to_omega_options(self, options):
         """
@@ -310,9 +310,9 @@ def _map_mpaso_to_omega_option(self, option, value):
         out_option = option
         found = False
 
-        assert self.map is not None
+        assert self.config_map is not None
         # traverse the map
-        for entry in self.map:
+        for entry in self.config_map:
             options_dict = entry['options']
             for mpaso_option, omega_option in options_dict.items():
                 if option == mpaso_option:
@@ -359,11 +359,11 @@ def _map_mpaso_to_omega_section_option(self, section, option, value):
         out_section = section
         out_option = option
 
-        assert self.map is not None
+        assert self.config_map is not None
 
         option_found = False
         # traverse the map
-        for entry in self.map:
+        for entry in self.config_map:
             section_dict = entry['section']
             try:
                 omega_section = section_dict[section]

polaris/ocean/tasks/manufactured_solution/analysis.py

@@ -1,5 +1,3 @@
-import xarray as xr
-
 from polaris.ocean.convergence import ConvergenceAnalysis
 from polaris.ocean.tasks.manufactured_solution.exact_solution import (
     ExactSolution,
@@ -69,7 +67,7 @@ def exact_solution(self, mesh_name, field_name, time, zidx=None):
         solution : xarray.DataArray
             The exact solution as derived from the initial condition
         """
-        init = xr.open_dataset(f'{mesh_name}_init.nc')
+        init = self.open_model_dataset(f'{mesh_name}_init.nc')
         exact = ExactSolution(self.config, init)
         if field_name != 'ssh':
             raise ValueError(f'{field_name} is not currently supported')

polaris/ocean/tasks/manufactured_solution/init.py

@@ -1,19 +1,18 @@
 import numpy as np
 import xarray as xr
-from mpas_tools.io import write_netcdf
 from mpas_tools.mesh.conversion import convert, cull
 from mpas_tools.planar_hex import make_planar_hex_mesh
 
-from polaris import Step
 from polaris.mesh.planar import compute_planar_hex_nx_ny
+from polaris.ocean.model import OceanIOStep
 from polaris.ocean.resolution import resolution_to_subdir
 from polaris.ocean.tasks.manufactured_solution.exact_solution import (
     ExactSolution,
 )
 from polaris.ocean.vertical import init_vertical_coord
 
 
-class Init(Step):
+class Init(OceanIOStep):
     """
     A step for creating a mesh and initial condition for the
     manufactured solution test cases
@@ -73,12 +72,12 @@ def run(self):
         ds_mesh = make_planar_hex_mesh(nx=nx, ny=ny, dc=dc,
                                        nonperiodic_x=False,
                                        nonperiodic_y=False)
-        write_netcdf(ds_mesh, 'base_mesh.nc')
+        self.write_model_dataset(ds_mesh, 'base_mesh.nc')
 
         ds_mesh = cull(ds_mesh, logger=logger)
         ds_mesh = convert(ds_mesh, graphInfoFileName='culled_graph.info',
                           logger=logger)
-        write_netcdf(ds_mesh, 'culled_mesh.nc')
+        self.write_model_dataset(ds_mesh, 'culled_mesh.nc')
 
         bottom_depth = config.getfloat('vertical_grid', 'bottom_depth')
 
@@ -112,4 +111,4 @@ def run(self):
                                                     'nVertLevels')
         ds['layerThickness'] = layer_thickness
 
-        write_netcdf(ds, 'initial_state.nc')
+        self.write_model_dataset(ds, 'initial_state.nc')