Add ocean model conservation checks

docs/developers_guide/api.md

@@ -161,6 +161,7 @@ seaice/api
    Step.add_output_file
    Step.add_dependency
    Step.validate_baselines
+   Step.check_properties
    Step.set_shared_config
 ```
 

docs/developers_guide/framework/validation.md

@@ -211,3 +211,51 @@ which need to have their variables renamed to the MPAS-Ocean names for use in
 Polaris.  The `ds1` and `ds2` keyword arguments are used to supply datasets
 corresponding to `filename1` and `filename2`, respectively, in such
 circumstances.
+
+# Property checks
+
+For some output files, you may wish to run checks of certain properties such as
+conservation of mass or energy. Currently, only conservation checks for the
+ocean are available.
+
+To run property checks, add a list of properties
+the keyword argument `verify_properties` to the
+:py:meth:`polaris.Step.add_output_file()` method.  As an example:
+
+```python
+from polaris import Step
+
+
+class Forward(OceanModelStep):
+    def __init__(self, task):
+        super().__init__(task=task, name='forward')
+
+        self.add_input_file(
+            filename='mesh.nc', target='../init/culled_mesh.nc'
+        )
+        self.add_output_file('output.nc',
+                             verify_properties=['mass conservation'])
+```
+
+## Conservation checks
+
+Mass, salt and energy conservation checks are available for ocean model output.
+The checks will fail if the relative change in the total quantity exceeds the
+tolerance given in
+
+```cfg
+# Options related the ocean component
+[ocean]
+
+# Tolerance for mass conservation, normalized by total mass
+mass_conservation_tolerance = 1e-8
+
+# Tolerance for salt conservation, normalized by total salt
+salt_conservation_tolerance = 1e-8
+
+# Tolerance for thermal energy conservation, normalized by total energy
+energy_conservation_tolerance = 1e-8
+```
+
+As shown in the previous example, we have added a mesh file with the name
+'mesh.nc' because conservation checks require the area of cells.

docs/developers_guide/ocean/api.md

@@ -461,6 +461,20 @@
 
 ## Ocean Framework
 
+### Conservation utilities
+
+```{eval-rst}
+.. currentmodule:: polaris.ocean.conservation
+
+.. autosummary::
+   :toctree: generated/
+
+   compute_total_mass
+   compute_total_mass_nonboussinesq
+   compute_total_salt
+   compute_total_energy
+```
+
 ### Convergence Tests
 
 ```{eval-rst}
@@ -533,6 +547,7 @@
 
    OceanModelStep
    OceanModelStep.setup
+   OceanModelStep.check_properties
    OceanModelStep.constrain_resources
    OceanModelStep.compute_cell_count
    OceanModelStep.map_yaml_options

docs/developers_guide/ocean/tasks/barotropic_channel.md

@@ -32,7 +32,7 @@ stream is also generated with wind stress fields.
 The class {py:class}`polaris.tasks.ocean.barotropic_channel.forward.Forward`
 defines a step for running the ocean from the initial condition produced in
 the `init` step. Namelist and streams files are updated in
-{py:meth}`polaris.tasks.ocean.overflow.forward.Forward.dynamic_model_config()`.
+{py:meth}`polaris.tasks.ocean.barotropic_channel.forward.Forward.dynamic_model_config()`.
 The number of cells is approximated from config options in
 {py:meth}`polaris.tasks.ocean.barotropic_channel.forward.Forward.compute_cell_count()`
 so that this can be used to constrain the number of MPI tasks that Polaris
@@ -47,7 +47,7 @@ The {py:class}`polaris.tasks.ocean.barotropic_channel.viz.Viz` plots the initial
 final velocity components, relative vorticity, and circulation at the bottommost
 vertical layer.
 
-(dev-ocean-overflow-default)=
+(dev-ocean-barotropic-channel-default)=
 
 ## default
 

docs/developers_guide/ocean/tasks/customizable_viz.md

@@ -1,4 +1,4 @@
-(dev-ocean-customizable_viz)=
+(dev-ocean-customizable-viz)=
 
 # customizable_viz
 
@@ -10,7 +10,7 @@ controlled via the task configuration.
 ## framework
 
 The config options for the `customizable_viz` tests are described in
-{ref}`ocean-customizable_viz` in the User's Guide.
+{ref}`ocean-customizable-viz` in the User's Guide.
 
 The test also makes use of `polaris.viz.helper` and `polaris.viz.mplstyle`.
 

docs/tutorials/dev_add_category_of_tasks/adding_forward.md

