Optionally produce coeffs_reconstruct stream for MPAS-O, link coeffs file from database for Omega

polaris/mesh/base/__init__.py

@@ -41,3 +41,10 @@ def get_base_mesh_steps():
             base_mesh_steps.append(base_mesh_step)
 
     return base_mesh_steps
+
+
+def parse_mesh_filepath(mesh_path):
+    component = mesh_path.split('/')[0]
+    database = '/'.join(mesh_path.split('/')[1:-1])
+    mesh_name = mesh_path.split('/')[-1]
+    return component, database, mesh_name

polaris/ocean/config/coeffs_reconstruct.yaml

@@ -0,0 +1,9 @@
+mpas-ocean:
+  streams:
+    coeffs_reconstruct:
+      type: output
+      filename_template: coeffs_reconstruct.nc
+      output_interval: 0000_00:00:01
+      clobber_mode: truncate
+      contents:
+      - coeffs_reconstruct

polaris/ocean/model/ocean_model_step.py

@@ -221,6 +221,22 @@ def dynamic_model_config(self, at_setup: bool) -> None:
 
         if self.update_eos:
             self.update_namelist_eos()
+        if self.config.has_option('ocean', 'write_coeffs_reconstruct'):
+            self.write_coeffs_reconstruct = self.config.get(
+                'ocean', 'write_coeffs_reconstruct'
+            )
+        else:
+            self.write_coeffs_reconstruct = False
+        if self.write_coeffs_reconstruct:
+            model = self.config.get('ocean', 'model')
+            if not model == 'mpas-ocean':
+                raise ValueError(
+                    'Coefficients for vector reconstruction can only be '
+                    'written for ocean model MPAS-Ocean'
+                )
+            self.add_yaml_file(
+                'polaris.ocean.config', 'coeffs_reconstruct.yaml'
+            )
 
     def constrain_resources(self, available_cores: Dict[str, Any]) -> None:
         """

polaris/ocean/mpas_ocean.cfg

@@ -1,5 +1,11 @@
 # This config file has default config options for MPAS-Ocean
 
+[ocean]
+
+# Whether to write coeffs for reconstructing vector quantities during
+# the forward run
+write_coeffs_reconstruct = True
+
 # The paths section points polaris to external paths
 [paths]
 

polaris/suites/ocean/omega_pr.txt

@@ -1,10 +1,17 @@
-#ocean/planar/barotropic_channel/default  # Supported but fails
-ocean/planar/merry_go_round/default
 ocean/planar/barotropic_gyre/munk/free-slip
 ocean/planar/manufactured_solution/convergence_both/default
 ocean/planar/manufactured_solution/convergence_both/del2
 ocean/planar/manufactured_solution/convergence_both/del4
+ocean/planar/merry_go_round/default
 ocean/spherical/qu/rotation_2d
 ocean/spherical/icos/rotation_2d
 ocean/spherical/icos/cosine_bell/decomp
 ocean/spherical/icos/cosine_bell/restart
+
+# Supported but fails
+#ocean/planar/barotropic_channel/default
+
+# Supported but expected to fail until higher-order advection is available
+#ocean/spherical/icos/nondivergent_2d
+#ocean/spherical/icos/divergent_2d
+#ocean/spherical/icos/correlated_tracers_2d

polaris/tasks/ocean/__init__.py

@@ -4,6 +4,7 @@
 
 import xarray as xr
 from mpas_tools.io import write_netcdf
+from mpas_tools.vector.reconstruct import reconstruct_variable
 from ruamel.yaml import YAML
 
 from polaris import Component
@@ -233,7 +234,14 @@ def map_var_list_from_native_model(self, var_list):
             ]
         return renamed_vars
 
-    def open_model_dataset(self, filename, **kwargs):
+    def open_model_dataset(
+        self,
+        filename,
+        mesh_filename=None,
+        reconstruct_variables=None,
+        coeffs_filename=None,
+        **kwargs,
+    ):
         """
         Open the given dataset, mapping variable and dimension names from Omega
         to MPAS-Ocean names if appropriate
@@ -243,6 +251,15 @@ def open_model_dataset(self, filename, **kwargs):
         filename : str
             The path for the NetCDF file to open
 
+        ds_mesh : xr.Dataset
+            The MPAS mesh dataset for the ocean model.
+
+        reconstruct_variables : list of str
+            List of variable names to reconstruct in the dataset.
+
+        coeffs_reconstruct : xarray.DataArray, optional
+            Coefficients used for reconstructing variables.
+
         kwargs
             keyword arguments passed to `xarray.open_dataset()`
 
