Source code for polaris.ocean.model.ocean_io_step

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


[docs] 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`` """
[docs] 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
[docs] 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()
[docs] 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 """ 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
[docs] 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)
[docs] 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 """ 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
[docs] 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'] 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()}