Add pseudothickness to write_model_dataset for Omega

polaris/ocean/convergence/analysis.py

@@ -426,7 +426,9 @@ def exact_solution(self, refinement_factor, field_name, time, zidx=None):
             The exact solution as derived from the initial condition
         """
 
-        ds_init = self.open_model_dataset(f'init_r{refinement_factor:02g}.nc')
+        ds_init = self.open_model_dataset(
+            f'init_r{refinement_factor:02g}.nc', self.config
+        )
         ds_init = ds_init.isel(Time=0)
         if zidx is not None:
             ds_init = ds_init.isel(nVertLevels=zidx)
@@ -458,7 +460,7 @@ def get_output_field(self, refinement_factor, field_name, time, zidx=None):
         """
         config = self.config
         ds_out = self.open_model_dataset(
-            f'output_r{refinement_factor:02g}.nc', decode_times=False
+            f'output_r{refinement_factor:02g}.nc', config, decode_times=False
         )
 
         model = config.get('ocean', 'model')

polaris/ocean/model/ocean_io_step.py

@@ -32,7 +32,7 @@ def map_to_native_model_vars(self, ds):
         """
         return self.component.map_to_native_model_vars(ds)
 
-    def write_model_dataset(self, ds, filename):
+    def write_model_dataset(self, ds, filename, config=None):
         """
         Write out the given dataset, mapping dimension and variable names from
         MPAS-Ocean to Omega names if appropriate
@@ -45,7 +45,7 @@ def write_model_dataset(self, ds, filename):
         filename : str
             The path for the NetCDF file to write
         """
-        self.component.write_model_dataset(ds, filename)
+        self.component.write_model_dataset(ds, filename, config=config)
 
     def map_from_native_model_vars(self, ds):
         """
@@ -65,7 +65,7 @@ def map_from_native_model_vars(self, ds):
         """
         return self.component.map_from_native_model_vars(ds)
 
-    def open_model_dataset(self, filename, **kwargs):
+    def open_model_dataset(self, filename, config=None, **kwargs):
         """
         Open the given dataset, mapping variable and dimension names from Omega
         to MPAS-Ocean names if appropriate
@@ -83,4 +83,6 @@ def open_model_dataset(self, filename, **kwargs):
         ds : xarray.Dataset
             The dataset with variables named as expected in MPAS-Ocean
         """
-        return self.component.open_model_dataset(filename, **kwargs)
+        return self.component.open_model_dataset(
+            filename, config=config, **kwargs
+        )

polaris/ocean/ocean.cfg

@@ -55,6 +55,9 @@ iterations = 10
 # Options related to the vertical grid
 [vertical_grid]
 
+# The initial surface pressure when not spatially-varying in Pascals
+surface_pressure = 0.
+
 # The minimum number of vertical levels for z-star coordinate
 min_vert_levels = 1
 

polaris/ocean/vertical/diagnostics.py

@@ -1,6 +1,52 @@
 import numpy as np
 import xarray as xr
 
+from polaris.constants import get_constant
+from polaris.ocean.vertical.ztilde import (
+    pressure_and_spec_vol_from_state_at_geom_height,
+    pseudothickness_from_pressure,
+)
+
+RhoSw = get_constant('seawater_density_reference')
+
+
+def geom_thickness_from_ds(ds, config):
+    if 'layerThickness' in ds.keys():
+        return ds['layerThickness']
+    elif 'SpecVol' in ds.keys() and 'PseudoThickness' in ds.keys():
+        return RhoSw * ds['SpecVol'] * ds['layerThickness']
+    else:
+        raise ValueError(
+            'Geometric layerThickness is not present in the '
+            'initial condition and SpecVol is not present '
+            'to compute it'
+        )
+
+
+def pseudothickness_from_ds(ds, config):
+    if 'temperature' not in ds.keys() or 'salinity' not in ds.keys():
+        print(
+            'PseudoThickness is not present in the '
+            'initial condition and T,S are not present '
+            'to compute it'
+        )
+        return
+
+    surface_pressure = config.getfloat('vertical_grid', 'surface_pressure')
+
+    p_interface, _, _ = pressure_and_spec_vol_from_state_at_geom_height(
+        config,
+        ds.layerThickness,
+        ds.temperature,
+        ds.salinity,
+        surface_pressure * xr.ones_like(ds.ssh),
+        iter_count=1,
+    )
+
+    pseudothickness = pseudothickness_from_pressure(p_interface)
+
+    return pseudothickness
+
 
 def depth_from_thickness(ds):
     """

