From e36f6a35616094ab14108e7f6e45307426ce6025 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Fri, 27 Sep 2024 07:56:44 -0500 Subject: [PATCH 1/4] Add streams_section to PolarisYaml and model step classes This is needed because Omega uses `IOStreams`, whereas MPAS components use `streams`. --- polaris/model_step.py | 31 ++++++++++++++----------- polaris/ocean/model/ocean_model_step.py | 3 ++- polaris/yaml.py | 19 +++++++++++---- 3 files changed, 34 insertions(+), 19 deletions(-) diff --git a/polaris/model_step.py b/polaris/model_step.py index 5fb3ebd54..a0fd402c5 100644 --- a/polaris/model_step.py +++ b/polaris/model_step.py @@ -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='<<>>') @@ -563,7 +566,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 +582,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 +646,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 +732,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,7 +741,8 @@ 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) @@ -773,10 +778,10 @@ 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) 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/yaml.py b/polaris/yaml.py index 0df2264f9..1c0016538 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 From 2ac8fc7cf02214af26d06bbebb427581f8c8d1bc Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Fri, 27 Sep 2024 07:59:29 -0500 Subject: [PATCH 2/4] Add Omega IOStreams to convergence and manufactured_solution --- polaris/ocean/convergence/forward.py | 4 ++ .../tasks/manufactured_solution/forward.py | 11 ----- .../tasks/manufactured_solution/forward.yaml | 42 ++++++++++++++----- 3 files changed, 36 insertions(+), 21 deletions(-) diff --git a/polaris/ocean/convergence/forward.py b/polaris/ocean/convergence/forward.py index 2866853e2..9ff06443b 100644 --- a/polaris/ocean/convergence/forward.py +++ b/polaris/ocean/convergence/forward.py @@ -138,6 +138,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_map = dict([('RK4', 'RungeKutta4')]) model = config.get('ocean', 'model') if model == 'omega': @@ -154,6 +157,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/tasks/manufactured_solution/forward.py b/polaris/ocean/tasks/manufactured_solution/forward.py index 62bd76a4d..fee1db144 100644 --- a/polaris/ocean/tasks/manufactured_solution/forward.py +++ b/polaris/ocean/tasks/manufactured_solution/forward.py @@ -44,17 +44,6 @@ def __init__(self, component, name, resolution, subdir, init): output_filename='output.nc', validate_vars=['layerThickness', 'normalVelocity']) - def setup(self): - """ - TEMP: symlink initial condition to name hard-coded in Omega - """ - super().setup() - config = self.config - model = config.get('ocean', 'model') - # TODO: remove as soon as Omega supports I/O streams - if model == 'omega': - self.add_input_file(filename='OmegaMesh.nc', target='init.nc') - def compute_cell_count(self): """ Compute the approximate number of cells in the mesh, used to constrain diff --git a/polaris/ocean/tasks/manufactured_solution/forward.yaml b/polaris/ocean/tasks/manufactured_solution/forward.yaml index f7aefa138..776359db2 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,31 @@ 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 Dimension: NVertLevels: 1 + IOStreams: + InitialState: + UsePointerFile: false + Filename: init.nc + Mode: read + Precision: double + Freq: 1 + FreqUnits: OnStartup + UseStartEnd: false + Contents: + - Restart + History: + UsePointerFile: false + Filename: output.nc + Mode: write + IfExists: replace + Precision: double + Freq: {{ output_freq }} + FreqUnits: Seconds + UseStartEnd: false + Contents: + - Tracers From 4c5ae6f5256cf10974999d23c35bbb71b253349d Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Sat, 28 Sep 2024 11:37:46 -0500 Subject: [PATCH 3/4] Add streams to yaml processing in model steps --- polaris/model_step.py | 72 +++++++++++++++++++++++++++++++++++++++---- polaris/yaml.py | 6 ---- 2 files changed, 66 insertions(+), 12 deletions(-) diff --git a/polaris/model_step.py b/polaris/model_step.py index a0fd402c5..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 @@ -278,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 @@ -304,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 @@ -328,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 @@ -443,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) @@ -484,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) @@ -749,7 +772,7 @@ def _process_yaml_streams(self, yaml_filename, package, replacements, 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``. @@ -764,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 ' @@ -785,7 +810,42 @@ def _process_yaml(self, quiet): 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/yaml.py b/polaris/yaml.py index 1c0016538..ded7bd4fa 100644 --- a/polaris/yaml.py +++ b/polaris/yaml.py @@ -153,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, From 60ddce2e227c672c7e313e6c89cb32605c596ba4 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Mon, 30 Sep 2024 04:12:39 -0500 Subject: [PATCH 4/4] Add back symlink OmegaMesh.nc, still required --- polaris/ocean/tasks/manufactured_solution/forward.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/polaris/ocean/tasks/manufactured_solution/forward.py b/polaris/ocean/tasks/manufactured_solution/forward.py index fee1db144..a57eeb9f2 100644 --- a/polaris/ocean/tasks/manufactured_solution/forward.py +++ b/polaris/ocean/tasks/manufactured_solution/forward.py @@ -44,6 +44,17 @@ def __init__(self, component, name, resolution, subdir, init): output_filename='output.nc', validate_vars=['layerThickness', 'normalVelocity']) + def setup(self): + """ + TEMP: symlink initial condition to name hard-coded in Omega + """ + super().setup() + config = self.config + model = config.get('ocean', 'model') + # 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') + def compute_cell_count(self): """ Compute the approximate number of cells in the mesh, used to constrain