Skip to content

Commit

Permalink
Update variable mapping methods
Browse files Browse the repository at this point in the history
This merge also adds or updates some attribute docstrings
  • Loading branch information
xylar committed Oct 17, 2024
1 parent 714fd8c commit e79d3cf
Showing 1 changed file with 37 additions and 19 deletions.
56 changes: 37 additions & 19 deletions polaris/ocean/model/ocean_model_step.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,14 @@ class OceanModelStep(ModelStep):
``min_tasks``) are computed dynamically from the number of cells
in the mesh
map : dict
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``
config_map : dict
A nested dictionary that maps from MPAS-Ocean to Omega model config
options
Expand Down Expand Up @@ -104,7 +111,8 @@ def __init__(self, component, name, subdir=None, indir=None, ntasks=None,

self.dynamic_ntasks = (ntasks is None and min_tasks is None)

self.var_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
self.config_map: Union[None, List[Dict[str, Dict[str, str]]]] = None
self.graph_target = graph_target

Expand All @@ -116,11 +124,11 @@ def setup(self):
"""
config = self.config
model = config.get('ocean', 'model')
self._read_var_map()
if model == 'omega':
self.make_yaml = True
self.config_models = ['ocean', 'Omega']
self.yaml = 'omega.yml'
self._read_var_map()
self._read_config_map()
self.partition_graph = False
elif model == 'mpas-ocean':
Expand Down Expand Up @@ -160,10 +168,11 @@ def compute_cell_count(self):
"""
return None

def rename_vars_mpaso_to_omega(self, ds):
def map_input_dataset(self, ds):
"""
Rename variables in a dataset from their MPAS-Ocean names to the Omega
equivalent
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
----------
Expand All @@ -173,31 +182,37 @@ def rename_vars_mpaso_to_omega(self, ds):
Returns
-------
ds : xarray.Dataset
The same dataset with variables renamed to their Omega equivalents
The same dataset with variables renamed as appropriate for the
ocean model being run
"""
assert self.var_map is not None
ds = ds.rename(self.var_map)
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 rename_vars_omega_to_mpaso(self, ds):
def map_output_dataset(self, ds):
"""
Rename variables in a dataset from their Omega names to the MPAS-Ocean
equivalent
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 Omega variable names
A dataset containing variable names native to either ocean model
Returns
-------
ds : xarray.Dataset
The same dataset with variables renamed to their MPAS-Ocean
equivalents
The same dataset with variables named as expected in MPAS-Ocean
"""
assert self.var_map is not None
inv_map = {v: k for k, v in self.var_map.items()}
ds = ds.rename(inv_map)
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 map_yaml_options(self, options, config_model):
Expand Down Expand Up @@ -324,7 +339,10 @@ def _read_var_map(self):

yaml_data = YAML(typ='rt')
nested_dict = yaml_data.load(text)
self.var_map = nested_dict['variables']
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()}

def _read_config_map(self):
"""
Expand Down

0 comments on commit e79d3cf

Please sign in to comment.