@@ -253,6 +270,45 @@ def open_model_dataset(self, filename, **kwargs):
         """
         ds = xr.open_dataset(filename, **kwargs)
         ds = self.map_from_native_model_vars(ds)
+        if reconstruct_variables is None:
+            reconstruct_variables = []
+        for variable in reconstruct_variables:
+            if variable in ds.keys():
+                if 'normal' in variable:
+                    out_var_name = variable.replace('normal', '').lower()
+                else:
+                    out_var_name = variable
+                if (
+                    f'{out_var_name}Zonal' in ds.keys()
+                    and f'{out_var_name}Meridional' in ds.keys()
+                ):
+                    return ds
+                if mesh_filename is None:
+                    raise ValueError(
+                        'mesh_filename must be provided to '
+                        f'open_model_dataset for reconstruction of {variable}'
+                    )
+                ds_mesh = xr.open_dataset(mesh_filename)
+                if coeffs_filename is None:
+                    raise ValueError(
+                        'coeffs_filename must be provided to '
+                        f'open_model_dataset for reconstruction of {variable}'
+                    )
+                ds_coeff = xr.open_dataset(coeffs_filename)
+                coeffs_reconstruct = ds_coeff.coeffs_reconstruct
+                reconstruct_variable(
+                    out_var_name,
+                    ds[variable],
+                    ds_mesh,
+                    coeffs_reconstruct,
+                    ds,
+                    quiet=True,
+                )
+            else:
+                raise ValueError(
+                    f'User requested vector reconstruction for {variable} but '
+                    f"it isn't present in the dataset."
+                )
         return ds
 
     def _has_ocean_io_model_steps(self, tasks) -> Tuple[bool, bool]:

polaris/tasks/ocean/barotropic_gyre/forward.py

@@ -277,6 +277,11 @@ def setup(self):
         # 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')
+            self.add_input_file(
+                target='coeffs.nc',
+                filename='coeffs.nc',
+                database='barotropic_gyre',
+            )
 
     def compute_max_time_step(self, config):
         """

polaris/tasks/ocean/cosine_bell/forward.py

@@ -1,3 +1,4 @@
+from polaris.mesh.base import parse_mesh_filepath
 from polaris.ocean.convergence.spherical import SphericalConvergenceForward
 
 
@@ -72,6 +73,7 @@ def __init__(
             refinement=refinement,
         )
         self.do_restart = do_restart
+        self.mesh_path = mesh.path
 
     def setup(self):
         """
@@ -82,3 +84,12 @@ def setup(self):
         # 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')
+            component, database, mesh_name = parse_mesh_filepath(
+                self.mesh_path
+            )
+            self.add_input_file(
+                filename='coeffs.nc',
+                target=f'{mesh_name}_coeffs.nc',
+                database_component=component,
+                database=database,
+            )

polaris/tasks/ocean/manufactured_solution/forward.py

@@ -89,6 +89,7 @@ def __init__(
             validate_vars=['layerThickness', 'normalVelocity'],
             check_properties=['mass conservation'],
         )
+        self.init = init
         self.del2 = del2
         self.del4 = del4
 
@@ -102,6 +103,12 @@ def setup(self):
         # 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')
+            mesh_name = self.init.path.split('/')[-1]
+            self.add_input_file(
+                target=f'{mesh_name}_coeffs.nc',
+                filename='coeffs.nc',
+                database='manufactured_solution',
+            )
 
     def compute_cell_count(self):
         """

polaris/tasks/ocean/merry_go_round/forward.py

@@ -75,6 +75,7 @@ def __init__(
         )
         self.order = vert_adv_order
         self.limiter = limiter
+        self.mesh_name = init.path.split('/')[-1]
 
     def setup(self):
         """
@@ -86,6 +87,11 @@ def setup(self):
         # 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')
+            self.add_input_file(
+                target=f'{self.mesh_name}_coeffs.nc',
+                filename='coeffs.nc',
+                database='merry_go_round',
+            )
 
     def dynamic_model_config(self, at_setup):
         """

polaris/tasks/ocean/sphere_transport/forward.py

@@ -1,3 +1,4 @@
+from polaris.mesh.base import parse_mesh_filepath
 from polaris.ocean.convergence.spherical import SphericalConvergenceForward
 
 