polaris/ocean/vertical/ztilde.py

@@ -2,9 +2,9 @@
 Conversions between Omega pseudo-height and pressure.
 
 Omega's vertical coordinate is pseudo-height
-    z_tilde = -p / (rho0 * g)
+    z_tilde = -p / (RhoSw * g)
 with z_tilde positive upward. Here, ``p`` is sea pressure in Pascals (Pa),
-``rho0`` is a reference density (kg m^-3), and ``g`` is gravitational
+``RhoSw`` is a reference density (kg m^-3), and ``g`` is gravitational
 acceleration as defined by ``polaris.constants`` via the Physical Constants
 Dictionary.
 """
@@ -20,6 +20,7 @@
 
 __all__ = [
     'z_tilde_from_pressure',
+    'pseudothickness_from_pressure',
     'pressure_from_z_tilde',
     'pressure_and_spec_vol_from_state_at_geom_height',
     'pressure_from_geom_thickness',
@@ -29,11 +30,48 @@
 RhoSw = get_constant('seawater_density_reference')
 
 
+def pseudothickness_from_pressure(
+    p: xr.DataArray,
+) -> xr.DataArray:
+    """
+    Convert sea pressure to pseudo-thickness.
+
+    z_tilde = -p / (RhoSw * g)
+
+    Parameters
+    ----------
+    p : xarray.DataArray
+        Sea pressure in Pascals (Pa) at layer interfaces.
+
+    Returns
+    -------
+    xarray.DataArray
+        Pseudo-thickness at layer mid-points with dimensions NCells by
+        NVertLayers (one less layer than ``p``) (units: m).
+    """
+
+    p_top = p.isel(nVertLevelsP1=slice(0, -1))
+    p_bot = p.isel(nVertLevelsP1=slice(1, None))
+    dims = list(p.dims)
+    dims = [item.replace('nVertLevelsP1', 'nVertLevels') for item in dims]
+    h = xr.DataArray(
+        (p_bot - p_top) / (RhoSw * Gravity),
+        dims=tuple(dims),
+    )
+    return h.assign_attrs(
+        {
+            'long_name': 'pseudo-thickness',
+            'units': 'm',
+            'note': 'h_tilde = -dp / (RhoSw * g)',
+        }
+    )
+
+
 def z_tilde_from_pressure(p: xr.DataArray) -> xr.DataArray:
     """
     Convert sea pressure to pseudo-height.
 
-    z_tilde = -p / (rho0 * g)
+    z_tilde = -p / (RhoSw * g)
 
     Parameters
     ----------
@@ -51,7 +89,7 @@ def z_tilde_from_pressure(p: xr.DataArray) -> xr.DataArray:
         {
             'long_name': 'pseudo-height',
             'units': 'm',
-            'note': 'z_tilde = -p / (rho0 * g)',
+            'note': 'z_tilde = -p / (RhoSw * g)',
         }
     )
 
