Add TEOS-10 and z-tilde vertical coordinate

docs/developers_guide/ocean/api.md

@@ -555,6 +555,21 @@
    get_time_interval_string
 ```
 
+### Equations of state (EOS)
+
+```{eval-rst}
+.. currentmodule:: polaris.ocean.eos
+
+.. autosummary::
+   :toctree: generated/
+
+   compute_density
+   compute_specvol
+
+   linear.compute_linear_density
+   teos10.compute_specvol
+```
+
 
 ### Reference Potential Energy (RPE)
 
@@ -577,13 +592,15 @@
    :toctree: generated/
 
    vertical.init_vertical_coord
+   vertical.compute_zint_zmid_from_layer_thickness
    vertical.diagnostics.depth_from_thickness
    vertical.grid_1d.generate_1d_grid
    vertical.grid_1d.write_1d_grid
    vertical.partial_cells.alter_bottom_depth
    vertical.partial_cells.alter_ssh
    vertical.sigma.init_sigma_vertical_coord
    vertical.sigma.update_sigma_layer_thickness
+   vertical.sigma.compute_sigma_layer_thickness
    vertical.update_layer_thickness
    vertical.zlevel.init_z_level_vertical_coord
    vertical.zlevel.update_z_level_layer_thickness
@@ -592,5 +609,10 @@
    vertical.zlevel.compute_z_level_resting_thickness
    vertical.zstar.init_z_star_vertical_coord
    vertical.zstar.update_z_star_layer_thickness
+   vertical.ztilde.z_tilde_from_pressure
+   vertical.ztilde.pressure_from_z_tilde
+   vertical.ztilde.pressure_from_geom_thickness
+   vertical.ztilde.pressure_and_spec_vol_from_state_at_geom_height
+   vertical.ztilde.geom_height_from_pseudo_height
 ```
 

docs/developers_guide/ocean/framework.md

@@ -195,6 +195,28 @@ output to double precision, adjusting sea surface height in ice-shelf cavities,
 and outputting variables related to frazil ice and land-ice fluxes.
 
 
+(dev-ocean-framework-eos)=
+
+## Equations of state (EOS)
+
+Polaris ocean tasks use EOS utilities from the `polaris.ocean.eos` package.
+
+The high-level APIs are {py:func}`polaris.ocean.eos.compute_density()` and
+{py:func}`polaris.ocean.eos.compute_specvol()`.  These functions dispatch
+based on `eos_type` in the `[ocean]` config section.
+
+- `eos_type = linear` uses the linear EOS and calls
+  {py:func}`polaris.ocean.eos.linear.compute_linear_density()`.
+- `eos_type = teos-10` uses TEOS-10 and calls
+  {py:func}`polaris.ocean.eos.teos10.compute_specvol()`, with density computed
+  as the inverse of specific volume.
+
+The linear EOS coefficients and reference values are set with ocean config
+options such as `eos_linear_alpha`, `eos_linear_beta`, `eos_linear_rhoref`,
+`eos_linear_Tref`, and `eos_linear_Sref`.  TEOS-10 requires pressure to be
+provided when calling the high-level EOS functions.
+
+
 (dev-ocean-spherical-meshes)=
 
 ## Quasi-uniform and Icosahedral Spherical Meshes
@@ -534,11 +556,32 @@ this section of the config file. The function
 {py:func}`polaris.ocean.vertical.init_vertical_coord()` can be used to compute
 `minLevelCell`, `maxLevelCell`, `cellMask`, `layerThickness`, `zMid`,
 and `restingThickness` variables for {ref}`ocean-z-level` and
-{ref}`ocean-z-star` coordinates using the `ssh` and `bottomDepth` as well
+{ref}`ocean-z-star` coordinates (including the `z-tilde` option handled
+through z-star infrastructure) using the `ssh` and `bottomDepth` as well
 as config options from `vertical_grid`. The function
-{py:func}`polaris.ocean.vertical.update_layer_thickness` can be used to update
+{py:func}`polaris.ocean.vertical.update_layer_thickness()` can be used to update
 `layerThickness` when either or both of `bottomDepth` and `ssh` have been
