diff --git a/e3sm_submodules/Omega b/e3sm_submodules/Omega index 8247b3712..399b7af25 160000 --- a/e3sm_submodules/Omega +++ b/e3sm_submodules/Omega @@ -1 +1 @@ -Subproject commit 8247b3712c50d90da5e7ce68dbc7284ce3445b76 +Subproject commit 399b7af2548979852262fd15950d3bb93a80b671 diff --git a/polaris/model_step.py b/polaris/model_step.py index 5fb3ebd54..c9d1ddf08 100644 --- a/polaris/model_step.py +++ b/polaris/model_step.py @@ -1,7 +1,7 @@ import os import shutil from collections import OrderedDict -from typing import List, Union +from typing import Dict, List, Union import numpy as np import xarray as xr @@ -69,6 +69,8 @@ class ModelStep(Step): Whether to create a yaml file with model config options and streams instead of MPAS namelist and streams files + streams_section : str + The name of the streams section in yaml files """ def __init__(self, component, name, subdir=None, indir=None, ntasks=None, min_tasks=None, openmp_threads=None, max_memory=None, @@ -173,6 +175,7 @@ def __init__(self, component, name, subdir=None, indir=None, ntasks=None, self.graph_filename = graph_filename self.make_yaml = make_yaml + self.streams_section = 'streams' self.add_input_file(filename='<<>>') @@ -275,7 +278,7 @@ def add_yaml_file(self, package, yaml, template_replacements=None): def map_yaml_options(self, options, config_model): """ - A mapping between model config options between different models. This + A mapping between model config options from different models. This method should be overridden for situations in which yaml config options have diverged in name or structure from their counterparts in another model (e.g. when translating from MPAS-Ocean namelist options @@ -301,7 +304,7 @@ def map_yaml_options(self, options, config_model): def map_yaml_configs(self, configs, config_model): """ - A mapping between model config options between different models. This + A mapping between model config options from different models. This method should be overridden for situations in which yaml config options have diverged in name or structure from their counterparts in another model (e.g. when translating from MPAS-Ocean namelist options @@ -325,6 +328,29 @@ def map_yaml_configs(self, configs, config_model): """ return configs + def map_yaml_streams(self, streams, config_model): + """ + A mapping between model streams from different models. This method + should be overridden for situations in which yaml streams have diverged + in name or structure from their counterparts in another model (e.g. + when translating from MPAS-Ocean streams to Omega IOStreams) + + Parameters + ---------- + streams : dict + A nested dictionary of streams data + + config_model : str or None + If streams are available for multiple models, the model that the + streams are from + + Returns + ------- + configs : dict + A revised nested dictionary of streams data + """ + return streams + def map_yaml_to_namelist(self, options): """ A mapping from yaml model config options to namelist options. This @@ -440,7 +466,7 @@ def runtime_setup(self): self.dynamic_model_config(at_setup=False) if self.make_yaml: - self._process_yaml(quiet=quiet) + self._process_yaml(quiet=quiet, remove_unrequested_streams=False) else: self._process_namelists(quiet=quiet) self._process_streams(quiet=quiet, remove_unrequested=False) @@ -481,7 +507,7 @@ def process_inputs_and_outputs(self): self._create_model_config() if self.make_yaml: - self._process_yaml(quiet=quiet) + self._process_yaml(quiet=quiet, remove_unrequested_streams=True) else: self._process_namelists(quiet=quiet) self._process_streams(quiet=quiet, remove_unrequested=True) @@ -563,7 +589,8 @@ def _create_model_config(self): config = self.config if self.make_yaml: defaults_filename = config.get('model_config', 'defaults') - self._yaml = PolarisYaml.read(defaults_filename) + self._yaml = PolarisYaml.read(defaults_filename, + streams_section=self.streams_section) else: defaults_filename = config.get('namelists', 'forward') self._namelist = polaris.namelist.ingest(defaults_filename) @@ -578,7 +605,8 @@ def _read_model_config(self): """ if self.make_yaml: filename = os.path.join(self.work_dir, self.yaml) - self._yaml = PolarisYaml.read(filename) + self._yaml = PolarisYaml.read(filename, + streams_section=self.streams_section) else: filename = os.path.join(self.work_dir, self.namelist) self._namelist = polaris.namelist.ingest(filename) @@ -641,10 +669,10 @@ def _process_namelists(self, quiet): options = self.map_yaml_to_namelist(options) replacements.update(options) if 'yaml' in entry: - yaml = PolarisYaml.read(filename=entry['yaml'], - package=entry['package'], - replacements=entry['replacements'], - model=config_model) + yaml = PolarisYaml.read( + filename=entry['yaml'], package=entry['package'], + replacements=entry['replacements'], model=config_model, + streams_section=self.streams_section) configs = self.map_yaml_configs(configs=yaml.configs, config_model=config_model) @@ -727,8 +755,7 @@ def _process_streams(self, quiet, remove_unrequested): if not found: defaults.remove(default) - @staticmethod - def _process_yaml_streams(yaml_filename, package, replacements, + def _process_yaml_streams(self, yaml_filename, package, replacements, config_model, processed_registry_filename, tree, quiet): if not quiet: @@ -737,14 +764,15 @@ def _process_yaml_streams(yaml_filename, package, replacements, yaml = PolarisYaml.read(filename=yaml_filename, package=package, replacements=replacements, - model=config_model) + model=config_model, + streams_section=self.streams_section) assert processed_registry_filename is not None new_tree = yaml_to_mpas_streams( processed_registry_filename, yaml) tree = polaris.streams.update_tree(tree, new_tree) return tree - def _process_yaml(self, quiet): + def _process_yaml(self, quiet, remove_unrequested_streams): """ Processes changes to a yaml file from the files and dictionaries in the step's ``model_config_data``. @@ -759,6 +787,8 @@ def _process_yaml(self, quiet): if not quiet: print(f'Warning: replacing yaml options in {self.yaml}') + streams: Dict[str, Dict[str, Union[str, float, int, List[str]]]] = {} + for entry in self.model_config_data: if 'namelist' in entry: raise ValueError('Cannot generate a yaml config from an MPAS ' @@ -773,14 +803,49 @@ def _process_yaml(self, quiet): config_model=config_model) self._yaml.update(options=options, quiet=quiet) if 'yaml' in entry: - yaml = PolarisYaml.read(filename=entry['yaml'], - package=entry['package'], - replacements=entry['replacements'], - model=config_model) + yaml = PolarisYaml.read( + filename=entry['yaml'], package=entry['package'], + replacements=entry['replacements'], model=config_model, + streams_section=self.streams_section) configs = self.map_yaml_configs(configs=yaml.configs, config_model=config_model) + new_streams = self.map_yaml_streams( + streams=yaml.streams, config_model=config_model) + self._update_yaml_streams(streams, new_streams, + quiet=quiet, + remove_unrequested=False) self._yaml.update(configs=configs, quiet=quiet) + self._update_yaml_streams( + self._yaml.streams, streams, quiet=quiet, + remove_unrequested=remove_unrequested_streams) + + @staticmethod + def _update_yaml_streams(streams, new_streams, quiet, remove_unrequested): + """ + Update yaml streams, optionally removing any streams that aren't in + new_streams + """ + + for stream_name, new_stream in new_streams.items(): + if stream_name in streams: + streams[stream_name].update(new_stream) + if not quiet: + print(f' updating: {stream_name}') + else: + if not quiet: + print(f' adding: {stream_name}') + streams[stream_name] = new_stream + + if remove_unrequested: + # during setup, we remove any default streams that aren't requested + # but at runtime we don't want to do this because we would lose any + # streams added only during setup. + for stream_name in list(streams.keys()): + if stream_name not in new_streams: + if not quiet: + print(f' dropping: {stream_name}') + streams.pop(stream_name) def make_graph_file(mesh_filename, graph_filename='graph.info', diff --git a/polaris/ocean/convergence/forward.py b/polaris/ocean/convergence/forward.py index 7a7ecbc2e..fff8f18bf 100644 --- a/polaris/ocean/convergence/forward.py +++ b/polaris/ocean/convergence/forward.py @@ -141,6 +141,9 @@ def dynamic_model_config(self, at_setup): output_interval_str = get_time_interval_string( seconds=output_interval * s_per_hour) + # For Omega, we want the output interval as a number of seconds + output_freq = int(output_interval * s_per_hour) + time_integrator = section.get('time_integrator') time_integrator_map = dict([('RK4', 'RungeKutta4')]) model = config.get('ocean', 'model') @@ -158,6 +161,7 @@ def dynamic_model_config(self, at_setup): btr_dt=btr_dt_str, run_duration=run_duration_str, output_interval=output_interval_str, + output_freq=f'{output_freq}' ) self.add_yaml_file(self.package, self.yaml_filename, diff --git a/polaris/ocean/model/mpaso_to_omega.yaml b/polaris/ocean/model/mpaso_to_omega.yaml index dbb7baa1c..0eea1c772 100644 --- a/polaris/ocean/model/mpaso_to_omega.yaml +++ b/polaris/ocean/model/mpaso_to_omega.yaml @@ -1,14 +1,33 @@ +dimensions: + Time: time + nCells: NCells + nEdges: NEdges + nVertices: NVertices + maxEdges: MaxEdges + maxEdges2: MaxEdges2 + TWO: MaxCellsOnEdge + vertexDegree: VertexDegree + nVertLevels: NVertLevels + nVertLevelsP1: NVertLevelsP1 + variables: - temperature: Temp - salinity: Salt + # tracers + temperature: Temperature + salinity: Salinity tracer1: Debug1 tracer2: Debug2 tracer3: Debug3 - ssh: SshCell + + # state + layerThickness: LayerThickness + normalVelocity: NormalVelocity + + # auxiliary state + ssh: SshCellDefault config: - section: - time_management: TimeManagement + time_management: TimeIntegration options: config_start_time: StartTime config_stop_time: StopTime @@ -47,3 +66,10 @@ config: options: config_use_mom_del4: VelHyperDiffTendencyEnable config_mom_del4: ViscDel4 + +- section: + manufactured_solution: ManufacturedSolution + options: + config_manufactured_solution_wavelength_x: WavelengthX + config_manufactured_solution_wavelength_y: WavelengthY + config_manufactured_solution_amplitude: Amplitude diff --git a/polaris/ocean/model/ocean_io_step.py b/polaris/ocean/model/ocean_io_step.py index 733258c86..65664f43a 100644 --- a/polaris/ocean/model/ocean_io_step.py +++ b/polaris/ocean/model/ocean_io_step.py @@ -14,6 +14,13 @@ class OceanIOStep(Step): 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 @@ -39,6 +46,8 @@ def __init__(self, component, name, **kwargs): 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 @@ -57,9 +66,9 @@ def setup(self): 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 - datasets like an initial condition) + 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 ---------- @@ -75,16 +84,20 @@ def map_to_native_model_vars(self, ds): 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 = {k: v for k, v in self.mpaso_to_omega_var_map.items() - if k in ds} + 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 write_model_dataset(self, ds, filename): """ - Write out the given dataset, mapping variable names from MPAS-Ocean - to Omega names if appropriate + Write out the given dataset, mapping dimension and variable names from + MPAS-Ocean to Omega names if appropriate Parameters ---------- @@ -99,9 +112,9 @@ def write_model_dataset(self, ds, filename): 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 - that are output from the model) + 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 ---------- @@ -116,16 +129,20 @@ def map_from_native_model_vars(self, ds): 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 = {k: v for k, v in self.omega_to_mpaso_var_map.items() - if k in ds} + 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 def open_model_dataset(self, filename, **kwargs): """ - Open the given dataset, mapping variable names from Omega to MPAS-Ocean - names if appropriate + Open the given dataset, mapping variable and dimension names from Omega + to MPAS-Ocean names if appropriate Parameters ---------- @@ -146,7 +163,7 @@ def open_model_dataset(self, filename, **kwargs): def _read_var_map(self): """ - Read the map from MPAS-Ocean to Omega config options + Read the map from MPAS-Ocean to Omega dimension and variable names """ package = 'polaris.ocean.model' filename = 'mpaso_to_omega.yaml' @@ -154,7 +171,11 @@ def _read_var_map(self): 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()} diff --git a/polaris/ocean/model/ocean_model_step.py b/polaris/ocean/model/ocean_model_step.py index 5c70269e0..d1f846980 100644 --- a/polaris/ocean/model/ocean_model_step.py +++ b/polaris/ocean/model/ocean_model_step.py @@ -119,6 +119,7 @@ def setup(self): self.make_yaml = True self.config_models = ['ocean', 'Omega'] self.yaml = 'omega.yml' + self.streams_section = 'IOStreams' self._read_config_map() self.partition_graph = False elif model == 'mpas-ocean': @@ -127,7 +128,7 @@ def setup(self): self.add_input_file( filename='graph.info', work_dir_target=self.graph_target) - + self.streams_section = 'streams' else: raise ValueError(f'Unexpected ocean model: {model}') diff --git a/polaris/ocean/tasks/manufactured_solution/forward.py b/polaris/ocean/tasks/manufactured_solution/forward.py index 012bea98f..2b7b9882b 100644 --- a/polaris/ocean/tasks/manufactured_solution/forward.py +++ b/polaris/ocean/tasks/manufactured_solution/forward.py @@ -65,7 +65,7 @@ def setup(self): super().setup() config = self.config model = config.get('ocean', 'model') - # TODO: remove as soon as Omega supports I/O streams + # TODO: remove as soon as Omega no longer hard-codes this file if model == 'omega': self.add_input_file(filename='OmegaMesh.nc', target='init.nc') diff --git a/polaris/ocean/tasks/manufactured_solution/forward.yaml b/polaris/ocean/tasks/manufactured_solution/forward.yaml index f7aefa138..4f0073794 100644 --- a/polaris/ocean/tasks/manufactured_solution/forward.yaml +++ b/polaris/ocean/tasks/manufactured_solution/forward.yaml @@ -5,6 +5,16 @@ ocean: time_integration: config_dt: {{ dt }} config_time_integrator: {{ time_integrator }} + +mpas-ocean: + bottom_drag: + config_bottom_drag_mode: implicit + config_implicit_bottom_drag_type: constant + config_implicit_constant_bottom_drag_coeff: 0.0 + manufactured_solution: + config_use_manufactured_solution: true + debug: + config_disable_vel_hmix: true streams: mesh: filename_template: init.nc @@ -22,19 +32,27 @@ ocean: - layerThickness - ssh -mpas-ocean: - bottom_drag: - config_bottom_drag_mode: implicit - config_implicit_bottom_drag_type: constant - config_implicit_constant_bottom_drag_coeff: 0.0 - manufactured_solution: - config_use_manufactured_solution: true - debug: - config_disable_vel_hmix: true - Omega: Tendencies: VelDiffTendencyEnable: false VelHyperDiffTendencyEnable: false + UseCustomTendency: true + ManufacturedSolutionTendency: true Dimension: NVertLevels: 1 + IOStreams: + InitialState: + UsePointerFile: false + Filename: init.nc + Contents: + - NormalVelocity + - LayerThickness + History: + Filename: output.nc + Freq: {{ output_freq }} + FreqUnits: Seconds + Contents: + - NormalVelocity + - LayerThickness + - SshCellDefault + RestartRead: {} diff --git a/polaris/ocean/tasks/manufactured_solution/manufactured_solution.cfg b/polaris/ocean/tasks/manufactured_solution/manufactured_solution.cfg index a4edc7e49..c6b270211 100644 --- a/polaris/ocean/tasks/manufactured_solution/manufactured_solution.cfg +++ b/polaris/ocean/tasks/manufactured_solution/manufactured_solution.cfg @@ -77,8 +77,9 @@ refinement_factors_time = 1., 0.5, 0.25 [convergence_forward] # time integrator -# mpas-ocean: {'split_explicit', 'RK4'} -# omega: {'Forward-Backward', 'RungeKutta4', 'RungeKutta2'} +# either: {'RK4'} +# mpas-ocean: {'split_explicit'} +# omega: {'Forward-Backward', 'RungeKutta2'} time_integrator = RK4 # RK4 time step per resolution (s/km), since dt is proportional to resolution diff --git a/polaris/yaml.py b/polaris/yaml.py index 0df2264f9..ded7bd4fa 100644 --- a/polaris/yaml.py +++ b/polaris/yaml.py @@ -21,6 +21,9 @@ class PolarisYaml: streams : dict Nested dictionaries containing data about streams + streams_section : str + The name of the streams section + model : str The name of the E3SM component """ @@ -30,11 +33,13 @@ def __init__(self): Create a yaml config object """ self.configs = dict() + self.streams_section = 'streams' self.streams = dict() self.model = None @classmethod - def read(cls, filename, package=None, replacements=None, model=None): + def read(cls, filename, package=None, replacements=None, model=None, + streams_section='streams'): """ Add config options from a yaml file @@ -55,6 +60,9 @@ def read(cls, filename, package=None, replacements=None, model=None): The name of the model to parse if the yaml file might have multiple models + streams_section : str, optional + The name of the streams section + Returns ------- yaml : polaris.yaml.PolarisYaml @@ -74,6 +82,7 @@ def read(cls, filename, package=None, replacements=None, model=None): text = template.render(**replacements) yaml = cls() + yaml.streams_section = streams_section yaml_data = YAML(typ='rt') configs = yaml_data.load(text) @@ -89,10 +98,10 @@ def read(cls, filename, package=None, replacements=None, model=None): yaml.streams = {} if model in configs: configs = configs[model] - if 'streams' in configs: - yaml.streams = configs['streams'] + if streams_section in configs: + yaml.streams = configs[streams_section] configs = dict(configs) - configs.pop('streams') + configs.pop(streams_section) else: configs = {} @@ -136,7 +145,7 @@ def write(self, filename): yaml = YAML(typ='rt') configs = dict(self.configs) if self.streams: - configs['streams'] = self.streams + configs[self.streams_section] = self.streams model_configs = dict() model_configs[self.model] = configs @@ -144,12 +153,6 @@ def write(self, filename): with open(filename, 'w') as outfile: yaml.dump(model_configs, outfile) - def _add_stream(self, stream_name, stream): - """ - Add stream from a dictionary - """ - self.streams[stream_name] = stream - def mpas_namelist_and_streams_to_yaml(model, namelist_template=None, namelist=None,