Updates related to Omega IOStreams

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='<<<model>>>')
 
@@ -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',

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,

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}')
 

polaris/ocean/tasks/manufactured_solution/forward.py

@@ -51,7 +51,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')
 

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

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,20 +145,14 @@ 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
 
         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,