@@ -60,7 +98,7 @@ def pressure_from_z_tilde(z_tilde: xr.DataArray) -> xr.DataArray:
     """
     Convert pseudo-height to sea pressure.
 
-    p = -z_tilde * (rho0 * g)
+    p = -z_tilde * (RhoSw * g)
 
     Parameters
     ----------
@@ -78,7 +116,7 @@ def pressure_from_z_tilde(z_tilde: xr.DataArray) -> xr.DataArray:
         {
             'long_name': 'sea pressure',
             'units': 'Pa',
-            'note': 'p = -rho0 * g * z_tilde',
+            'note': 'p = -RhoSw * g * z_tilde',
         }
     )
 

polaris/tasks/ocean/__init__.py

@@ -7,6 +7,10 @@
 from ruamel.yaml import YAML
 
 from polaris import Component
+from polaris.ocean.vertical.diagnostics import (
+    geom_thickness_from_ds,
+    pseudothickness_from_ds,
+)
 
 
 class Ocean(Component):
@@ -154,7 +158,7 @@ def map_var_list_to_native_model(self, var_list):
             ]
         return renamed_vars
 
-    def write_model_dataset(self, ds, filename):
+    def write_model_dataset(self, ds, filename, config=None):
         """
         Write out the given dataset, mapping dimension and variable names from
         MPAS-Ocean to Omega names if appropriate
@@ -167,6 +171,14 @@ def write_model_dataset(self, ds, filename):
         filename : str
             The path for the NetCDF file to write
         """
+        if (
+            self.model == 'omega'
+            and 'layerThickness' in ds.keys()
+            and 'PseudoThickness' not in ds.keys()
+            and config is not None
+        ):
+            ds['PseudoThickness'] = pseudothickness_from_ds(ds, config=config)
+            ds['layerThickness'] = ds['PseudoThickness'].copy()
         ds = self.map_to_native_model_vars(ds)
         write_netcdf(ds=ds, fileName=filename)
 
@@ -233,7 +245,7 @@ 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, config=None, **kwargs):
         """
         Open the given dataset, mapping variable and dimension names from Omega
         to MPAS-Ocean names if appropriate
@@ -252,6 +264,13 @@ def open_model_dataset(self, filename, **kwargs):
             The dataset with variables named as expected in MPAS-Ocean
         """
         ds = xr.open_dataset(filename, **kwargs)
+        if (
+            self.model == 'omega'
+            and 'LayerThickness' in ds.keys()
+            and config is not None
+        ):
+            ds['PseudoThickness'] = ds.LayerThickness
+            ds['LayerThickness'] = geom_thickness_from_ds(ds, config=config)
         ds = self.map_from_native_model_vars(ds)
         return ds
 

polaris/tasks/ocean/barotropic_channel/init.py

@@ -4,12 +4,12 @@
 from mpas_tools.mesh.conversion import convert, cull
 from mpas_tools.planar_hex import make_planar_hex_mesh
 
-from polaris import Step
 from polaris.mesh.planar import compute_planar_hex_nx_ny
+from polaris.ocean.model import OceanIOStep
 from polaris.ocean.vertical import init_vertical_coord
 
 
-class Init(Step):
+class Init(OceanIOStep):
     """
     A step for creating a mesh and initial condition for barotropic channel
     tasks
@@ -102,7 +102,7 @@ def run(self):
         ds['fCell'] = xr.zeros_like(ds.xCell)
         ds['fEdge'] = xr.zeros_like(ds.xEdge)
         ds['fVertex'] = xr.zeros_like(ds.xVertex)
-        write_netcdf(ds, 'initial_state.nc')
+        self.write_model_dataset(ds, 'initial_state.nc', config)
 
         # set the wind stress forcing
         ds_forcing = xr.Dataset()
@@ -114,4 +114,4 @@ def run(self):
         ds_forcing['windStressMeridional'] = (
             wind_stress_meridional.expand_dims(dim='Time', axis=0)
         )
-        write_netcdf(ds_forcing, 'forcing.nc')
+        self.write_model_dataset(ds_forcing, 'forcing.nc')

polaris/tasks/ocean/barotropic_channel/viz.py