-changed.
+changed.  After thicknesses are updated, tasks can call
+{py:func}`polaris.ocean.vertical.compute_zint_zmid_from_layer_thickness()` to
+recover `zInterface` and `zMid` from the resulting layer thickness.
+
+For workflows that need pseudo-height/pressure conversion, the
+`polaris.ocean.vertical.ztilde` module provides utilities:
+
+- {py:func}`polaris.ocean.vertical.ztilde.z_tilde_from_pressure()` and
+  {py:func}`polaris.ocean.vertical.ztilde.pressure_from_z_tilde()` convert
+  between pseudo-height and pressure.
+- {py:func}`polaris.ocean.vertical.ztilde.pressure_from_geom_thickness()` and
+  {py:func}`polaris.ocean.vertical.ztilde.pressure_and_spec_vol_from_state_at_geom_height()`
+  compute hydrostatic pressure (and specific volume) from geometric layer
+  thickness and state variables.
+- {py:func}`polaris.ocean.vertical.ztilde.geom_height_from_pseudo_height()`
+  reconstructs geometric layer-interface and midpoint heights from
+  pseudo-thickness and specific volume.
+
+For sigma coordinates, shared functionality for direct thickness computation is
+available in
+{py:func}`polaris.ocean.vertical.sigma.compute_sigma_layer_thickness()`.
 
 (dev-ocean-rpe)=
 

docs/tutorials/dev_add_category_of_tasks/fleshing_out_step.md

@@ -168,7 +168,7 @@ $ vim polaris/tasks/ocean/my_overflow/init.py
 
 ...
 
-from polaris.ocean.model.eos import compute_density
+from polaris.ocean.eos import compute_density
 from polaris.ocean.vertical import init_vertical_coord
 
 

polaris/ocean/eos/__init__.py

@@ -0,0 +1,102 @@
+import xarray as xr
+
+from polaris.config import PolarisConfigParser
+
+from .linear import compute_linear_density
+from .teos10 import compute_specvol as compute_teos10_specvol
+
+
+def compute_density(
+    config: PolarisConfigParser,
+    temperature: xr.DataArray,
+    salinity: xr.DataArray,
+    pressure: xr.DataArray | None = None,
+) -> xr.DataArray:
+    """
+    Compute the density of seawater based on the equation of state specified
+    in the configuration.
+
+    Parameters
+    ----------
+    config : polaris.config.PolarisConfigParser
+        Configuration object containing ocean parameters.
+
+    temperature : float or xarray.DataArray
+        Temperature (conservative, potential or in-situ) of the seawater.
+
+    salinity : float or xarray.DataArray
+        Salinity (practical or absolute) of the seawater.
+
+    pressure : float or xarray.DataArray, optional
+        Pressure (in-situ or reference) of the seawater.
+
+    Returns
+    -------
+    density : float or xarray.DataArray
+        Computed density (in-situ or reference) of the seawater.
+    """
+    eos_type = config.get('ocean', 'eos_type')
+    if eos_type == 'linear':
+        density = compute_linear_density(config, temperature, salinity)
+    elif eos_type == 'teos-10':
+        if pressure is None:
+            raise ValueError(
+                'Pressure must be provided when using the TEOS-10 equation of '
+                'state.'
+            )
+        density = 1.0 / compute_teos10_specvol(
+            sa=salinity, ct=temperature, p=pressure
+        )
+    else:
+        raise ValueError(f'Unsupported equation of state type: {eos_type}')
+    density.attrs['units'] = 'kg m-3'
+    density.attrs['long_name'] = 'density'
+    return density
+
+
+def compute_specvol(
+    config: PolarisConfigParser,
+    temperature: xr.DataArray,
+    salinity: xr.DataArray,
+    pressure: xr.DataArray | None = None,
+) -> xr.DataArray:
+    """
+    Compute the specific volume of seawater based on the equation of state
+    specified in the configuration.
+
+    Parameters
+    ----------
+    config : polaris.config.PolarisConfigParser
+        Configuration object containing ocean parameters.
+
+    temperature : float or xarray.DataArray
+        Temperature (conservative, potential or in-situ) of the seawater.
+
+    salinity : float or xarray.DataArray
+        Salinity (practical or absolute) of the seawater.
+
+    pressure : float or xarray.DataArray, optional
+        Pressure (in-situ or reference) of the seawater.
+
+    Returns
+    -------
+    specvol : float or xarray.DataArray
+        Computed specific volume (in-situ or reference) of the seawater.
+    """
+    eos_type = config.get('ocean', 'eos_type')
+    if eos_type == 'linear':
+        specvol = 1.0 / compute_linear_density(config, temperature, salinity)
+    elif eos_type == 'teos-10':
+        if pressure is None:
+            raise ValueError(
+                'Pressure must be provided when using the TEOS-10 equation of '
+                'state.'
+            )
+        specvol = compute_teos10_specvol(
+            sa=salinity, ct=temperature, p=pressure
+        )
+    else:
+        raise ValueError(f'Unsupported equation of state type: {eos_type}')
+    specvol.attrs['units'] = 'm3 kg-1'
+    specvol.attrs['long_name'] = 'specific volume'
+    return specvol

