From 714fd8c7ba1f4030258107173e3e56fed52e2936 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Thu, 17 Oct 2024 17:01:34 -0500 Subject: [PATCH 01/12] Add map between MPAS-Ocean and Omega variables --- polaris/ocean/model/mpaso_to_omega.yaml | 9 +++- polaris/ocean/model/ocean_model_step.py | 70 ++++++++++++++++++++++--- 2 files changed, 70 insertions(+), 9 deletions(-) diff --git a/polaris/ocean/model/mpaso_to_omega.yaml b/polaris/ocean/model/mpaso_to_omega.yaml index 067e9bd21..5d64de043 100644 --- a/polaris/ocean/model/mpaso_to_omega.yaml +++ b/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: diff --git a/polaris/ocean/model/ocean_model_step.py b/polaris/ocean/model/ocean_model_step.py index 598d01215..0bb6cf85f 100644 --- a/polaris/ocean/model/ocean_model_step.py +++ b/polaris/ocean/model/ocean_model_step.py @@ -104,7 +104,8 @@ 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.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 def setup(self): @@ -115,11 +116,12 @@ 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_map() + self._read_config_map() self.partition_graph = False elif model == 'mpas-ocean': self.config_models = ['ocean', 'mpas-ocean'] @@ -158,6 +160,46 @@ def compute_cell_count(self): """ return None + def rename_vars_mpaso_to_omega(self, ds): + """ + Rename variables in a dataset from their MPAS-Ocean names to the Omega + equivalent + + Parameters + ---------- + ds : xarray.Dataset + A dataset containing MPAS-Ocean variable names + + Returns + ------- + ds : xarray.Dataset + The same dataset with variables renamed to their Omega equivalents + """ + assert self.var_map is not None + ds = ds.rename(self.var_map) + return ds + + def rename_vars_omega_to_mpaso(self, ds): + """ + Rename variables in a dataset from their Omega names to the MPAS-Ocean + equivalent + + Parameters + ---------- + ds : xarray.Dataset + A dataset containing Omega variable names + + Returns + ------- + ds : xarray.Dataset + The same dataset with variables renamed to their MPAS-Ocean + equivalents + """ + assert self.var_map is not None + inv_map = {v: k for k, v in self.var_map.items()} + ds = ds.rename(inv_map) + return ds + def map_yaml_options(self, options, config_model): """ A mapping between model config options from MPAS-Ocean to Omega @@ -272,7 +314,19 @@ def _update_ntasks(self): self.min_tasks = max(1, 4 * round(cell_count / (4 * max_cells_per_core))) - def _read_map(self): + 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.var_map = nested_dict['variables'] + + def _read_config_map(self): """ Read the map from MPAS-Ocean to Omega config options """ @@ -282,7 +336,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 +364,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 +413,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] From e79d3cfa13bafb0f701ab36c27596d648c74cf7b Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Thu, 17 Oct 2024 18:07:36 -0500 Subject: [PATCH 02/12] Update variable mapping methods This merge also adds or updates some attribute docstrings --- polaris/ocean/model/ocean_model_step.py | 56 ++++++++++++++++--------- 1 file changed, 37 insertions(+), 19 deletions(-) diff --git a/polaris/ocean/model/ocean_model_step.py b/polaris/ocean/model/ocean_model_step.py index 0bb6cf85f..69d53d8c4 100644 --- a/polaris/ocean/model/ocean_model_step.py +++ b/polaris/ocean/model/ocean_model_step.py @@ -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 @@ -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 @@ -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': @@ -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 ---------- @@ -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): @@ -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): """ From 68c036f4a9cadaacd53124ac36fb8fc78d1c8913 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Thu, 17 Oct 2024 18:08:57 -0500 Subject: [PATCH 03/12] Add documentation of the new feature --- docs/developers_guide/ocean/api.md | 2 ++ docs/developers_guide/ocean/framework.md | 16 ++++++++++++++++ 2 files changed, 18 insertions(+) diff --git a/docs/developers_guide/ocean/api.md b/docs/developers_guide/ocean/api.md index c308eb85f..f6fe9f298 100644 --- a/docs/developers_guide/ocean/api.md +++ b/docs/developers_guide/ocean/api.md @@ -341,6 +341,8 @@ OceanModelStep.setup OceanModelStep.constrain_resources OceanModelStep.compute_cell_count + OceanModelStep.map_input_dataset + OceanModelStep.map_output_dataset OceanModelStep.map_yaml_options OceanModelStep.map_yaml_configs diff --git a/docs/developers_guide/ocean/framework.md b/docs/developers_guide/ocean/framework.md index 6f93ec09e..94ec1b769 100644 --- a/docs/developers_guide/ocean/framework.md +++ b/docs/developers_guide/ocean/framework.md @@ -84,6 +84,22 @@ map in the [mpaso_to_omega.yaml](https://github.com/E3SM-Project/polaris/blob/ma file. Note that `config_model='Omega'` must be capitalized since this is the convention on the model name in Omega's own YAML files. +#### Mapping between MPAS-Ocean to Omega variable names + +Similarly to config options, it is necessary to be able to map between +MPAS-Ocean and Omega variable names in both input (e.g. initial condition) +and output datasets. + +We implement a mapping between the MPAS-Ocean and Omega variable names in the +methods +{py:meth}`polaris.ocean.model.OceanModelStep.map_input_dataset()` and +{py:meth}`polaris.ocean.model.OceanModelStep.map_output_dataset()`. These +methods should be called in Polaris immediatly before writing out input files +and immediately after reading in output files, respectively. +As new variables 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) +file. + #### Setting MPI resources The target and minimum number of MPI tasks (`ntasks` and `min_tasks`, From f19cea7af9447bdc542e3c10656731fc91acacba Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Thu, 17 Oct 2024 20:20:14 -0500 Subject: [PATCH 04/12] Add OceanIOStep for reading/writing ocean model files This step is appropriate for most ocean steps that are not OceanModelSteps, since often variables should be translated to or from the native model names to MPAS-Ocean names for use in Polaris --- polaris/ocean/model/__init__.py | 1 + polaris/ocean/model/ocean_io_step.py | 154 ++++++++++++++++++++++++ polaris/ocean/model/ocean_model_step.py | 72 ----------- 3 files changed, 155 insertions(+), 72 deletions(-) create mode 100644 polaris/ocean/model/ocean_io_step.py diff --git a/polaris/ocean/model/__init__.py b/polaris/ocean/model/__init__.py index 3bcb9b7f4..018976797 100644 --- a/polaris/ocean/model/__init__.py +++ b/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 diff --git a/polaris/ocean/model/ocean_io_step.py b/polaris/ocean/model/ocean_io_step.py new file mode 100644 index 000000000..3401bcd0d --- /dev/null +++ b/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()} diff --git a/polaris/ocean/model/ocean_model_step.py b/polaris/ocean/model/ocean_model_step.py index 69d53d8c4..5c70269e0 100644 --- a/polaris/ocean/model/ocean_model_step.py +++ b/polaris/ocean/model/ocean_model_step.py @@ -17,13 +17,6 @@ class OceanModelStep(ModelStep): ``min_tasks``) are computed dynamically from the number of cells in the mesh - 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 @@ -111,8 +104,6 @@ def __init__(self, component, name, subdir=None, indir=None, ntasks=None, self.dynamic_ntasks = (ntasks is None and min_tasks is 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 @@ -128,7 +119,6 @@ def setup(self): 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': @@ -168,53 +158,6 @@ def compute_cell_count(self): """ return None - def map_input_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 map_output_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 map_yaml_options(self, options, config_model): """ A mapping between model config options from MPAS-Ocean to Omega @@ -329,21 +272,6 @@ def _update_ntasks(self): self.min_tasks = max(1, 4 * round(cell_count / (4 * max_cells_per_core))) - 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()} - def _read_config_map(self): """ Read the map from MPAS-Ocean to Omega config options From 702a178a2fc5194a72913462baff92568fb15868 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Thu, 17 Oct 2024 20:22:01 -0500 Subject: [PATCH 05/12] Add OceanIOStep to the docs --- docs/developers_guide/ocean/api.md | 9 ++++-- docs/developers_guide/ocean/framework.md | 41 ++++++++++++++---------- 2 files changed, 31 insertions(+), 19 deletions(-) diff --git a/docs/developers_guide/ocean/api.md b/docs/developers_guide/ocean/api.md index f6fe9f298..4abd4e2f2 100644 --- a/docs/developers_guide/ocean/api.md +++ b/docs/developers_guide/ocean/api.md @@ -337,12 +337,17 @@ .. 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 OceanModelStep.compute_cell_count - OceanModelStep.map_input_dataset - OceanModelStep.map_output_dataset OceanModelStep.map_yaml_options OceanModelStep.map_yaml_configs diff --git a/docs/developers_guide/ocean/framework.md b/docs/developers_guide/ocean/framework.md index 94ec1b769..ac1d6935e 100644 --- a/docs/developers_guide/ocean/framework.md +++ b/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,26 +102,11 @@ 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. -#### Mapping between MPAS-Ocean to Omega variable names - -Similarly to config options, it is necessary to be able to map between -MPAS-Ocean and Omega variable names in both input (e.g. initial condition) -and output datasets. - -We implement a mapping between the MPAS-Ocean and Omega variable names in the -methods -{py:meth}`polaris.ocean.model.OceanModelStep.map_input_dataset()` and -{py:meth}`polaris.ocean.model.OceanModelStep.map_output_dataset()`. These -methods should be called in Polaris immediatly before writing out input files -and immediately after reading in output files, respectively. -As new variables 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) -file. - #### Setting MPI resources The target and minimum number of MPI tasks (`ntasks` and `min_tasks`, From e60ced5fcad1b6b5b4676725a920f85e12bab46a Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Thu, 17 Oct 2024 20:23:13 -0500 Subject: [PATCH 06/12] Update convergence framework and manufactured solution --- polaris/ocean/convergence/analysis.py | 11 +++++------ polaris/ocean/tasks/manufactured_solution/analysis.py | 4 +--- polaris/ocean/tasks/manufactured_solution/init.py | 11 +++++------ polaris/ocean/tasks/manufactured_solution/viz.py | 11 +++++------ 4 files changed, 16 insertions(+), 21 deletions(-) diff --git a/polaris/ocean/convergence/analysis.py b/polaris/ocean/convergence/analysis.py index 2e39fcc45..f2bbf85a1 100644 --- a/polaris/ocean/convergence/analysis.py +++ b/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) diff --git a/polaris/ocean/tasks/manufactured_solution/analysis.py b/polaris/ocean/tasks/manufactured_solution/analysis.py index 25a17e717..e7947192e 100644 --- a/polaris/ocean/tasks/manufactured_solution/analysis.py +++ b/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') diff --git a/polaris/ocean/tasks/manufactured_solution/init.py b/polaris/ocean/tasks/manufactured_solution/init.py index 788cae069..63cf166f2 100644 --- a/polaris/ocean/tasks/manufactured_solution/init.py +++ b/polaris/ocean/tasks/manufactured_solution/init.py @@ -1,11 +1,10 @@ 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, @@ -13,7 +12,7 @@ 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') diff --git a/polaris/ocean/tasks/manufactured_solution/viz.py b/polaris/ocean/tasks/manufactured_solution/viz.py index 65a232f7b..3fe4b8ef1 100644 --- a/polaris/ocean/tasks/manufactured_solution/viz.py +++ b/polaris/ocean/tasks/manufactured_solution/viz.py @@ -3,9 +3,8 @@ import cmocean # noqa: F401 import matplotlib.pyplot as plt import numpy as np -import xarray as xr -from polaris import Step +from polaris.ocean.model import OceanIOStep from polaris.ocean.resolution import resolution_to_subdir from polaris.ocean.tasks.manufactured_solution.exact_solution import ( ExactSolution, @@ -13,7 +12,7 @@ from polaris.viz import plot_horiz_field, use_mplstyle -class Viz(Step): +class Viz(OceanIOStep): """ A step for visualizing the output from the manufactured solution test case @@ -73,9 +72,9 @@ def run(self): error_range = None for i, res in enumerate(resolutions): mesh_name = resolution_to_subdir(res) - ds_mesh = xr.open_dataset(f'mesh_{mesh_name}.nc') - ds_init = xr.open_dataset(f'init_{mesh_name}.nc') - ds = xr.open_dataset(f'output_{mesh_name}.nc') + ds_mesh = self.open_model_dataset(f'mesh_{mesh_name}.nc') + ds_init = self.open_model_dataset(f'init_{mesh_name}.nc') + ds = self.open_model_dataset(f'output_{mesh_name}.nc') exact = ExactSolution(config, ds_init) t0 = datetime.datetime.strptime(ds.xtime.values[0].decode(), From 8127134a269dcb74e6caa4e87ac4c3609801995f Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Sat, 19 Oct 2024 12:31:31 -0500 Subject: [PATCH 07/12] Rename methods and fix docstring based on review --- polaris/ocean/model/ocean_io_step.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/polaris/ocean/model/ocean_io_step.py b/polaris/ocean/model/ocean_io_step.py index 3401bcd0d..304c59bb2 100644 --- a/polaris/ocean/model/ocean_io_step.py +++ b/polaris/ocean/model/ocean_io_step.py @@ -44,9 +44,8 @@ def __init__(self, component, name, **kwargs): 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 + If the ocean model is Omega, set up maps between Omega and MPAS-Ocean + variable names """ config = self.config model = config.get('ocean', 'model') @@ -56,7 +55,7 @@ def setup(self): raise ValueError(f'Unexpected ocean model: {model}') super().setup() - def map_to_model_dataset(self, ds): + def map_to_native_model_vars(self, ds): """ If the model is Omega, rename variables in a dataset from their MPAS-Ocean names to the Omega equivalent (appropriate for input @@ -93,10 +92,10 @@ def write_model_dataset(self, ds, filename): filename : str The path for the NetCDF file to write """ - ds = self.map_to_model_dataset(ds) + ds = self.map_to_native_model_vars(ds) write_netcdf(ds=ds, fileName=filename) - def map_from_model_dataset(self, ds): + def map_from_native_model_vars(self, ds): """ If the model is Omega, rename variables in a dataset from their Omega names to the MPAS-Ocean equivalent (appropriate for datasets @@ -135,7 +134,7 @@ def open_model_dataset(self, filename): The dataset with variables named as expected in MPAS-Ocean """ ds = xr.open_dataset(filename) - ds = self.map_from_model_dataset(ds) + ds = self.map_from_native_model_vars(ds) return ds def _read_var_map(self): From 76135a58b6c06c722db5a9efc1d89e0fb0dd838c Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Sat, 19 Oct 2024 12:33:45 -0500 Subject: [PATCH 08/12] Pass kwargs to xr.open_dataset() --- polaris/ocean/model/ocean_io_step.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/polaris/ocean/model/ocean_io_step.py b/polaris/ocean/model/ocean_io_step.py index 304c59bb2..711fa31c5 100644 --- a/polaris/ocean/model/ocean_io_step.py +++ b/polaris/ocean/model/ocean_io_step.py @@ -118,7 +118,7 @@ def map_from_native_model_vars(self, ds): ds = ds.rename(self.omega_to_mpaso_var_map) return ds - def open_model_dataset(self, filename): + def open_model_dataset(self, filename, **kwargs): """ Open the given dataset, mapping variable names from Omega to MPAS-Ocean names if appropriate @@ -128,12 +128,15 @@ def open_model_dataset(self, filename): 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) + ds = xr.open_dataset(filename, **kwargs) ds = self.map_from_native_model_vars(ds) return ds From 97242b469d706b3c9f5df3f940b725cc957347cc Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Sat, 19 Oct 2024 12:56:44 -0500 Subject: [PATCH 09/12] Fix renaming --- polaris/ocean/model/ocean_io_step.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/polaris/ocean/model/ocean_io_step.py b/polaris/ocean/model/ocean_io_step.py index 711fa31c5..733258c86 100644 --- a/polaris/ocean/model/ocean_io_step.py +++ b/polaris/ocean/model/ocean_io_step.py @@ -76,7 +76,9 @@ def map_to_native_model_vars(self, ds): 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) + rename = {k: v for k, v in self.mpaso_to_omega_var_map.items() + if k in ds} + ds = ds.rename(rename) return ds def write_model_dataset(self, ds, filename): @@ -115,7 +117,9 @@ def map_from_native_model_vars(self, ds): 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) + rename = {k: v for k, v in self.omega_to_mpaso_var_map.items() + if k in ds} + ds = ds.rename(rename) return ds def open_model_dataset(self, filename, **kwargs): From 62154006714d4156ed9c99a97b0340184229d409 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Sat, 19 Oct 2024 13:34:23 -0500 Subject: [PATCH 10/12] Add ssh to mpaso --> omega map --- polaris/ocean/model/mpaso_to_omega.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/polaris/ocean/model/mpaso_to_omega.yaml b/polaris/ocean/model/mpaso_to_omega.yaml index 5d64de043..dbb7baa1c 100644 --- a/polaris/ocean/model/mpaso_to_omega.yaml +++ b/polaris/ocean/model/mpaso_to_omega.yaml @@ -4,6 +4,7 @@ variables: tracer1: Debug1 tracer2: Debug2 tracer3: Debug3 + ssh: SshCell config: - section: From 0a17c6d8fcf88484654c0b48a4195628836e4c24 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Sat, 19 Oct 2024 14:06:46 -0500 Subject: [PATCH 11/12] Update docs --- docs/developers_guide/ocean/api.md | 4 ++-- docs/developers_guide/ocean/framework.md | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/developers_guide/ocean/api.md b/docs/developers_guide/ocean/api.md index 4abd4e2f2..e1d8d50a2 100644 --- a/docs/developers_guide/ocean/api.md +++ b/docs/developers_guide/ocean/api.md @@ -339,9 +339,9 @@ OceanIOStep OceanIOStep.setup - OceanIOStep.map_to_model_dataset + OceanIOStep.map_to_native_model_vars OceanIOStep.write_model_dataset - OceanIOStep.map_from_model_dataset + OceanIOStep.map_from_native_model_vars OceanIOStep.open_model_dataset OceanModelStep diff --git a/docs/developers_guide/ocean/framework.md b/docs/developers_guide/ocean/framework.md index ac1d6935e..4f16e5b54 100644 --- a/docs/developers_guide/ocean/framework.md +++ b/docs/developers_guide/ocean/framework.md @@ -17,8 +17,8 @@ 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 +{py:meth}`polaris.ocean.model.OceanIOStep.map_to_native_model_vars()` and +{py:meth}`polaris.ocean.model.OceanIOStep.map_from_native_model_vars()`. 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 From 72fa2e45395fe29c89626452206a58a96ac260ca Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Sat, 19 Oct 2024 14:49:51 -0500 Subject: [PATCH 12/12] Move ocean-specific docs to ocean framework --- docs/developers_guide/framework/model.md | 106 +++++++++++------------ docs/developers_guide/ocean/framework.md | 4 + 2 files changed, 56 insertions(+), 54 deletions(-) diff --git a/docs/developers_guide/framework/model.md b/docs/developers_guide/framework/model.md index a407296bd..f10ea70ab 100644 --- a/docs/developers_guide/framework/model.md +++ b/docs/developers_guide/framework/model.md @@ -7,41 +7,39 @@ Steps that run a standalone build of an E3SM component should descend from the {py:class}`polaris.ModelStep` class. -By default (if the attribute `update_pio = True`), at runtime, the namelist -options associated with the [PIO library](https://ncar.github.io/ParallelIO/) -will automatically be set. If the `make_graph = True`, the mesh will be +By default (if the attribute `update_pio = True`), at runtime, the namelist +options associated with the [PIO library](https://ncar.github.io/ParallelIO/) +will automatically be set. If the `make_graph = True`, the mesh will be partitioned across MPI tasks at runtime. -During construction or by setting attributes directly, you can can provide +During construction or by setting attributes directly, you can can provide non-default names for the graph, namelist and streams files. At construction -or with the {py:meth}`polaris.ModelStep.set_model_resources()` -method, you can set number of tasks, threads, etc. determined from the -`ntasks`, `min_tasks`, `cpus_per_task`, `` min_cpus_per_task` `` and +or with the {py:meth}`polaris.ModelStep.set_model_resources()` +method, you can set number of tasks, threads, etc. determined from the +`ntasks`, `min_tasks`, `cpus_per_task`, `` min_cpus_per_task` `` and `openmp_threads` attributes. These resources need to be set at construction or -in the {ref}`dev-step-setup` method (i.e. before calling {ref}`dev-step-run`) -so that the polaris framework can ensure that the required resources are -available. If the graph partition file has been constructed prior to the ocean -model step, the path to the graph file should be provided in the -`graph_target` argument to {py:meth}`polaris.ocean.OceanModelStep.__init__()`. +in the {ref}`dev-step-setup` method (i.e. before calling {ref}`dev-step-run`) +so that the polaris framework can ensure that the required resources are +available. (dev-model-yaml-namelists-and-streams)= ### Adding yaml, namelist and streams files -Components and tasks can provide yaml config options, -namelist and streams files that are used to replace default model config -options and streams definitions before the E3SM component gets run. Namelist -and streams files within the `polaris` package must start with the prefix -`namelist.` and `streams.`, respectively, to ensure that they are included +Components and tasks can provide yaml config options, +namelist and streams files that are used to replace default model config +options and streams definitions before the E3SM component gets run. Namelist +and streams files within the `polaris` package must start with the prefix +`namelist.` and `streams.`, respectively, to ensure that they are included when we build the package. Yaml files must end with `.yaml` or `.yml` for the same reason. You can make calls to {py:meth}`polaris.ModelStep.add_namelist_file()`, {py:meth}`polaris.ModelStep.add_yaml_file()`, {py:meth}`polaris.ModelStep.add_model_config_options()` and -{py:meth}`polaris.ModelStep.add_streams_file()` as described below to indicate -how yaml, namelist and streams file should be built up by modifying the -defaults for the E3SM component. The yaml, namelists and streams files +{py:meth}`polaris.ModelStep.add_streams_file()` as described below to indicate +how yaml, namelist and streams file should be built up by modifying the +defaults for the E3SM component. The yaml, namelists and streams files themselves are generated automatically (which of these depends on the E3SM component in question) as part of setting up the task. @@ -49,11 +47,11 @@ component in question) as part of setting up the task. #### Adding a yaml file -Typically, a step that runs an E3SM component will include one or more calls +Typically, a step that runs an E3SM component will include one or more calls to {py:meth}`polaris.ModelStep.add_namelist_file()` or {py:meth}`polaris.ModelStep.add_yaml_file()` within the {ref}`dev-step-init` or {ref}`dev-step-setup` method. Calling one of these methods simply adds the -file to a list that will be parsed if and when the step gets set up. (This +file to a list that will be parsed if and when the step gets set up. (This way, it is safe to add namelist files to a step in init even if that task will never get set up or run.) @@ -72,7 +70,7 @@ ocean: config_block_decomp_file_prefix: graph.info.part. time_integration: config_time_integrator: RK4 - + streams: mesh: filename_template: init.nc @@ -113,14 +111,14 @@ self.add_yaml_file('polaris.ocean.tasks.global_convergence.cosine_bell', Model config values are replaced by the files (or options, see below) in the sequence they are given. This way, you can add the model config substitutions -common to related tasks first, and then override those with the replacements +common to related tasks first, and then override those with the replacements specific to the task or step. (dev-model-add-namelists-file)= #### Adding a namelist file -Typically, a step that runs the E3SM component will include one or more calls +Typically, a step that runs the E3SM component will include one or more calls to {py:meth}`polaris.ModelStep.add_namelist_file()` or {py:meth}`polaris.ModelStep.add_yaml_file()` within the {ref}`dev-step-init` or {ref}`dev-step-setup` method. Calling this method simply adds the file to @@ -155,7 +153,7 @@ self.add_namelist_file('polaris.ocean.tasks.baroclinic_channel', Namelist values are replaced by the files (or options, see below) in the sequence they are given. This way, you can add the namelist substitutions for -that are common to related tasks first, and then override those with the +that are common to related tasks first, and then override those with the replacements that are specific to the task or step. (dev-model-add-model-config-options)= @@ -163,17 +161,17 @@ replacements that are specific to the task or step. #### Adding model config options Sometimes, it is easier to replace yaml or namelist options (together referred -to as model config options) using a dictionary within the code, rather than -a yaml or namelist file. This is appropriate when there are only 1 or 2 +to as model config options) using a dictionary within the code, rather than +a yaml or namelist file. This is appropriate when there are only 1 or 2 options to replace (so creating a file seems like overkill) or when the -model config options rely on values that are determined by the code (e.g. +model config options rely on values that are determined by the code (e.g. different values for different resolutions). Simply create a dictionary -replacements and call {py:meth}`polaris.ModelStep.add_model_config_options()` +replacements and call {py:meth}`polaris.ModelStep.add_model_config_options()` either at init or in the `setup()` method of the step. These replacements are -parsed, along with replacements from files, in the order they are added. +parsed, along with replacements from files, in the order they are added. Thus, you could add replacements from a model config file common to multiple -tasks, specific to a task, and/or specific to step. Then, you could override -them with namelist options in a dictionary for the task or step, as in this +tasks, specific to a task, and/or specific to step. Then, you could override +them with namelist options in a dictionary for the task or step, as in this example: ```python @@ -189,7 +187,7 @@ self.add_yaml_file('polaris.ocean.tasks.baroclinic_channel', ``` -Here, we set the viscosity `nu`, which is an option passed in when creating +Here, we set the viscosity `nu`, which is an option passed in when creating this step. Then, we get default model config options for ocean model output (`output.yaml`) and for baroclinic channel forward steps (`forward.yaml`). @@ -206,7 +204,7 @@ file. It is sometimes useful to have model config options that are based on Polaris config options and/or algorithms. In such cases, the model config options need to be computed once at setup and again (possibly based on updated -config options) at runtime. A step needs to override the +config options) at runtime. A step needs to override the {py:meth}`polaris.ModelStep.dynamic_model_config()` method, e.g.: ```python @@ -262,8 +260,8 @@ streams files are XML documents, requiring some slightly more sophisticated parsing. Typically, a step that runs the E3SM component will include one or more calls -to {py:meth}`polaris.ModelStep.add_streams_file()` within the -{ref}`dev-step-init` or {ref}`dev-step-setup` method. Calling this function +to {py:meth}`polaris.ModelStep.add_streams_file()` within the +{ref}`dev-step-init` or {ref}`dev-step-setup` method. Calling this function simply adds the file to a list within the `step` dictionary that will be parsed if an when the step gets set up. (This way, it is safe to add streams files to a step at init even if that task will never get set up or run.) @@ -378,20 +376,20 @@ Thus, calls to {py:meth}`polaris.ModelStep.add_streams_file()` with ### Adding E3SM component as an input -If a step involves running the E3SM component, it should descend from -:py:class`polaris.ModelStep`. The model executable will automatically be +If a step involves running the E3SM component, it should descend from +:py:class`polaris.ModelStep`. The model executable will automatically be linked and added as an input to the step. This way, if the user has forgotten -to compile the model, this will be obvious by the broken symlink and the step +to compile the model, this will be obvious by the broken symlink and the step will immediately fail because of the missing input. The path to the executable -is automatically detected based on the work directory for the step and the +is automatically detected based on the work directory for the step and the config options. ## Partitioning the mesh -The method {py:meth}`polaris.ModelStep.partition()` calls the graph -partitioning executable ([gpmetis](https://arc.vt.edu/userguide/metis/) by +The method {py:meth}`polaris.ModelStep.partition()` calls the graph +partitioning executable ([gpmetis](https://arc.vt.edu/userguide/metis/) by default) to divide up the MPAS mesh across MPI tasks. If you create a -`ModelStep` with `partition_graph=True` (the default), this method is called +`ModelStep` with `partition_graph=True` (the default), this method is called automatically. In some circumstances, a step may need to partition the mesh separately from @@ -403,18 +401,18 @@ provide `partition_graph=False` and then call ## Updating PIO namelist options -You can use {py:meth}`polaris.ModelStep.update_namelist_pio()` to -automatically set the MPAS namelist options `config_pio_num_iotasks` and -`config_pio_stride` such that there is 1 PIO task per node of the MPAS run. -This is particularly useful for PIO v1, which we have found performs much -better in this configuration than when there is 1 PIO task per core, the MPAS +You can use {py:meth}`polaris.ModelStep.update_namelist_pio()` to +automatically set the MPAS namelist options `config_pio_num_iotasks` and +`config_pio_stride` such that there is 1 PIO task per node of the MPAS run. +This is particularly useful for PIO v1, which we have found performs much +better in this configuration than when there is 1 PIO task per core, the MPAS default. When running with PIO v2, we have found little performance difference -between the MPAS default and the polaris default of one task per node, so we +between the MPAS default and the polaris default of one task per node, so we feel this is a safe default. -By default, this function is called in -{py:meth}`polaris.ModelStep.runtime_setup()`. If the same namelist -file is used for multiple model runs, it may be useful to update the number of +By default, this function is called in +{py:meth}`polaris.ModelStep.runtime_setup()`. If the same namelist +file is used for multiple model runs, it may be useful to update the number of PIO tasks only once. In this case, set the attribute `update_pio = False` and {py:meth}`polaris.ModelStep.update_namelist_pio()` yourself. @@ -429,6 +427,6 @@ Some polaris tasks take advantage of the fact that the can produce a graph file as part of the process of culling cells from an MPAS mesh. In tasks that do not require cells to be culled, you can call {py:func}`polaris.model_step.make_graph_file()` to produce a graph file -from an MPAS mesh file. Optionally, you can provide the name of an MPAS field +from an MPAS mesh file. Optionally, you can provide the name of an MPAS field on cells in the mesh file that gives different weight to different cells (`weight_field`) in the partitioning process. diff --git a/docs/developers_guide/ocean/framework.md b/docs/developers_guide/ocean/framework.md index 4f16e5b54..cca5f7986 100644 --- a/docs/developers_guide/ocean/framework.md +++ b/docs/developers_guide/ocean/framework.md @@ -37,6 +37,10 @@ Steps that run either Omega or MPAS-Ocean should descend from the from {py:class}`polaris.ModelStep`, so there is a lot of relevant discussion in {ref}`dev-model`. +If the graph partition file has been constructed prior to the ocean model step, +the path to the graph file should be provided in the `graph_target` argument +to the constructor {py:class}`polaris.ocean.model.OceanModelStep()`. + #### YAML files vs. namelists and streams In order to have the same tasks support Omega or MPAS-Ocean, we want