"""
General utilities for Omega's z-tilde (pseudo-height) vertical coordinate.
Omega's vertical coordinate is pseudo-height
z_tilde = -p / (RhoSw * g)
with z_tilde positive upward. Here, ``p`` is sea gauge pressure (pressure
relative to the atmosphere, zero at the free surface) in Pascals (Pa),
``RhoSw`` is a reference density (kg m^-3), and ``g`` is gravitational
acceleration as defined by ``polaris.constants`` via the Physical Constants
Dictionary.
These utilities are shared across any z-tilde–based coordinate. For the
p-star coordinate specifically (which uses ``RefPseudoThickness`` and
``VertCoordMovementWeights``), see
:py:mod:`polaris.ocean.vertical.pstar`.
"""
import logging
import numpy as np
import xarray as xr
from polaris.config import PolarisConfigParser
from polaris.constants import get_constant
from polaris.ocean.eos import compute_specvol
__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',
]
Gravity = get_constant('standard_acceleration_of_gravity')
RhoSw = get_constant('seawater_density_reference')
def pseudothickness_from_pressure(
p: xr.DataArray,
) -> xr.DataArray:
"""
Convert sea gauge pressure to pseudo-thickness.
z_tilde = -p / (RhoSw * g)
Parameters
----------
p : xarray.DataArray
Sea gauge 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)',
}
)
[docs]
def z_tilde_from_pressure(p: xr.DataArray) -> xr.DataArray:
"""
Convert sea gauge pressure to pseudo-height.
z_tilde = -p / (RhoSw * g)
Parameters
----------
p : xarray.DataArray
Sea gauge pressure in Pascals (Pa).
Returns
-------
xarray.DataArray
Pseudo-height with the same shape and coords as ``p`` (units: m).
"""
z = -(p) / (RhoSw * Gravity)
return z.assign_attrs(
{
'long_name': 'pseudo-height',
'units': 'm',
'note': 'z_tilde = -p / (RhoSw * g)',
}
)
[docs]
def pressure_from_z_tilde(z_tilde: xr.DataArray) -> xr.DataArray:
"""
Convert pseudo-height to sea gauge pressure.
p = -z_tilde * (RhoSw * g)
Pressure here is gauge pressure (relative to the atmosphere), which is
zero at the free surface (z_tilde = 0) and positive downward.
Parameters
----------
z_tilde : xarray.DataArray
Pseudo-height in meters (m), positive upward.
Returns
-------
xarray.DataArray
Sea gauge pressure with the same shape and coords as ``z_tilde``
(Pa).
"""
p = -(z_tilde) * (RhoSw * Gravity)
return p.assign_attrs(
{
'long_name': 'sea gauge pressure',
'units': 'Pa',
'note': 'p = -RhoSw * g * z_tilde',
}
)
[docs]
def pressure_from_geom_thickness(
surf_pressure: xr.DataArray,
geom_layer_thickness: xr.DataArray,
spec_vol: xr.DataArray,
) -> tuple[xr.DataArray, xr.DataArray]:
"""
Compute gauge pressure at layer interfaces and midpoints given surface
gauge pressure, geometric layer thicknesses, and specific volume. This
calculation assumes a constant specific volume within each layer.
Parameters
----------
surf_pressure : xarray.DataArray
The surface gauge pressure at the top of the water column (zero for
a free surface open to the atmosphere).
geom_layer_thickness : xarray.DataArray
The geometric thickness of each layer, set to zero for invalid layers.
spec_vol : xarray.DataArray
The specific volume at each layer.
Returns
-------
p_interface : xarray.DataArray
The gauge pressure at layer interfaces.
p_mid : xarray.DataArray
The gauge pressure at layer midpoints.
"""
dp = Gravity / spec_vol * geom_layer_thickness
p_interface = dp.cumsum(dim='nVertLevels').pad(
nVertLevels=(1, 0), mode='constant', constant_values=0.0
)
p_interface = surf_pressure + p_interface
p_interface = p_interface.rename({'nVertLevels': 'nVertLevelsP1'})
p_interface_top = p_interface.isel(nVertLevelsP1=slice(0, -1)).rename(
{'nVertLevelsP1': 'nVertLevels'}
)
p_mid = p_interface_top + 0.5 * dp
dims = list(geom_layer_thickness.dims)
interface_dims = [dim for dim in dims if dim != 'nVertLevels']
interface_dims.append('nVertLevelsP1')
p_interface = p_interface.transpose(*interface_dims)
p_mid = p_mid.transpose(*dims)
return p_interface, p_mid
[docs]
def pressure_and_spec_vol_from_state_at_geom_height(
config: PolarisConfigParser,
geom_layer_thickness: xr.DataArray,
temperature: xr.DataArray,
salinity: xr.DataArray,
surf_pressure: xr.DataArray,
iter_count: int,
logger: logging.Logger | None = None,
) -> tuple[xr.DataArray, xr.DataArray, xr.DataArray]:
"""
Compute gauge pressure at layer interfaces and midpoints, as well as the
specific volume at midpoints given geometric layer thicknesses,
temperature and salinity at layer midpoints (i.e. constant in geometric
height, not pseudo-height), and surface gauge pressure. The solution is
found iteratively starting from a specific volume calculated from the
reference density.
Requires config options needed by
{py:func}`polaris.ocean.eos.compute_specvol()`.
Parameters
----------
config : polaris.config.PolarisConfigParser
Configuration options with parameters defining the equation of state.
geom_layer_thickness : xarray.DataArray
The geometric thickness of each layer.
temperature : xarray.DataArray
The temperature at layer midpoints.
salinity : xarray.DataArray
The salinity at layer midpoints.
surf_pressure : xarray.DataArray
The surface gauge pressure at the top of the water column (zero for
a free surface open to the atmosphere).
iter_count : int
The number of iterations to perform.
logger : logging.Logger, optional
A logger for logging iteration information.
Returns
-------
p_interface : xarray.DataArray
The gauge pressure at layer interfaces.
p_mid : xarray.DataArray
The gauge pressure at layer midpoints.
spec_vol : xarray.DataArray
The specific volume at layer midpoints.
"""
spec_vol = 1.0 / RhoSw * xr.ones_like(geom_layer_thickness)
p_interface, p_mid = pressure_from_geom_thickness(
surf_pressure=surf_pressure,
geom_layer_thickness=geom_layer_thickness,
spec_vol=spec_vol,
)
prev_spec_vol = spec_vol
for iter in range(iter_count):
spec_vol = compute_specvol(
config=config,
temperature=temperature,
salinity=salinity,
pressure=p_mid,
)
if logger is not None:
delta_spec_vol = spec_vol - prev_spec_vol
max_delta = np.abs(delta_spec_vol).max().item()
prev_spec_vol = spec_vol
logger.info(
f'Max change in specific volume during EOS iteration {iter}: '
f'{max_delta:.3e} m3 kg-1'
)
p_interface, p_mid = pressure_from_geom_thickness(
surf_pressure=surf_pressure,
geom_layer_thickness=geom_layer_thickness,
spec_vol=spec_vol,
)
return p_interface, p_mid, spec_vol
[docs]
def geom_height_from_pseudo_height(
geom_z_bot: xr.DataArray,
h_tilde: xr.DataArray,
spec_vol: xr.DataArray,
min_level_cell: np.ndarray,
max_level_cell: np.ndarray,
) -> tuple[xr.DataArray, xr.DataArray]:
"""
Sum geometric heights from pseudo-heights and specific volume.
Parameters
----------
geom_z_bot : xarray.DataArray
Geometric height at the bathymetry for each water column.
h_tilde : xarray.DataArray
Pseudo-thickness of vertical layers, set to zero for invalid layers.
spec_vol : xarray.DataArray
Specific volume at midpoints of vertical layers.
min_level_cell : xarray.DataArray
Minimum valid zero-based level index for each cell.
max_level_cell : xarray.DataArray
Maximum valid zero-based level index for each cell.
Returns
-------
geom_z_inter : xarray.DataArray
Geometric height at layer interfaces.
geom_z_mid : xarray.DataArray
Geometric height at layer midpoints.
"""
# geometric height starts with geom_z_bot at the bottom
# and adds up the contributions from each layer above
n_vert_levels = spec_vol.sizes['nVertLevels']
z_index = xr.DataArray(np.arange(n_vert_levels), dims=['nVertLevels'])
mid_mask = np.logical_and(
z_index >= min_level_cell, z_index <= max_level_cell
)
geom_thickness = (spec_vol * h_tilde * RhoSw).where(mid_mask, 0.0)
dz_rev = geom_thickness.isel(nVertLevels=slice(None, None, -1))
sum_from_level = dz_rev.cumsum(dim='nVertLevels').isel(
nVertLevels=slice(None, None, -1)
)
geom_z_inter_top = geom_z_bot + sum_from_level
geom_z_inter = geom_z_inter_top.pad(
nVertLevels=(0, 1), mode='constant', constant_values=0.0
)
geom_z_inter[dict(nVertLevels=n_vert_levels)] = geom_z_bot
geom_z_inter = geom_z_inter.rename({'nVertLevels': 'nVertLevelsP1'})
z_index_p1 = xr.DataArray(
np.arange(n_vert_levels + 1), dims=['nVertLevelsP1']
)
inter_mask = np.logical_and(
z_index_p1 >= min_level_cell,
z_index_p1 <= max_level_cell + 1,
)
geom_z_inter = geom_z_inter.where(inter_mask)
geom_z_inter_lower = geom_z_inter.isel(
nVertLevelsP1=slice(1, None)
).rename({'nVertLevelsP1': 'nVertLevels'})
geom_z_mid = (geom_z_inter_lower + 0.5 * geom_thickness).where(mid_mask)
# transpose to match h_tilde. For geom_z_inter, replace nVertLevels with
# nVertLevelsP1 at the same index in the list of dimensions
z_mid_dims = list(h_tilde.dims)
z_inter_dims = z_mid_dims.copy()
z_inter_dims[z_mid_dims.index('nVertLevels')] = 'nVertLevelsP1'
geom_z_inter = geom_z_inter.transpose(*z_inter_dims)
geom_z_mid = geom_z_mid.transpose(*z_mid_dims)
return geom_z_inter, geom_z_mid
def get_iter_count_for_eos(config: PolarisConfigParser) -> int:
"""
Get the number of iterations to perform when adjusting the
pseudo-bottom-depth to hit the right geometric bottom depth. Default is
from the `pseudothickness_iter_count` config option for TEOS-10 and 1 for
other equations of state.
Parameters
----------
config : polaris.config.PolarisConfigParser
Configuration options with parameters defining the equation of state.
Returns
-------
int
The number of iterations to perform.
"""
eos_type = config.get('ocean', 'eos_type')
if eos_type == 'teos-10':
return config.getint('vertical_grid', 'pseudothickness_iter_count')
else:
return 1