polaris/ocean/eos/linear.py

@@ -0,0 +1,47 @@
+import xarray as xr
+
+from polaris.config import PolarisConfigParser
+
+
+def compute_linear_density(
+    config: PolarisConfigParser,
+    temperature: xr.DataArray,
+    salinity: xr.DataArray,
+) -> xr.DataArray:
+    """
+    Compute the density of seawater based on the the linear equation of state
+    with coefficients specified in the configuration. The distinction between
+    conservative, potential, and in-situ temperature or between absolute and
+    practical salinity is not relevant for the linear EOS.
+
+    Parameters
+    ----------
+    config : polaris.config.PolarisConfigParser
+        Configuration object containing ocean parameters.
+
+    temperature : float or xarray.DataArray
+        Temperature of the seawater.
+
+    salinity : float or xarray.DataArray
+        Salinity of the seawater.
+
+    Returns
+    -------
+    density : float or xarray.DataArray
+        Computed density of the seawater.
+    """
+    section = config['ocean']
+    alpha = section.getfloat('eos_linear_alpha')
+    beta = section.getfloat('eos_linear_beta')
+    rhoref = section.getfloat('eos_linear_rhoref')
+    Tref = section.getfloat('eos_linear_Tref')
+    Sref = section.getfloat('eos_linear_Sref')
+    assert (
+        alpha is not None
+        and beta is not None
+        and rhoref is not None
+        and Tref is not None
+        and Sref is not None
+    ), 'All linear EOS parameters must be specified in the config options.'
+    density = rhoref + -alpha * (temperature - Tref) + beta * (salinity - Sref)
+    return density

polaris/ocean/eos/teos10.py

@@ -0,0 +1,57 @@
+import gsw
+import xarray as xr
+
+
+def compute_specvol(
+    sa: xr.DataArray, ct: xr.DataArray, p: xr.DataArray
+) -> xr.DataArray:
+    """
+    Compute specific volume from co-located p, CT and SA.
+
+    Notes
+    -----
+    - This function converts inputs to NumPy arrays and calls
+        ``gsw.specvol`` directly for performance. Inputs must fit in
+        memory.
+
+    - Any parallelization should be handled by the caller (e.g., splitting
+        over outer dimensions and calling this function per chunk).
+
+    Parameters
+    ----------
+    sa : xarray.DataArray
+        Absolute Salinity at the same points as p and ct.
+
+    ct : xarray.DataArray
+        Conservative Temperature at the same points as p and sa.
+
+    p : xarray.DataArray
+        Sea pressure in Pascals (Pa) at the same points as ct and sa.
+
+    Returns
+    -------
+    xarray.DataArray
+        Specific volume with the same dims/coords as ct and sa (m^3/kg).
+    """
+
+    # Check sizes/dims match exactly
+    if not (p.sizes == ct.sizes == sa.sizes):
+        raise ValueError(
+            'p, ct and sa must have identical dimensions and sizes; '
+            f'got p={p.sizes}, ct={ct.sizes}, sa={sa.sizes}'
+        )
+
+    # Ensure coordinates align identically (names and labels)
+    p, ct, sa = xr.align(p, ct, sa, join='exact')
+
+    # Convert to NumPy and call gsw directly for performance
+    p_dbar = (p / 1.0e4).to_numpy()
+    ct_np = ct.to_numpy()
+    sa_np = sa.to_numpy()
+    specvol_np = gsw.specvol(sa_np, ct_np, p_dbar)
+
+    specvol = xr.DataArray(
+        specvol_np, dims=ct.dims, coords=ct.coords, name='specvol'
+    )
+
+    return specvol