@@ -139,6 +139,9 @@ class Forward(OceanModelStep):
                 'normalVelocity',
                 'temperature',
             ],
+            check_properties=[
+                'mass conservation',
+            ],
         )
 ```
 

polaris/ocean/conservation.py

@@ -0,0 +1,72 @@
+from mpas_tools.cime.constants import constants
+
+# Should match config_density0
+rho_sw = 1026.0
+
+cp_sw = constants['SHR_CONST_CPSW']
+
+
+def compute_total_mass(ds_mesh, ds):
+    """
+    Compute the total mass in an ocean model output file using a constant
+    density
+    """
+    ds = _reduce_dataset_time_dim(ds)
+    area_cell = ds_mesh.areaCell
+    layer_thickness = ds.layerThickness
+    total_mass = rho_sw * (
+        area_cell * layer_thickness.sum(dim='nVertLevels')
+    ).sum(dim='nCells')
+    return total_mass
+
+
+def compute_total_energy(ds_mesh, ds):
+    """
+    Compute the total heat content an ocean model output file
+    """
+    ds = _reduce_dataset_time_dim(ds)
+    area_cell = ds_mesh.areaCell
+    layer_thickness = ds.layerThickness
+    temperature = ds.temperature
+    total_energy = (
+        rho_sw
+        * cp_sw
+        * (
+            area_cell * (layer_thickness * temperature).sum(dim='nVertLevels')
+        ).sum(dim='nCells')
+    )
+    return total_energy
+
+
+def compute_total_tracer(ds_mesh, ds, tracer_name='tracer1'):
+    """
+    Compute the total salt in an ocean model output file
+    """
+    ds = _reduce_dataset_time_dim(ds)
+    area_cell = ds_mesh.areaCell
+    layer_thickness = ds.layerThickness
+    tracer = ds[tracer_name]
+    total_tracer = (
+        area_cell * (layer_thickness * tracer).sum(dim='nVertLevels')
+    ).sum(dim='nCells')
+    return total_tracer
+
+
+def compute_total_salt(ds_mesh, ds):
+    """
+    Compute the total salt in an ocean model output file
+    """
+    return compute_total_tracer(ds_mesh, ds, tracer_name='salinity')
+
+
+def _reduce_dataset_time_dim(ds):
+    for time_dim in ['time', 'Time']:
+        if time_dim in ds.dims:
+            if ds.sizes[time_dim] > 1:
+                raise ValueError(
+                    'compute_total_energy is designed to work on a dataset '
+                    'with one time slice'
+                )
+            else:
+                ds = ds.isel(time_dim=0)
+    return ds

polaris/ocean/convergence/forward.py

@@ -38,6 +38,7 @@ def __init__(
         graph_target=None,
         output_filename='output.nc',
         validate_vars=None,
+        check_properties=None,
         refinement='both',
     ):
         """
@@ -108,9 +109,14 @@ def __init__(
         self.add_input_file(
             filename='init.nc', work_dir_target=f'{init.path}/initial_state.nc'
         )
+        self.add_input_file(
+            filename='mesh.nc', work_dir_target=f'{init.path}/mesh.nc'
+        )
 
         self.add_output_file(
-            filename=output_filename, validate_vars=validate_vars
+            filename=output_filename,
+            validate_vars=validate_vars,
+            check_properties=check_properties,
         )
 
     def compute_cell_count(self):

polaris/ocean/model/ocean_model_step.py

@@ -1,10 +1,18 @@
 import importlib.resources as imp_res
+import os
 from types import ModuleType
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 from ruamel.yaml import YAML
 
 from polaris.model_step import ModelStep
+from polaris.ocean.conservation import (
+    compute_total_energy,
+    compute_total_mass,
+    # compute_total_mass_nonboussinesq, # Add when Omega EOS is used
+    compute_total_salt,
+    compute_total_tracer,
+)
 from polaris.tasks.ocean import Ocean
 
 OptionValue = Union[str, int, float, bool]
@@ -362,6 +370,126 @@ def update_namelist_eos(self) -> None:
             options=replacements, config_model='ocean'
         )
 
+    def check_properties(self):
+        checked = False
+        success = True
+        if self.work_dir is None:
+            raise ValueError(
+                'The work directory must be set before the step '
+                'output properties can be checked.'
+            )
+        passed_properties = []
+        failed_properties = []
+        for filename, properties in self.properties_to_check.items():
+            filename = str(filename)
+            mesh_filename = os.path.join(self.work_dir, 'mesh.nc')
+            this_filename = os.path.join(self.work_dir, filename)
+            ds_mesh = self.component.open_model_dataset(mesh_filename)
+            ds = self.component.open_model_dataset(this_filename)
+            if 'tracer conservation' in properties:
+                # All tracers in mpaso_to_omega.yaml
+                tracers_to_check = [
+                    'temperature',
+                    'salinity',
+                    'tracer1',
+                    'tracer2',
+                    'tracer3',
+                ]
+                # Expand 'tracer conservation' into list of tracers to check
+                properties = [
+                    item
+                    for item in properties
+                    if item != 'tracer conservation'
+                ]
+                for tracer in tracers_to_check:
+                    if tracer in ds.keys():
+                        properties.append(f'tracer conservation-{tracer}')
+            for output_property in properties:
+                if output_property == 'mass conservation':
+                    tol = self.config.getfloat(
+                        'ocean', 'mass_conservation_tolerance'
+                    )
+                    # Add when Omega EOS is used
+                    # if self.config.get('ocean', 'model') == 'omega':
+                    # relative_error  = self._compute_rel_err(
+                    #    compute_total_mass_nonboussinesq, ds_mesh, ds
+                    # )
+                    # else:
+                    relative_error = self._compute_rel_err(
+                        compute_total_mass, ds_mesh=ds_mesh, ds=ds
+                    )
+                elif output_property == 'salt conservation':
+                    tol = self.config.getfloat(
+                        'ocean', 'salt_conservation_tolerance'
+                    )
+                    relative_error = self._compute_rel_err(
+                        compute_total_mass, ds_mesh, ds
+                    )
+                    relative_error = self._compute_rel_err(
+                        compute_total_salt, ds_mesh, ds
+                    )
+                elif output_property.split('-')[0] == 'tracer conservation':
+                    tol = self.config.getfloat(
+                        'ocean', 'tracer_conservation_tolerance'
+                    )
+                    tracer = output_property.split('-')[1]
+                    relative_error = self._compute_rel_err(
+                        compute_total_tracer,
+                        ds_mesh,
+                        ds,
+                        tracer_name=tracer,
+                    )
+                elif output_property == 'energy conservation':
+                    tol = self.config.getfloat(
+                        'ocean', 'energy_conservation_tolerance'
+                    )
+                    relative_error = self._compute_rel_err(
+                        compute_total_energy, ds_mesh, ds
+                    )
+                else:
+                    raise ValueError(
+                        'Could not find method to execute property check '
+                        f'{output_property}'
+                    )
+
+                result = relative_error < tol
+                success = success and result
+                checked = True
+                # We already appended log strings for tracer conservation
+                if output_property != 'tracer conservation':
+                    if not result:
+                        failed_properties.append(
+                            f'{output_property} relative error '
+                            f'{relative_error:.3e} exceeds {tol}'
+                        )
+                    else:
+                        passed_properties.append(
+                            f'{output_property} relative error '
+                            f'{relative_error:.3e}'
+                        )
+        if checked and success:
+            log_filename = os.path.join(
+                self.work_dir, 'property_check_passed.log'
+            )
+            passed_properties_str = '\n  '.join(passed_properties)
+            with open(log_filename, 'w') as result_log_file:
+                result_log_file.write(
+                    f'Output file {filename} passed property checks.\n'
+                    f'{passed_properties_str}\n'
+                )
+        elif checked and not success:
+            log_filename = os.path.join(
+                self.work_dir, 'property_check_failed.log'
+            )
+            failed_properties_str = '\n  '.join(failed_properties)
+            with open(log_filename, 'w') as result_log_file:
+                result_log_file.write(
+                    f'Property checks on {filename} failed for:\n '
+                    f'{failed_properties_str}\n'
+                )
+
+        return checked, success
+
     def _update_ntasks(self) -> None:
         """
         Update ``ntasks`` and ``min_tasks`` for the step based on the estimated
@@ -546,6 +674,20 @@ def _map_mpaso_to_omega_section_option(
 
         return out_sections, out_option, out_value
 
+    def _compute_rel_err(
+        self,
+        func,
+        ds_mesh,
+        ds,
+        time_index_start=0,
+        time_index_end=-1,
+        **kwargs,
+    ):
+        init_val = func(ds_mesh, ds.isel(Time=time_index_start), **kwargs)
+        final_val = func(ds_mesh, ds.isel(Time=time_index_end), **kwargs)
+        val_change = final_val - init_val
+        return abs(val_change) / (final_val + 1.0)
+
     @staticmethod
     def _warn_not_found(not_found: List[str]) -> None:
         """Warn about options that were not found in the map"""