diff --git a/docs/developers_guide/ocean/api.md b/docs/developers_guide/ocean/api.md index fa02a083cd..97a0ccea26 100644 --- a/docs/developers_guide/ocean/api.md +++ b/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,6 +592,7 @@ :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 @@ -584,6 +600,7 @@ 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 ``` diff --git a/docs/developers_guide/ocean/framework.md b/docs/developers_guide/ocean/framework.md index 4fc90ac867..f07ef035be 100644 --- a/docs/developers_guide/ocean/framework.md +++ b/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)= diff --git a/docs/tutorials/dev_add_category_of_tasks/fleshing_out_step.md b/docs/tutorials/dev_add_category_of_tasks/fleshing_out_step.md index 7147099ca7..be6b5204c5 100644 --- a/docs/tutorials/dev_add_category_of_tasks/fleshing_out_step.md +++ b/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 diff --git a/polaris/ocean/eos/__init__.py b/polaris/ocean/eos/__init__.py new file mode 100644 index 0000000000..0afbfbd268 --- /dev/null +++ b/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 diff --git a/polaris/ocean/eos/linear.py b/polaris/ocean/eos/linear.py new file mode 100644 index 0000000000..2ac76bd1d6 --- /dev/null +++ b/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 diff --git a/polaris/ocean/eos/teos10.py b/polaris/ocean/eos/teos10.py new file mode 100644 index 0000000000..75e5ddace1 --- /dev/null +++ b/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 diff --git a/polaris/ocean/model/eos.py b/polaris/ocean/model/eos.py deleted file mode 100644 index 7099c1379a..0000000000 --- a/polaris/ocean/model/eos.py +++ /dev/null @@ -1,16 +0,0 @@ -def compute_density(config, temperature, salinity): - eos_type = config.get('ocean', 'eos_type') - if eos_type == 'linear': - density = _compute_linear_density(config, temperature, salinity) - return density - - -def _compute_linear_density(config, temperature, salinity): - 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') - density = rhoref + -alpha * (temperature - Tref) + beta * (salinity - Sref) - return density diff --git a/polaris/ocean/vertical/__init__.py b/polaris/ocean/vertical/__init__.py index 40f422872e..04cea553e0 100644 --- a/polaris/ocean/vertical/__init__.py +++ b/polaris/ocean/vertical/__init__.py @@ -85,7 +85,7 @@ def init_vertical_coord(config, ds): if coord_type == 'z-level': init_z_level_vertical_coord(config, ds) - elif coord_type == 'z-star': + elif coord_type == 'z-star' or coord_type == 'z-tilde': init_z_star_vertical_coord(config, ds) elif coord_type == 'sigma': init_sigma_vertical_coord(config, ds) @@ -110,8 +110,11 @@ def init_vertical_coord(config, ds): dim='Time', axis=0 ) - ds['zMid'] = _compute_zmid_from_layer_thickness( - ds.layerThickness, ds.ssh, ds.cellMask + ds['zInterface'], ds['zMid'] = compute_zint_zmid_from_layer_thickness( + layer_thickness=ds.layerThickness, + bottom_depth=ds.bottomDepth, + min_level_cell=ds.minLevelCell, + max_level_cell=ds.maxLevelCell, ) # fortran 1-based indexing @@ -148,7 +151,7 @@ def update_layer_thickness(config, ds): if coord_type == 'z-level': update_z_level_layer_thickness(config, ds) - elif coord_type == 'z-star': + elif coord_type == 'z-star' or coord_type == 'z-tilde': update_z_star_layer_thickness(config, ds) elif coord_type == 'sigma': update_sigma_layer_thickness(config, ds) @@ -162,47 +165,89 @@ def update_layer_thickness(config, ds): ds['layerThickness'] = ds.layerThickness.expand_dims(dim='Time', axis=0) -def _compute_cell_mask(minLevelCell, maxLevelCell, nVertLevels): - cellMask = [] - for zIndex in range(nVertLevels): - mask = np.logical_and(zIndex >= minLevelCell, zIndex <= maxLevelCell) - cellMask.append(mask) - cellMaskArray = xr.DataArray(cellMask, dims=['nVertLevels', 'nCells']) - cellMaskArray = cellMaskArray.transpose('nCells', 'nVertLevels') - return cellMaskArray - - -def _compute_zmid_from_layer_thickness(layerThickness, ssh, cellMask): +def compute_zint_zmid_from_layer_thickness( + layer_thickness: xr.DataArray, + bottom_depth: xr.DataArray, + min_level_cell: xr.DataArray, + max_level_cell: xr.DataArray, +) -> tuple[xr.DataArray, xr.DataArray]: """ - Compute zMid from ssh and layerThickness for any vertical coordinate + Compute height z at layer interfaces and midpoints given layer thicknesses + and bottom depth. Parameters ---------- - layerThickness : xarray.DataArray - The thickness of each layer + layer_thickness : xarray.DataArray + The layer thickness of each layer. - ssh : xarray.DataArray - The sea surface height + bottom_depth : xarray.DataArray + The positive-down depth of the seafloor. - cellMask : xarray.DataArray - A boolean mask of where there are valid cells + min_level_cell : xarray.DataArray + The zero-based minimum vertical index from each column. + + max_level_cell : xarray.DataArray + The zero-based maximum vertical index from each column. Returns ------- - zMid : xarray.DataArray - The elevation of layer centers + z_interface : xarray.DataArray + The elevation of layer interfaces. + + z_mid : xarray.DataArray + The elevation of layer midpoints. """ - zTop = ssh.copy() - nVertLevels = layerThickness.sizes['nVertLevels'] - zMid = [] - for zIndex in range(nVertLevels): - mask = cellMask.isel(nVertLevels=zIndex) - thickness = layerThickness.isel(nVertLevels=zIndex).where(mask, 0.0) - z = (zTop - 0.5 * thickness).where(mask) - zMid.append(z) - zTop -= thickness - zMid = xr.concat(zMid, dim='nVertLevels').transpose( - 'Time', 'nCells', 'nVertLevels' + n_vert_levels = layer_thickness.sizes['nVertLevels'] + + z_index = xr.DataArray(np.arange(n_vert_levels), dims=['nVertLevels']) + mask_mid = np.logical_and( + z_index >= min_level_cell, z_index <= max_level_cell ) - return zMid + + dz = layer_thickness.where(mask_mid, 0.0) + dz_rev = dz.isel(nVertLevels=slice(None, None, -1)) + sum_from_level = dz_rev.cumsum(dim='nVertLevels').isel( + nVertLevels=slice(None, None, -1) + ) + + z_bot = ( + xr.zeros_like(layer_thickness.isel(nVertLevels=0, drop=True)) + - bottom_depth + ) + z_interface_top = z_bot + sum_from_level + z_interface = z_interface_top.pad(nVertLevels=(0, 1), mode='constant') + z_interface[dict(nVertLevels=n_vert_levels)] = z_bot + z_interface = z_interface.rename({'nVertLevels': 'nVertLevelsP1'}) + + z_index_p1 = xr.DataArray( + np.arange(n_vert_levels + 1), dims=['nVertLevelsP1'] + ) + mask_interface = np.logical_and( + z_index_p1 >= min_level_cell, + z_index_p1 - 1 <= max_level_cell, + ) + z_interface = z_interface.where(mask_interface) + + z_interface_upper = z_interface.isel(nVertLevelsP1=slice(0, -1)).rename( + {'nVertLevelsP1': 'nVertLevels'} + ) + z_interface_lower = z_interface.isel(nVertLevelsP1=slice(1, None)).rename( + {'nVertLevelsP1': 'nVertLevels'} + ) + z_mid = (0.5 * (z_interface_upper + z_interface_lower)).where(mask_mid) + + dims = list(layer_thickness.dims) + interface_dims = [dim for dim in dims if dim != 'nVertLevels'] + interface_dims.append('nVertLevelsP1') + z_interface = z_interface.transpose(*interface_dims) + z_mid = z_mid.transpose(*dims) + + return z_interface, z_mid + + +def _compute_cell_mask(minLevelCell, maxLevelCell, nVertLevels): + z_index = xr.DataArray(np.arange(nVertLevels), dims=['nVertLevels']) + return np.logical_and( + z_index >= minLevelCell, z_index <= maxLevelCell + ).transpose('nCells', 'nVertLevels') diff --git a/polaris/ocean/vertical/zlevel.py b/polaris/ocean/vertical/zlevel.py index 6725b31cdb..eeb1a32f4f 100644 --- a/polaris/ocean/vertical/zlevel.py +++ b/polaris/ocean/vertical/zlevel.py @@ -118,10 +118,10 @@ def update_z_level_layer_thickness(config, ds): def compute_min_max_level_cell( - refTopDepth, - refBottomDepth, + ref_top_depth, + ref_bottom_depth, ssh, - bottomDepth, + bottom_depth, min_vert_levels=None, min_layer_thickness=None, ): @@ -131,16 +131,16 @@ def compute_min_max_level_cell( Parameters ---------- - refTopDepth : xarray.DataArray + ref_top_depth : xarray.DataArray A 1D array of positive-down depths of the top of each z level - refBottomDepth : xarray.DataArray + ref_bottom_depth : xarray.DataArray A 1D array of positive-down depths of the bottom of each z level ssh : xarray.DataArray The sea surface height - bottomDepth : xarray.DataArray + bottom_depth : xarray.DataArray The positive-down depth of the seafloor @@ -156,12 +156,14 @@ def compute_min_max_level_cell( A boolean mask of where there are valid cells """ if min_layer_thickness is not None: - valid = bottomDepth + min_layer_thickness * min_vert_levels >= -ssh + valid = bottom_depth + min_layer_thickness * min_vert_levels >= -ssh else: - valid = bottomDepth > -ssh + valid = bottom_depth > -ssh - aboveTopMask = (refBottomDepth <= -ssh).transpose('nCells', 'nVertLevels') - aboveBottomMask = (refTopDepth < bottomDepth).transpose( + aboveTopMask = (ref_bottom_depth <= -ssh).transpose( + 'nCells', 'nVertLevels' + ) + aboveBottomMask = (ref_top_depth < bottom_depth).transpose( 'nCells', 'nVertLevels' ) aboveBottomMask = numpy.logical_and(aboveBottomMask, valid) @@ -181,29 +183,34 @@ def compute_min_max_level_cell( def compute_z_level_layer_thickness( - refTopDepth, refBottomDepth, ssh, bottomDepth, minLevelCell, maxLevelCell + ref_top_depth, + ref_bottom_depth, + ssh, + bottom_depth, + min_level_cell, + max_level_cell, ): """ Compute z-level layer thickness from ssh and bottomDepth Parameters ---------- - refTopDepth : xarray.DataArray + ref_top_depth : xarray.DataArray A 1D array of positive-down depths of the top of each z level - refBottomDepth : xarray.DataArray + ref_bottom_depth : xarray.DataArray A 1D array of positive-down depths of the bottom of each z level ssh : xarray.DataArray The sea surface height - bottomDepth : xarray.DataArray + bottom_depth : xarray.DataArray The positive-down depth of the seafloor - minLevelCell : xarray.DataArray + min_level_cell : xarray.DataArray The zero-based index of the top valid level - maxLevelCell : xarray.DataArray + max_level_cell : xarray.DataArray The zero-based index of the bottom valid level Returns @@ -212,27 +219,31 @@ def compute_z_level_layer_thickness( The thickness of each layer (level) """ - nVertLevels = refBottomDepth.sizes['nVertLevels'] - layerThickness = [] - for zIndex in range(nVertLevels): - mask = numpy.logical_and( - zIndex >= minLevelCell, zIndex <= maxLevelCell - ) - zTop = numpy.minimum(ssh, -refTopDepth[zIndex]) - zBot = numpy.maximum(-bottomDepth, -refBottomDepth[zIndex]) - thickness = (zTop - zBot).where(mask, 0.0) - layerThickness.append(thickness) - layerThicknessArray = xarray.DataArray( - layerThickness, dims=['nVertLevels', 'nCells'] + n_vert_levels = ref_bottom_depth.sizes['nVertLevels'] + z_index = xarray.DataArray( + numpy.arange(n_vert_levels), dims=['nVertLevels'] ) - layerThicknessArray = layerThicknessArray.transpose( - 'nCells', 'nVertLevels' + mask = numpy.logical_and( + z_index >= min_level_cell, z_index <= max_level_cell + ).transpose('nCells', 'nVertLevels') + + z_top = xarray.where(ssh < -ref_top_depth, ssh, -ref_top_depth) + z_bot = xarray.where( + -bottom_depth > -ref_bottom_depth, + -bottom_depth, + -ref_bottom_depth, ) - return layerThicknessArray + thickness = (z_top - z_bot).transpose('nCells', 'nVertLevels') + + return thickness.where(mask, 0.0) def compute_z_level_resting_thickness( - layerThickness, ssh, bottomDepth, minLevelCell, maxLevelCell + layer_thickness, + ssh, + bottom_depth, + min_level_cell, + max_level_cell, ): """ Compute z-level resting thickness by "unstretching" layerThickness @@ -240,19 +251,19 @@ def compute_z_level_resting_thickness( Parameters ---------- - layerThickness : xarray.DataArray + layer_thickness : xarray.DataArray The thickness of each layer (level) ssh : xarray.DataArray The sea surface height - bottomDepth : xarray.DataArray + bottom_depth : xarray.DataArray The positive-down depth of the seafloor - minLevelCell : xarray.DataArray + min_level_cell : xarray.DataArray The zero-based index of the top valid level - maxLevelCell : xarray.DataArray + max_level_cell : xarray.DataArray The zero-based index of the bottom valid level Returns @@ -261,21 +272,15 @@ def compute_z_level_resting_thickness( The thickness of z-star layers when ssh = 0 """ - nVertLevels = layerThickness.sizes['nVertLevels'] - restingThickness = [] - - layerStretch = bottomDepth / (ssh + bottomDepth) - for zIndex in range(nVertLevels): - mask = numpy.logical_and( - zIndex >= minLevelCell, zIndex <= maxLevelCell - ) - thickness = layerStretch * layerThickness.isel(nVertLevels=zIndex) - thickness = thickness.where(mask, 0.0) - restingThickness.append(thickness) - restingThicknessArray = xarray.DataArray( - restingThickness, dims=['nVertLevels', 'nCells'] + n_vert_levels = layer_thickness.sizes['nVertLevels'] + z_index = xarray.DataArray( + numpy.arange(n_vert_levels), dims=['nVertLevels'] ) - restingThicknessArray = restingThicknessArray.transpose( - 'nCells', 'nVertLevels' - ) - return restingThicknessArray + mask = numpy.logical_and( + z_index >= min_level_cell, z_index <= max_level_cell + ).transpose('nCells', 'nVertLevels') + + layer_stretch = bottom_depth / (ssh + bottom_depth) + resting_thickness = layer_stretch * layer_thickness + + return resting_thickness.where(mask, 0.0) diff --git a/polaris/ocean/vertical/zstar.py b/polaris/ocean/vertical/zstar.py index f3c93a3267..23ff410d52 100644 --- a/polaris/ocean/vertical/zstar.py +++ b/polaris/ocean/vertical/zstar.py @@ -127,7 +127,11 @@ def update_z_star_layer_thickness(config, ds): def _compute_z_star_layer_thickness( - restingThickness, ssh, bottomDepth, minLevelCell, maxLevelCell + resting_thickness, + ssh, + bottom_depth, + min_level_cell, + max_level_cell, ): """ Compute z-star layer thickness by stretching restingThickness based on ssh @@ -135,19 +139,19 @@ def _compute_z_star_layer_thickness( Parameters ---------- - restingThickness : xarray.DataArray + resting_thickness : xarray.DataArray The thickness of z-star layers when ssh = 0 ssh : xarray.DataArray The sea surface height - bottomDepth : xarray.DataArray + bottom_depth : xarray.DataArray The positive-down depth of the seafloor - minLevelCell : xarray.DataArray + min_level_cell : xarray.DataArray The zero-based index of the top valid level - maxLevelCell : xarray.DataArray + max_level_cell : xarray.DataArray The zero-based index of the bottom valid level Returns @@ -156,23 +160,17 @@ def _compute_z_star_layer_thickness( The thickness of each layer (level) """ - nVertLevels = restingThickness.sizes['nVertLevels'] - layerThickness = [] + n_vert_levels = resting_thickness.sizes['nVertLevels'] + z_index = xarray.DataArray( + numpy.arange(n_vert_levels), dims=['nVertLevels'] + ) + mask = numpy.logical_and( + z_index >= min_level_cell, z_index <= max_level_cell + ).transpose('nCells', 'nVertLevels') - layerStretch = (ssh + bottomDepth) / restingThickness.sum( + layer_stretch = (ssh + bottom_depth) / resting_thickness.sum( dim='nVertLevels' ) - for zIndex in range(nVertLevels): - mask = numpy.logical_and( - zIndex >= minLevelCell, zIndex <= maxLevelCell - ) - thickness = layerStretch * restingThickness.isel(nVertLevels=zIndex) - thickness = thickness.where(mask, 0.0) - layerThickness.append(thickness) - layerThicknessArray = xarray.DataArray( - layerThickness, dims=['nVertLevels', 'nCells'] - ) - layerThicknessArray = layerThicknessArray.transpose( - 'nCells', 'nVertLevels' - ) - return layerThicknessArray + layer_thickness = layer_stretch * resting_thickness + + return layer_thickness.where(mask, 0.0) diff --git a/polaris/ocean/vertical/ztilde.py b/polaris/ocean/vertical/ztilde.py new file mode 100644 index 0000000000..550ce0632b --- /dev/null +++ b/polaris/ocean/vertical/ztilde.py @@ -0,0 +1,328 @@ +""" +Conversions between Omega pseudo-height and pressure. + +Omega's vertical coordinate is pseudo-height + z_tilde = -p / (rho0 * g) +with z_tilde positive upward. Here, ``p`` is sea pressure in Pascals (Pa), +``rho0`` is a reference density (kg m^-3) supplied by the caller, and +``g`` is gravitational acceleration obtained from +``mpas_tools.cime.constants``. +""" + +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', + 'pressure_from_z_tilde', + 'pressure_and_spec_vol_from_state_at_geom_height', + 'pressure_from_geom_thickness', +] + +Gravity = get_constant('standard_acceleration_of_gravity') + + +def z_tilde_from_pressure(p: xr.DataArray, rho0: float) -> xr.DataArray: + """ + Convert sea pressure to pseudo-height. + + z_tilde = -p / (rho0 * g) + + Parameters + ---------- + p : xarray.DataArray + Sea pressure in Pascals (Pa). + + rho0 : float + Reference density in kg m^-3. + + Returns + ------- + xarray.DataArray + Pseudo-height with the same shape and coords as ``p`` (units: m). + """ + + z = -(p) / (rho0 * Gravity) + return z.assign_attrs( + { + 'long_name': 'pseudo-height', + 'units': 'm', + 'note': 'z_tilde = -p / (rho0 * g)', + } + ) + + +def pressure_from_z_tilde(z_tilde: xr.DataArray, rho0: float) -> xr.DataArray: + """ + Convert pseudo-height to sea pressure. + + p = -z_tilde * (rho0 * g) + + Parameters + ---------- + z_tilde : xarray.DataArray + Pseudo-height in meters (m), positive upward. + + rho0 : float + Reference density in kg m^-3. + + Returns + ------- + xarray.DataArray + Sea pressure with the same shape and coords as ``z_tilde`` (Pa). + """ + + p = -(z_tilde) * (rho0 * Gravity) + return p.assign_attrs( + { + 'long_name': 'sea pressure', + 'units': 'Pa', + 'note': 'p = -rho0 * g * z_tilde', + } + ) + + +def pressure_from_geom_thickness( + surf_pressure: xr.DataArray, + geom_layer_thickness: xr.DataArray, + spec_vol: xr.DataArray, +) -> tuple[xr.DataArray, xr.DataArray]: + """ + Compute the pressure at layer interfaces and midpoints given surface + pressure, geometric layer thicknesses, and specific volume. This + calculation assumes a constant specific volume within each layer. + + Parameters + ---------- + surf_pressure : xarray.DataArray + The surface pressure at the top of the water column. + + 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 pressure at layer interfaces. + + p_mid : xarray.DataArray + The 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 + + +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 the 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 pressure. The solution is found + iteratively starting from a specific volume calculated from the reference + density. + + Requires config option ``[vertical_grid] rho0`` and those required + for {py:func}`polaris.ocean.eos.compute_specvol()`. + + Parameters + ---------- + config : polaris.config.PolarisConfigParser + Configuration options with parameters defining the equation of state + and ``rho0`` for the pseudo-height. + + 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 pressure at the top of the water column. + + 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 pressure at layer interfaces. + + p_mid : xarray.DataArray + The pressure at layer midpoints. + + spec_vol : xarray.DataArray + The specific volume at layer midpoints. + """ + + rho0 = config.getfloat('vertical_grid', 'rho0') + if rho0 is None: + raise ValueError( + 'Config option [vertical_grid] rho0 must be set to use ' + 'pressure_and_spec_vol_from_state_at_geom_height().' + ) + + spec_vol = 1.0 / rho0 * 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 + + +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, + rho0: float, +) -> 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. + + rho0 : float + Reference density in kg m^-3. + + 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 * rho0).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 diff --git a/polaris/tasks/ocean/overflow/init.py b/polaris/tasks/ocean/overflow/init.py index b3ee04a363..48f05f36e1 100644 --- a/polaris/tasks/ocean/overflow/init.py +++ b/polaris/tasks/ocean/overflow/init.py @@ -4,8 +4,8 @@ from mpas_tools.planar_hex import make_planar_hex_mesh from polaris.mesh.planar import compute_planar_hex_nx_ny +from polaris.ocean.eos import compute_density from polaris.ocean.model import OceanIOStep -from polaris.ocean.model.eos import compute_density from polaris.ocean.vertical import init_vertical_coord @@ -112,15 +112,7 @@ def run(self): salinity = section.getfloat('salinity') * np.ones_like(temperature) ds['salinity'] = salinity * xr.ones_like(ds.temperature) - density = compute_density(config, temperature, salinity) - ds['density'] = ( - ( - 'Time', - 'nCells', - 'nVertLevels', - ), - np.expand_dims(density, axis=0), - ) + ds['density'] = compute_density(config, ds.temperature, ds.salinity) # initial velocity on edges is stationary ds['normalVelocity'] = (