diff --git a/docs/developers_guide/ocean/api.md b/docs/developers_guide/ocean/api.md index 20bb6c098..e6f3af643 100644 --- a/docs/developers_guide/ocean/api.md +++ b/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 diff --git a/polaris/ocean/model/ocean_io_step.py b/polaris/ocean/model/ocean_io_step.py index 92c140be7..360b8b5f9 100644 --- a/polaris/ocean/model/ocean_io_step.py +++ b/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) diff --git a/polaris/ocean/model/ocean_model_step.py b/polaris/ocean/model/ocean_model_step.py index 835368126..157335775 100644 --- a/polaris/ocean/model/ocean_model_step.py +++ b/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 diff --git a/polaris/tasks/ocean/__init__.py b/polaris/tasks/ocean/__init__.py index 2954fe71f..1ff62db03 100644 --- a/polaris/tasks/ocean/__init__.py +++ b/polaris/tasks/ocean/__init__.py @@ -1,9 +1,27 @@ +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 Component class Ocean(Component): """ The collection of all test case for the MPAS-Ocean core + + Attributes + ---------- + model : str + The ocean model being used, either 'mpas-ocean' or 'omega' + + mpaso_to_omega_dim_map : dict + A map from MPAS-Ocean dimension names to their Omega equivalents + + mpaso_to_omega_var_map : dict + A map from MPAS-Ocean variable names to their Omega equivalents """ def __init__(self): @@ -11,6 +29,9 @@ def __init__(self): Construct the collection of MPAS-Ocean test cases """ super().__init__(name='ocean') + self.model: Union[None, str] = None + self.mpaso_to_omega_dim_map: Union[None, Dict[str, str]] = None + self.mpaso_to_omega_var_map: Union[None, Dict[str, str]] = None def configure(self, config): """ @@ -29,6 +50,184 @@ def configure(self, config): config.add_from_package('polaris.ocean', configs[model]) + if model == 'omega': + self._read_var_map() + self.model = model + + def map_to_native_model_vars(self, ds): + """ + If the model is Omega, rename dimensions and 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 + """ + model = self.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 + + def map_var_list_to_native_model(self, var_list): + """ + If the model is Omega, rename variables from their MPAS-Ocean names to + the Omega equivalent (appropriate for validation variable lists) + + Parameters + ---------- + var_list : list of str + A list of MPAS-Ocean variable names + + Returns + ------- + renamed_vars : list of str + The same list with variables renamed as appropriate for the + ocean model being run + """ + renamed_vars = var_list + model = self.model + if model == 'omega': + assert self.mpaso_to_omega_var_map is not None + renamed_vars = [ + v + for k, v in self.mpaso_to_omega_var_map.items() + if k in var_list + ] + return renamed_vars + + def write_model_dataset(self, ds, filename): + """ + Write out the given dataset, mapping dimension and 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_native_model_vars(ds) + write_netcdf(ds=ds, fileName=filename) + + def map_from_native_model_vars(self, ds): + """ + If the model is Omega, rename dimensions and 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 + """ + model = self.model + if model == 'omega': + # switch keys and values in mpaso_to_omega maps to get + # omega to mpaso maps + assert self.mpaso_to_omega_dim_map is not None + rename = { + k: v + for v, k 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 v, k in self.mpaso_to_omega_var_map.items() if k in ds + } + rename.update(rename_vars) + ds = ds.rename(rename) + return ds + + def map_var_list_from_native_model(self, var_list): + """ + If the model is Omega, rename variables from their Omega names to + the MPAS-Ocean equivalent + + Parameters + ---------- + var_list : list of str + A list of MPAS-Ocean variable names + + Returns + ------- + renamed_vars : list of str + The same list with variables renamed as appropriate for the + ocean model being run + """ + renamed_vars = var_list + model = self.model + if model == 'omega': + # switch keys and values in mpaso_to_omega maps to get + # omega to mpaso maps + assert self.mpaso_to_omega_var_map is not None + renamed_vars = [ + v + for v, k in self.mpaso_to_omega_var_map.items() + if k in var_list + ] + return renamed_vars + + def open_model_dataset(self, filename, **kwargs): + """ + Open the given dataset, mapping variable and dimension names from Omega + to MPAS-Ocean names if appropriate + + Parameters + ---------- + filename : str + The path for the NetCDF file to open + + kwargs + keyword arguments passed to `xarray.open_dataset()` + + Returns + ------- + 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'] + # create a single module-level instance available to other components ocean = Ocean()