@@ -78,6 +79,7 @@ def __init__(
             refinement_factor=refinement_factor,
             refinement=refinement,
         )
+        self.mesh_path = base_mesh.path
 
     def setup(self):
         """
@@ -88,3 +90,12 @@ def setup(self):
         # 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')
+            component, database, mesh_name = parse_mesh_filepath(
+                self.mesh_path
+            )
+            self.add_input_file(
+                filename='coeffs.nc',
+                target=f'{mesh_name}_coeffs.nc',
+                database_component=component,
+                database=database,
+            )

polaris/tasks/ocean/sphere_transport/init.py

@@ -192,6 +192,8 @@ def run(self):
         )
         normalVelocity, _ = xr.broadcast(normalVelocity, ds.refZMid)
         ds['normalVelocity'] = normalVelocity.expand_dims(dim='Time', axis=0)
+        ds['velocityZonal'] = np.nan * xr.zeros_like(ds.temperature)
+        ds['velocityMeridional'] = np.nan * xr.zeros_like(ds.temperature)
 
         ds['fCell'] = xr.zeros_like(ds_mesh.xCell)
         ds['fEdge'] = xr.zeros_like(ds_mesh.xEdge)

polaris/tasks/ocean/sphere_transport/sphere_transport.cfg

@@ -138,3 +138,30 @@ norm_type = linear
 
 # A dictionary with keywords for the norm
 norm_args = {'vmin': -0.25, 'vmax': 0.25}
+
+# options for velocity visualization for the sphere transport test case
+[sphere_transport_viz_velocity]
+
+# colormap options
+# colormap
+colormap_name = cmo.balance
+
+# the type of norm used in the colormap
+norm_type = linear
+
+# A dictionary with keywords for the norm
+norm_args = {'vmin': -50., 'vmax': 50.}
+
+
+# options for plotting tracer differences from sphere transport tests
+[sphere_transport_viz_velocity_diff]
+
+# colormap options
+# colormap
+colormap_name = cmo.balance
+
+# the type of norm used in the colormap
+norm_type = linear
+
+# A dictionary with keywords for the norm
+norm_args = {'vmin': -0.1, 'vmax': 0.1}

polaris/tasks/ocean/sphere_transport/viz.py

@@ -57,13 +57,16 @@ def __init__(
         self.add_input_file(
             filename='output.nc', work_dir_target=f'{forward.path}/output.nc'
         )
+        self.forward = forward
         self.mesh_name = mesh_name
         variables_to_plot = dict(
             {
                 'tracer1': 'tracer',
                 'tracer2': 'tracer',
                 'tracer3': 'tracer',
                 'layerThickness': 'h',
+                'velocityZonal': 'velocity',
+                'velocityMeridional': 'velocity',
             }
         )
         self.variables_to_plot = variables_to_plot
@@ -72,6 +75,15 @@ def __init__(
             self.add_output_file(f'{var}_final.png')
             self.add_output_file(f'{var}_diff.png')
 
+    def setup(self):
+        model = self.config.get('ocean', 'model')
+        # TODO: remove as soon as Omega no longer needs this file
+        if model == 'omega':
+            self.add_input_file(
+                filename='coeffs.nc',
+                work_dir_target=f'{self.forward.path}/coeffs.nc',
+            )
+
     def run(self):
         """
         Run this step of the test case
@@ -82,10 +94,22 @@ def run(self):
         run_duration = config.getfloat('convergence_forward', 'run_duration')
 
         variables_to_plot = self.variables_to_plot
-        ds_init = self.open_model_dataset('initial_state.nc')
+        ds_init = self.open_model_dataset(
+            'initial_state.nc',
+            decode_times=False,
+            mesh_filename='mesh.nc',
+            reconstruct_variables=['normalVelocity'],
+            coeffs_filename='coeffs.nc',
+        )
         ds_init = ds_init[variables_to_plot.keys()].isel(Time=0, nVertLevels=0)
 
-        ds_out = self.open_model_dataset('output.nc', decode_times=False)
+        ds_out = self.open_model_dataset(
+            'output.nc',
+            decode_times=False,
+            mesh_filename='mesh.nc',
+            reconstruct_variables=['normalVelocity'],
+            coeffs_filename='coeffs.nc',
+        )
         s_per_hour = 3600.0
 
         # Visualization at halfway around the globe (provided run duration is