@@ -40,8 +40,8 @@ def run(self):
         Run this step of the task
         """
         ds_mesh = self.open_model_dataset('mesh.nc')
-        ds_init = self.open_model_dataset('init.nc')
-        ds_out = self.open_model_dataset('output.nc')
+        ds_init = self.open_model_dataset('init.nc', self.config)
+        ds_out = self.open_model_dataset('output.nc', self.config)
 
         cell_mask = ds_init.maxLevelCell >= 1
         vertex_mask = ds_init.boundaryVertex == 0

polaris/tasks/ocean/barotropic_gyre/init.py

@@ -109,6 +109,8 @@ def run(self):
 
         # use polaris framework functions to initialize the vertical coordinate
         init_vertical_coord(config, ds)
+        ds['temperature'] = 20.0 * xr.ones_like(ds.zMid)
+        ds['salinity'] = 35.0 * xr.ones_like(ds.zMid)
 
         # set the coriolis values
         for loc in ['Cell', 'Edge', 'Vertex']:
@@ -140,7 +142,7 @@ def run(self):
         )
 
         # write the initial condition file
-        self.write_model_dataset(ds, 'init.nc')
+        self.write_model_dataset(ds, 'init.nc', config)
 
         cell_mask = ds.maxLevelCell >= 1
 

polaris/tasks/ocean/cosine_bell/analysis.py

@@ -83,7 +83,9 @@ def exact_solution(self, refinement_factor, field_name, time, zidx=None):
         ds_mesh = xr.open_dataset(f'mesh_r{refinement_factor:02g}.nc')
         sphere_radius = ds_mesh.sphere_radius
 
-        ds_init = self.open_model_dataset(f'init_r{refinement_factor:02g}.nc')
+        ds_init = self.open_model_dataset(
+            f'init_r{refinement_factor:02g}.nc', config
+        )
         latCell = ds_init.latCell.values
         lonCell = ds_init.lonCell.values
 

polaris/tasks/ocean/cosine_bell/init.py

@@ -112,7 +112,7 @@ def run(self):
         ds['fEdge'] = xr.zeros_like(ds_mesh.xEdge)
         ds['fVertex'] = xr.zeros_like(ds_mesh.xVertex)
 
-        self.write_model_dataset(ds, 'initial_state.nc')
+        self.write_model_dataset(ds, 'initial_state.nc', config)
 
 
 def cosine_bell(max_value, ri, r):

polaris/tasks/ocean/cosine_bell/validate.py

@@ -58,8 +58,8 @@ def run(self):
                 all_pass = False
 
         if all_pass:
-            ds1 = self.open_model_dataset(filename1)
-            ds2 = self.open_model_dataset(filename2)
+            ds1 = self.open_model_dataset(filename1, self.config)
+            ds2 = self.open_model_dataset(filename2, self.config)
 
             all_pass = compare_variables(
                 component=self.component,

polaris/tasks/ocean/cosine_bell/viz.py

@@ -67,7 +67,7 @@ def run(self):
         mesh_name = self.mesh_name
         run_duration = config.getfloat('convergence_forward', 'run_duration')
 
-        ds_init = self.open_model_dataset('initial_state.nc')
+        ds_init = self.open_model_dataset('initial_state.nc', config)
         da = ds_init['tracer1'].isel(Time=0, nVertLevels=0)
 
         plot_global_mpas_field(

polaris/tasks/ocean/manufactured_solution/analysis.py

@@ -67,7 +67,9 @@ def exact_solution(self, refinement_factor, field_name, time, zidx=None):
         solution : xarray.DataArray
             The exact solution as derived from the initial condition
         """
-        init = self.open_model_dataset(f'init_r{refinement_factor:02g}.nc')
+        init = self.open_model_dataset(
+            f'init_r{refinement_factor:02g}.nc', self.config
+        )
         exact = ExactSolution(self.config, init)
         if field_name != 'ssh':
             raise ValueError(f'{field_name} is not currently supported')

polaris/tasks/ocean/manufactured_solution/init.py

@@ -109,5 +109,7 @@ def run(self):
             'Time', 'nCells', 'nVertLevels'
         )
         ds['layerThickness'] = layer_thickness
+        ds['temperature'] = xr.zeros_like(layer_thickness)
+        ds['salinity'] = xr.ones_like(layer_thickness)
 
-        self.write_model_dataset(ds, 'initial_state.nc')
+        self.write_model_dataset(ds, 'initial_state.nc', config)