From 73035ff58a65a8e8fd0a7b852e2a65fbb4f0a1fa Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Mon, 3 Nov 2025 10:54:33 +0100 Subject: [PATCH 01/23] Move eos module and add optional pressure argument This merge also adds docstrings and adds a compute_specvol() funciton --- .../fleshing_out_step.md | 2 +- polaris/ocean/eos/__init__.py | 65 +++++++++++++++++++ polaris/ocean/eos/linear.py | 31 +++++++++ polaris/ocean/model/eos.py | 16 ----- polaris/tasks/ocean/overflow/init.py | 2 +- 5 files changed, 98 insertions(+), 18 deletions(-) create mode 100644 polaris/ocean/eos/__init__.py create mode 100644 polaris/ocean/eos/linear.py delete mode 100644 polaris/ocean/model/eos.py 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..9f1206a6e5 --- /dev/null +++ b/polaris/ocean/eos/__init__.py @@ -0,0 +1,65 @@ +from .linear import compute_linear_density + + +def compute_density(config, temperature, salinity, pressure=None): + """ + 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) + else: + raise ValueError(f'Unsupported equation of state type: {eos_type}') + return density + + +def compute_specvol(config, temperature, salinity, pressure=None): + """ + 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) + else: + raise ValueError(f'Unsupported equation of state type: {eos_type}') + return specvol diff --git a/polaris/ocean/eos/linear.py b/polaris/ocean/eos/linear.py new file mode 100644 index 0000000000..ea53eefdb4 --- /dev/null +++ b/polaris/ocean/eos/linear.py @@ -0,0 +1,31 @@ +def compute_linear_density(config, temperature, salinity): + """ + 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') + density = rhoref + -alpha * (temperature - Tref) + beta * (salinity - Sref) + return density 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/tasks/ocean/overflow/init.py b/polaris/tasks/ocean/overflow/init.py index b3ee04a363..fbbcf71faf 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 From 3335d166475ceff1c314b4a62729455ad28bd179 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Mon, 3 Nov 2025 10:56:15 +0100 Subject: [PATCH 02/23] Add TEOS-10 as supported EOS --- polaris/ocean/eos/__init__.py | 15 +++++++++ polaris/ocean/eos/teos10.py | 63 +++++++++++++++++++++++++++++++++++ 2 files changed, 78 insertions(+) create mode 100644 polaris/ocean/eos/teos10.py diff --git a/polaris/ocean/eos/__init__.py b/polaris/ocean/eos/__init__.py index 9f1206a6e5..6a7dbd767f 100644 --- a/polaris/ocean/eos/__init__.py +++ b/polaris/ocean/eos/__init__.py @@ -1,4 +1,5 @@ from .linear import compute_linear_density +from .teos10 import compute_specvol as compute_teos10_specvol def compute_density(config, temperature, salinity, pressure=None): @@ -28,6 +29,13 @@ def compute_density(config, temperature, salinity, pressure=None): 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(salinity, temperature, pressure) else: raise ValueError(f'Unsupported equation of state type: {eos_type}') return density @@ -60,6 +68,13 @@ def compute_specvol(config, temperature, salinity, pressure=None): 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(temperature, salinity, pressure) else: raise ValueError(f'Unsupported equation of state type: {eos_type}') return specvol diff --git a/polaris/ocean/eos/teos10.py b/polaris/ocean/eos/teos10.py new file mode 100644 index 0000000000..19d708fab7 --- /dev/null +++ b/polaris/ocean/eos/teos10.py @@ -0,0 +1,63 @@ +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.assign_attrs( + { + 'long_name': 'specific volume', + 'units': 'm^3 kg^-1', + 'standard_name': 'specific_volume', + } + ) From 6d9dccbb5e68c80ae0c63cc8d29df5d859711c6c Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Tue, 4 Nov 2025 08:41:30 -0600 Subject: [PATCH 03/23] Add a module for converting p <--> z_tilde --- polaris/ocean/vertical/ztilde.py | 84 ++++++++++++++++++++++++++++++++ 1 file changed, 84 insertions(+) create mode 100644 polaris/ocean/vertical/ztilde.py diff --git a/polaris/ocean/vertical/ztilde.py b/polaris/ocean/vertical/ztilde.py new file mode 100644 index 0000000000..f9640bb904 --- /dev/null +++ b/polaris/ocean/vertical/ztilde.py @@ -0,0 +1,84 @@ +""" +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``. +""" + +from typing import Union + +import xarray as xr +from mpas_tools.cime.constants import constants + +ArrayLike = Union[xr.DataArray] + +__all__ = [ + 'z_tilde_from_pressure', + 'pressure_from_z_tilde', +] + + +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). + """ + + g = constants['SHR_CONST_G'] + + z = -(p) / (rho0 * g) + 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). + """ + + g = constants['SHR_CONST_G'] + + p = -(z_tilde) * (rho0 * g) + return p.assign_attrs( + { + 'long_name': 'sea pressure', + 'units': 'Pa', + 'note': 'p = -z_tilde * (rho0 * g)', + } + ) From d4d0229593d50b54e4a2725d50fb0d69a0d50ef0 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Wed, 5 Nov 2025 05:26:39 -0600 Subject: [PATCH 04/23] Add z-tilde option to vertical module It is the same as z-star in this context --- polaris/ocean/vertical/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/polaris/ocean/vertical/__init__.py b/polaris/ocean/vertical/__init__.py index 40f422872e..52f11ccd0a 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) @@ -148,7 +148,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) From 49e8b2b0a8f5598c980c5a462cd08bad3e3e1b8f Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Wed, 5 Nov 2025 05:28:12 -0600 Subject: [PATCH 05/23] Add a function for computing geometric z from z_tilde --- polaris/ocean/vertical/ztilde.py | 68 ++++++++++++++++++++++++++++++-- 1 file changed, 64 insertions(+), 4 deletions(-) diff --git a/polaris/ocean/vertical/ztilde.py b/polaris/ocean/vertical/ztilde.py index f9640bb904..4dda4a57c0 100644 --- a/polaris/ocean/vertical/ztilde.py +++ b/polaris/ocean/vertical/ztilde.py @@ -9,16 +9,13 @@ ``mpas_tools.cime.constants``. """ -from typing import Union - import xarray as xr from mpas_tools.cime.constants import constants -ArrayLike = Union[xr.DataArray] - __all__ = [ 'z_tilde_from_pressure', 'pressure_from_z_tilde', + 'z_from_z_tilde', ] @@ -82,3 +79,66 @@ def pressure_from_z_tilde(z_tilde: xr.DataArray, rho0: float) -> xr.DataArray: 'note': 'p = -z_tilde * (rho0 * g)', } ) + + +def z_from_z_tilde( + layer_thickness: xr.DataArray, + bottom_depth: xr.DataArray, + spec_vol: xr.DataArray, + rho0: float, +) -> tuple[xr.DataArray, xr.DataArray]: + """ + Compute geometric height z at layer interfaces and midpoints given the + layer thicknesses, bottom depth, specific volume and reference density. + This calculation assumes a constant specific volume within each layer. + + Parameters + ---------- + layer_thickness : xarray.DataArray + The pseudo thickness of each layer. + + spec_vol : xarray.DataArray + The specific volume at each layer. + + bottom_depth : xarray.DataArray + The positive-down depth of the seafloor. + + rho0 : float + Reference density in kg m^-3. + + Returns + ------- + z_interface : xarray.DataArray + The elevation of layer interfaces. + + z_mid : xarray.DataArray + The elevation of layer centers. + """ + + n_vert_levels = layer_thickness.sizes['nVertLevels'] + + z_bot = -bottom_depth + z_interface_list = [z_bot] + z_mid_list = [] + + for k in range(n_vert_levels - 1, -1, -1): + dz = ( + rho0 + * spec_vol.isel(nVertLevels=k) + * layer_thickness.isel(nVertLevels=k) + ) + z_top = z_bot + dz + z_interface_list.append(z_top) + z_mid_list.append(z_bot + 0.5 * dz) + z_bot = z_top + + dims = list(layer_thickness.dims) + interface_dims = list(dims) + ['nVertLevelsP1'] + interface_dims.remove('nVertLevels') + + z_interface = xr.concat( + reversed(z_interface_list), dim='nVertLevelsP1' + ).transpose(*interface_dims) + z_mid = xr.concat(reversed(z_mid_list), dim='nVertLevels').transpose(*dims) + + return z_interface, z_mid From 09e2ed05d7b5d071dfa825771e4a4cf725ff6e6e Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Wed, 5 Nov 2025 07:39:02 -0600 Subject: [PATCH 06/23] Add functions for computing pressure for z-tilde One function computes the pressure given geometric layer thickness and specific volume, while the other computes pressure and specific volume iteratively, given geometric thickness, temperature and salinity (among other things). --- polaris/ocean/vertical/ztilde.py | 181 +++++++++++++++++++++++++++++-- 1 file changed, 172 insertions(+), 9 deletions(-) diff --git a/polaris/ocean/vertical/ztilde.py b/polaris/ocean/vertical/ztilde.py index 4dda4a57c0..da0ed8493f 100644 --- a/polaris/ocean/vertical/ztilde.py +++ b/polaris/ocean/vertical/ztilde.py @@ -9,12 +9,20 @@ ``mpas_tools.cime.constants``. """ +import logging + +import numpy as np import xarray as xr from mpas_tools.cime.constants import constants +from polaris.config import PolarisConfigParser +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', 'z_from_z_tilde', ] @@ -81,20 +89,175 @@ def pressure_from_z_tilde(z_tilde: xr.DataArray, rho0: float) -> xr.DataArray: ) +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. + + 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. + """ + + n_vert_levels = geom_layer_thickness.sizes['nVertLevels'] + g = constants['SHR_CONST_G'] + + p_top = surf_pressure + p_interface_list = [p_top] + p_mid_list = [] + + for k in range(n_vert_levels): + dp = ( + g + / spec_vol.isel(nVertLevels=k) + * geom_layer_thickness.isel(nVertLevels=k) + ) + p_bot = p_top + dp + p_interface_list.append(p_bot) + p_mid_list.append(p_top + 0.5 * dp) + p_top = p_bot + + dims = list(geom_layer_thickness.dims) + interface_dims = list(dims) + ['nVertLevelsP1'] + interface_dims.remove('nVertLevels') + + p_interface = xr.concat(p_interface_list, dim='nVertLevelsP1').transpose( + *interface_dims + ) + p_mid = xr.concat(p_mid_list, dim='nVertLevels').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') + + 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 z_from_z_tilde( - layer_thickness: xr.DataArray, + pseudo_thickness: xr.DataArray, bottom_depth: xr.DataArray, spec_vol: xr.DataArray, rho0: float, ) -> tuple[xr.DataArray, xr.DataArray]: """ - Compute geometric height z at layer interfaces and midpoints given the - layer thicknesses, bottom depth, specific volume and reference density. - This calculation assumes a constant specific volume within each layer. + Compute geometric height z at layer interfaces and midpoints given + pseudo-layer thicknesses, bottom depth, specific volume and reference + density. This calculation assumes a constant specific volume within each + layer. Parameters ---------- - layer_thickness : xarray.DataArray + pseudo_thickness : xarray.DataArray The pseudo thickness of each layer. spec_vol : xarray.DataArray @@ -112,10 +275,10 @@ def z_from_z_tilde( The elevation of layer interfaces. z_mid : xarray.DataArray - The elevation of layer centers. + The elevation of layer midpoints. """ - n_vert_levels = layer_thickness.sizes['nVertLevels'] + n_vert_levels = pseudo_thickness.sizes['nVertLevels'] z_bot = -bottom_depth z_interface_list = [z_bot] @@ -125,14 +288,14 @@ def z_from_z_tilde( dz = ( rho0 * spec_vol.isel(nVertLevels=k) - * layer_thickness.isel(nVertLevels=k) + * pseudo_thickness.isel(nVertLevels=k) ) z_top = z_bot + dz z_interface_list.append(z_top) z_mid_list.append(z_bot + 0.5 * dz) z_bot = z_top - dims = list(layer_thickness.dims) + dims = list(pseudo_thickness.dims) interface_dims = list(dims) + ['nVertLevelsP1'] interface_dims.remove('nVertLevels') From 44f1a86ac247e06bee58e07285f401082538698e Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Thu, 6 Nov 2025 09:10:34 -0600 Subject: [PATCH 07/23] Add type annotation to eos module --- polaris/ocean/eos/__init__.py | 18 ++++++++++++++++-- polaris/ocean/eos/linear.py | 18 +++++++++++++++++- 2 files changed, 33 insertions(+), 3 deletions(-) diff --git a/polaris/ocean/eos/__init__.py b/polaris/ocean/eos/__init__.py index 6a7dbd767f..6fdba65d9e 100644 --- a/polaris/ocean/eos/__init__.py +++ b/polaris/ocean/eos/__init__.py @@ -1,8 +1,17 @@ +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, temperature, salinity, pressure=None): +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. @@ -41,7 +50,12 @@ def compute_density(config, temperature, salinity, pressure=None): return density -def compute_specvol(config, temperature, salinity, pressure=None): +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. diff --git a/polaris/ocean/eos/linear.py b/polaris/ocean/eos/linear.py index ea53eefdb4..2ac76bd1d6 100644 --- a/polaris/ocean/eos/linear.py +++ b/polaris/ocean/eos/linear.py @@ -1,4 +1,13 @@ -def compute_linear_density(config, temperature, salinity): +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 @@ -27,5 +36,12 @@ def compute_linear_density(config, temperature, salinity): 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 From 94be927966e58df291484aee61fed657921def4a Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Thu, 6 Nov 2025 09:11:10 -0600 Subject: [PATCH 08/23] Rename geom_z_from_z_tilde() --- polaris/ocean/vertical/ztilde.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/polaris/ocean/vertical/ztilde.py b/polaris/ocean/vertical/ztilde.py index da0ed8493f..9a7836a211 100644 --- a/polaris/ocean/vertical/ztilde.py +++ b/polaris/ocean/vertical/ztilde.py @@ -23,7 +23,7 @@ 'pressure_from_z_tilde', 'pressure_and_spec_vol_from_state_at_geom_height', 'pressure_from_geom_thickness', - 'z_from_z_tilde', + 'geom_z_from_z_tilde', ] @@ -243,7 +243,7 @@ def pressure_and_spec_vol_from_state_at_geom_height( return p_interface, p_mid, spec_vol -def z_from_z_tilde( +def geom_z_from_z_tilde( pseudo_thickness: xr.DataArray, bottom_depth: xr.DataArray, spec_vol: xr.DataArray, From 6da806e76b6f7b2c24d635ee55689086f568f68f Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Thu, 6 Nov 2025 09:47:28 -0600 Subject: [PATCH 09/23] Fix teos10 variable names and calls --- polaris/ocean/eos/__init__.py | 8 ++++++-- polaris/ocean/eos/teos10.py | 30 +++++++++++++++--------------- 2 files changed, 21 insertions(+), 17 deletions(-) diff --git a/polaris/ocean/eos/__init__.py b/polaris/ocean/eos/__init__.py index 6fdba65d9e..60d7d2d19a 100644 --- a/polaris/ocean/eos/__init__.py +++ b/polaris/ocean/eos/__init__.py @@ -44,7 +44,9 @@ def compute_density( 'Pressure must be provided when using the TEOS-10 equation of ' 'state.' ) - density = 1.0 / compute_teos10_specvol(salinity, temperature, pressure) + density = 1.0 / compute_teos10_specvol( + sa=salinity, ct=temperature, p=pressure + ) else: raise ValueError(f'Unsupported equation of state type: {eos_type}') return density @@ -88,7 +90,9 @@ def compute_specvol( 'Pressure must be provided when using the TEOS-10 equation of ' 'state.' ) - specvol = compute_teos10_specvol(temperature, salinity, pressure) + specvol = compute_teos10_specvol( + sa=salinity, ct=temperature, p=pressure + ) else: raise ValueError(f'Unsupported equation of state type: {eos_type}') return specvol diff --git a/polaris/ocean/eos/teos10.py b/polaris/ocean/eos/teos10.py index 19d708fab7..59a49c11de 100644 --- a/polaris/ocean/eos/teos10.py +++ b/polaris/ocean/eos/teos10.py @@ -3,7 +3,7 @@ def compute_specvol( - SA: xr.DataArray, CT: xr.DataArray, p: xr.DataArray + sa: xr.DataArray, ct: xr.DataArray, p: xr.DataArray ) -> xr.DataArray: """ Compute specific volume from co-located p, CT and SA. @@ -19,39 +19,39 @@ def compute_specvol( Parameters ---------- - SA : xarray.DataArray - Absolute Salinity at the same points as p and CT. + 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. + 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. + 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). + 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): + 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}' + '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') + 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) + 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' + specvol_np, dims=ct.dims, coords=ct.coords, name='specvol' ) return specvol.assign_attrs( From a23844032dffcc257876faff03acf9ce876fff4b Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Thu, 6 Nov 2025 09:48:04 -0600 Subject: [PATCH 10/23] Fix masking in geom_z_from_z_tilde() --- polaris/ocean/vertical/ztilde.py | 20 +++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/polaris/ocean/vertical/ztilde.py b/polaris/ocean/vertical/ztilde.py index 9a7836a211..1b7e81a859 100644 --- a/polaris/ocean/vertical/ztilde.py +++ b/polaris/ocean/vertical/ztilde.py @@ -248,6 +248,8 @@ def geom_z_from_z_tilde( bottom_depth: xr.DataArray, spec_vol: xr.DataArray, rho0: float, + min_level_cell: xr.DataArray, + max_level_cell: xr.DataArray, ) -> tuple[xr.DataArray, xr.DataArray]: """ Compute geometric height z at layer interfaces and midpoints given @@ -269,6 +271,12 @@ def geom_z_from_z_tilde( rho0 : float Reference density in kg m^-3. + 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 ------- z_interface : xarray.DataArray @@ -281,7 +289,9 @@ def geom_z_from_z_tilde( n_vert_levels = pseudo_thickness.sizes['nVertLevels'] z_bot = -bottom_depth - z_interface_list = [z_bot] + k = n_vert_levels + mask_bot = np.logical_and(k >= min_level_cell, k <= max_level_cell) + z_interface_list = [z_bot.where(mask_bot)] z_mid_list = [] for k in range(n_vert_levels - 1, -1, -1): @@ -290,9 +300,13 @@ def geom_z_from_z_tilde( * spec_vol.isel(nVertLevels=k) * pseudo_thickness.isel(nVertLevels=k) ) + mask_mid = np.logical_and(k >= min_level_cell, k < max_level_cell) + mask_top = np.logical_and(k >= min_level_cell, k <= max_level_cell) + dz = dz.where(mask_mid, 0.0) z_top = z_bot + dz - z_interface_list.append(z_top) - z_mid_list.append(z_bot + 0.5 * dz) + z_interface_list.append(z_top.where(mask_top)) + z_mid = (z_bot + 0.5 * dz).where(mask_mid) + z_mid_list.append(z_mid) z_bot = z_top dims = list(pseudo_thickness.dims) From 952447c62e067b5bfffae7680ab8e470d83d2beb Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Thu, 6 Nov 2025 12:10:28 -0600 Subject: [PATCH 11/23] Switch to general compute_zint_zmid_from_layer_thickness() --- polaris/ocean/vertical/__init__.py | 99 ++++++++++++++++++++---------- polaris/ocean/vertical/ztilde.py | 79 ------------------------ 2 files changed, 65 insertions(+), 113 deletions(-) diff --git a/polaris/ocean/vertical/__init__.py b/polaris/ocean/vertical/__init__.py index 52f11ccd0a..99b90987d5 100644 --- a/polaris/ocean/vertical/__init__.py +++ b/polaris/ocean/vertical/__init__.py @@ -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['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 @@ -162,47 +165,75 @@ 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. + + bottom_depth : xarray.DataArray + The positive-down depth of the seafloor. - ssh : xarray.DataArray - The sea surface height + min_level_cell : xarray.DataArray + The zero-based minimum vertical index from each column. - cellMask : xarray.DataArray - A boolean mask of where there are valid cells + 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 = [] + n_vert_levels = layer_thickness.sizes['nVertLevels'] + + z_bot = -bottom_depth + k = n_vert_levels + mask_bot = np.logical_and(k >= min_level_cell, k - 1 <= max_level_cell) + z_interface_list = [z_bot.where(mask_bot)] + z_mid_list = [] + + for k in range(n_vert_levels - 1, -1, -1): + dz = layer_thickness.isel(nVertLevels=k) + mask_mid = np.logical_and(k >= min_level_cell, k <= max_level_cell) + mask_top = np.logical_and(k >= min_level_cell, k - 1 <= max_level_cell) + dz = dz.where(mask_mid, 0.0) + z_top = z_bot + dz + z_interface_list.append(z_top.where(mask_top)) + z_mid = (z_bot + 0.5 * dz).where(mask_mid) + z_mid_list.append(z_mid) + z_bot = z_top + + dims = list(layer_thickness.dims) + interface_dims = list(dims) + ['nVertLevelsP1'] + interface_dims.remove('nVertLevels') + + z_interface = xr.concat( + reversed(z_interface_list), dim='nVertLevelsP1' + ).transpose(*interface_dims) + z_mid = xr.concat(reversed(z_mid_list), dim='nVertLevels').transpose(*dims) + + return z_interface, z_mid + + +def _compute_cell_mask(minLevelCell, maxLevelCell, nVertLevels): + cellMask = [] 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' - ) - return zMid + 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 diff --git a/polaris/ocean/vertical/ztilde.py b/polaris/ocean/vertical/ztilde.py index 1b7e81a859..77b368b752 100644 --- a/polaris/ocean/vertical/ztilde.py +++ b/polaris/ocean/vertical/ztilde.py @@ -23,7 +23,6 @@ 'pressure_from_z_tilde', 'pressure_and_spec_vol_from_state_at_geom_height', 'pressure_from_geom_thickness', - 'geom_z_from_z_tilde', ] @@ -241,81 +240,3 @@ def pressure_and_spec_vol_from_state_at_geom_height( ) return p_interface, p_mid, spec_vol - - -def geom_z_from_z_tilde( - pseudo_thickness: xr.DataArray, - bottom_depth: xr.DataArray, - spec_vol: xr.DataArray, - rho0: float, - min_level_cell: xr.DataArray, - max_level_cell: xr.DataArray, -) -> tuple[xr.DataArray, xr.DataArray]: - """ - Compute geometric height z at layer interfaces and midpoints given - pseudo-layer thicknesses, bottom depth, specific volume and reference - density. This calculation assumes a constant specific volume within each - layer. - - Parameters - ---------- - pseudo_thickness : xarray.DataArray - The pseudo thickness of each layer. - - spec_vol : xarray.DataArray - The specific volume at each layer. - - bottom_depth : xarray.DataArray - The positive-down depth of the seafloor. - - rho0 : float - Reference density in kg m^-3. - - 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 - ------- - z_interface : xarray.DataArray - The elevation of layer interfaces. - - z_mid : xarray.DataArray - The elevation of layer midpoints. - """ - - n_vert_levels = pseudo_thickness.sizes['nVertLevels'] - - z_bot = -bottom_depth - k = n_vert_levels - mask_bot = np.logical_and(k >= min_level_cell, k <= max_level_cell) - z_interface_list = [z_bot.where(mask_bot)] - z_mid_list = [] - - for k in range(n_vert_levels - 1, -1, -1): - dz = ( - rho0 - * spec_vol.isel(nVertLevels=k) - * pseudo_thickness.isel(nVertLevels=k) - ) - mask_mid = np.logical_and(k >= min_level_cell, k < max_level_cell) - mask_top = np.logical_and(k >= min_level_cell, k <= max_level_cell) - dz = dz.where(mask_mid, 0.0) - z_top = z_bot + dz - z_interface_list.append(z_top.where(mask_top)) - z_mid = (z_bot + 0.5 * dz).where(mask_mid) - z_mid_list.append(z_mid) - z_bot = z_top - - dims = list(pseudo_thickness.dims) - interface_dims = list(dims) + ['nVertLevelsP1'] - interface_dims.remove('nVertLevels') - - z_interface = xr.concat( - reversed(z_interface_list), dim='nVertLevelsP1' - ).transpose(*interface_dims) - z_mid = xr.concat(reversed(z_mid_list), dim='nVertLevels').transpose(*dims) - - return z_interface, z_mid From be3a2f886f05a793fe86b8cb4d1446e597f131a3 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Mon, 10 Nov 2025 15:36:17 +0100 Subject: [PATCH 12/23] Add zInterface along with zMid in vertical coordinate --- polaris/ocean/vertical/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/polaris/ocean/vertical/__init__.py b/polaris/ocean/vertical/__init__.py index 99b90987d5..ec137faec3 100644 --- a/polaris/ocean/vertical/__init__.py +++ b/polaris/ocean/vertical/__init__.py @@ -110,7 +110,7 @@ def init_vertical_coord(config, ds): dim='Time', axis=0 ) - _, ds['zMid'] = compute_zint_zmid_from_layer_thickness( + ds['zInterface'], ds['zMid'] = compute_zint_zmid_from_layer_thickness( layer_thickness=ds.layerThickness, bottom_depth=ds.bottomDepth, min_level_cell=ds.minLevelCell, From 294c246f4879dd4aabd0246a062b007347f31d5a Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Fri, 19 Dec 2025 15:22:10 +0100 Subject: [PATCH 13/23] Add geom_height_from_pseudo_height() function This computes geometric height from pseudo-height, complimenting the functionality we already have to compute pseudo-height from geometric height. --- polaris/ocean/vertical/ztilde.py | 64 +++++++++++++++++++++++++++++++- 1 file changed, 62 insertions(+), 2 deletions(-) diff --git a/polaris/ocean/vertical/ztilde.py b/polaris/ocean/vertical/ztilde.py index 77b368b752..36bef3692b 100644 --- a/polaris/ocean/vertical/ztilde.py +++ b/polaris/ocean/vertical/ztilde.py @@ -36,6 +36,7 @@ def z_tilde_from_pressure(p: xr.DataArray, rho0: float) -> xr.DataArray: ---------- p : xarray.DataArray Sea pressure in Pascals (Pa). + rho0 : float Reference density in kg m^-3. @@ -67,6 +68,7 @@ def pressure_from_z_tilde(z_tilde: xr.DataArray, rho0: float) -> xr.DataArray: ---------- z_tilde : xarray.DataArray Pseudo-height in meters (m), positive upward. + rho0 : float Reference density in kg m^-3. @@ -83,7 +85,7 @@ def pressure_from_z_tilde(z_tilde: xr.DataArray, rho0: float) -> xr.DataArray: { 'long_name': 'sea pressure', 'units': 'Pa', - 'note': 'p = -z_tilde * (rho0 * g)', + 'note': 'p = -rho0 * g * z_tilde', } ) @@ -104,7 +106,7 @@ def pressure_from_geom_thickness( The surface pressure at the top of the water column. geom_layer_thickness : xarray.DataArray - The geometric thickness of each layer. + The geometric thickness of each layer, set to zero for invalid layers. spec_vol : xarray.DataArray The specific volume at each layer. @@ -205,6 +207,11 @@ def pressure_and_spec_vol_from_state_at_geom_height( """ 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) @@ -240,3 +247,56 @@ def pressure_and_spec_vol_from_state_at_geom_height( ) 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, + 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. + + 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'] + + geom_thickness = spec_vol * h_tilde * rho0 + + geom_z_inter_list: list[xr.DataArray] = [geom_z_bot] + geom_z_mid_list: list[xr.DataArray] = [] + geom_z_prev = geom_z_bot + for z_index in range(n_vert_levels): + thickness = geom_thickness.isel(nVertLevels=z_index) + geom_z_next = geom_z_prev + thickness + geom_z_inter_list.append(geom_z_next) + geom_z_mid_list.append(geom_z_prev + 0.5 * thickness) + geom_z_prev = geom_z_next + + geom_z_inter = xr.concat(geom_z_inter_list, dim='nVertLevelsP1') + geom_z_mid = xr.concat(geom_z_mid_list, dim='nVertLevels') + return geom_z_inter, geom_z_mid From 5bc5fc16a1ff1e7519f1e2a9678c2bffa8888653 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Tue, 20 Jan 2026 14:23:19 +0100 Subject: [PATCH 14/23] Fix masking and order in geom_height_from_pseudo_height() --- polaris/ocean/vertical/ztilde.py | 39 ++++++++++++++++++++++++++++---- 1 file changed, 34 insertions(+), 5 deletions(-) diff --git a/polaris/ocean/vertical/ztilde.py b/polaris/ocean/vertical/ztilde.py index 36bef3692b..c44ff63cf2 100644 --- a/polaris/ocean/vertical/ztilde.py +++ b/polaris/ocean/vertical/ztilde.py @@ -253,6 +253,8 @@ 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]: """ @@ -269,6 +271,12 @@ def geom_height_from_pseudo_height( 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. @@ -287,16 +295,37 @@ def geom_height_from_pseudo_height( geom_thickness = spec_vol * h_tilde * rho0 - geom_z_inter_list: list[xr.DataArray] = [geom_z_bot] + inter_mask = np.logical_and( + n_vert_levels >= min_level_cell, n_vert_levels <= max_level_cell + 1 + ) + geom_z_inter_list: list[xr.DataArray] = [geom_z_bot.where(inter_mask)] geom_z_mid_list: list[xr.DataArray] = [] geom_z_prev = geom_z_bot - for z_index in range(n_vert_levels): - thickness = geom_thickness.isel(nVertLevels=z_index) + for z_index in range(n_vert_levels - 1, -1, -1): + mid_mask = np.logical_and( + z_index >= min_level_cell, z_index <= max_level_cell + ) + inter_mask = np.logical_and( + z_index >= min_level_cell, z_index <= max_level_cell + 1 + ) + thickness = xr.where( + mid_mask, geom_thickness.isel(nVertLevels=z_index), 0.0 + ) geom_z_next = geom_z_prev + thickness - geom_z_inter_list.append(geom_z_next) - geom_z_mid_list.append(geom_z_prev + 0.5 * thickness) + geom_z_mid = geom_z_prev + 0.5 * thickness + geom_z_inter_list.insert(0, geom_z_next.where(inter_mask)) + geom_z_mid_list.insert(0, geom_z_mid.where(mid_mask)) geom_z_prev = geom_z_next geom_z_inter = xr.concat(geom_z_inter_list, dim='nVertLevelsP1') geom_z_mid = xr.concat(geom_z_mid_list, dim='nVertLevels') + + # 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 From e268c9027943fafc003be03c5879f859cb57252b Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Tue, 20 Jan 2026 14:38:36 +0100 Subject: [PATCH 15/23] Add units and long name to density and specific volume --- polaris/ocean/eos/__init__.py | 4 ++++ polaris/ocean/eos/teos10.py | 8 +------- 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/polaris/ocean/eos/__init__.py b/polaris/ocean/eos/__init__.py index 60d7d2d19a..0afbfbd268 100644 --- a/polaris/ocean/eos/__init__.py +++ b/polaris/ocean/eos/__init__.py @@ -49,6 +49,8 @@ def compute_density( ) else: raise ValueError(f'Unsupported equation of state type: {eos_type}') + density.attrs['units'] = 'kg m-3' + density.attrs['long_name'] = 'density' return density @@ -95,4 +97,6 @@ def compute_specvol( ) 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/teos10.py b/polaris/ocean/eos/teos10.py index 59a49c11de..75e5ddace1 100644 --- a/polaris/ocean/eos/teos10.py +++ b/polaris/ocean/eos/teos10.py @@ -54,10 +54,4 @@ def compute_specvol( specvol_np, dims=ct.dims, coords=ct.coords, name='specvol' ) - return specvol.assign_attrs( - { - 'long_name': 'specific volume', - 'units': 'm^3 kg^-1', - 'standard_name': 'specific_volume', - } - ) + return specvol From 042e2dbadb122ecf1c94652e936b0ee1dfba1ce5 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Mon, 16 Feb 2026 10:32:10 -0600 Subject: [PATCH 16/23] Temporarily hard-code Omega's acceleration of gravity --- polaris/ocean/vertical/ztilde.py | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/polaris/ocean/vertical/ztilde.py b/polaris/ocean/vertical/ztilde.py index c44ff63cf2..ff6db949f5 100644 --- a/polaris/ocean/vertical/ztilde.py +++ b/polaris/ocean/vertical/ztilde.py @@ -13,7 +13,6 @@ import numpy as np import xarray as xr -from mpas_tools.cime.constants import constants from polaris.config import PolarisConfigParser from polaris.ocean.eos import compute_specvol @@ -25,6 +24,10 @@ 'pressure_from_geom_thickness', ] +# TODO: replace with value looked up in GCD YAML file when possible +# Temporarily hard-coded with the GCD/Omega value +Gravity = 9.80665 + def z_tilde_from_pressure(p: xr.DataArray, rho0: float) -> xr.DataArray: """ @@ -46,9 +49,7 @@ def z_tilde_from_pressure(p: xr.DataArray, rho0: float) -> xr.DataArray: Pseudo-height with the same shape and coords as ``p`` (units: m). """ - g = constants['SHR_CONST_G'] - - z = -(p) / (rho0 * g) + z = -(p) / (rho0 * Gravity) return z.assign_attrs( { 'long_name': 'pseudo-height', @@ -78,9 +79,7 @@ def pressure_from_z_tilde(z_tilde: xr.DataArray, rho0: float) -> xr.DataArray: Sea pressure with the same shape and coords as ``z_tilde`` (Pa). """ - g = constants['SHR_CONST_G'] - - p = -(z_tilde) * (rho0 * g) + p = -(z_tilde) * (rho0 * Gravity) return p.assign_attrs( { 'long_name': 'sea pressure', @@ -121,7 +120,6 @@ def pressure_from_geom_thickness( """ n_vert_levels = geom_layer_thickness.sizes['nVertLevels'] - g = constants['SHR_CONST_G'] p_top = surf_pressure p_interface_list = [p_top] @@ -129,7 +127,7 @@ def pressure_from_geom_thickness( for k in range(n_vert_levels): dp = ( - g + Gravity / spec_vol.isel(nVertLevels=k) * geom_layer_thickness.isel(nVertLevels=k) ) From 8925f707123663f9701227d3ccb8bfdb466c73d0 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Wed, 18 Feb 2026 06:36:42 -0600 Subject: [PATCH 17/23] Improve performance of vertical coordinate init --- polaris/ocean/vertical/__init__.py | 74 ++++++++++++++++++------------ polaris/ocean/vertical/zlevel.py | 56 ++++++++++------------ polaris/ocean/vertical/zstar.py | 28 +++++------ 3 files changed, 79 insertions(+), 79 deletions(-) diff --git a/polaris/ocean/vertical/__init__.py b/polaris/ocean/vertical/__init__.py index ec137faec3..04cea553e0 100644 --- a/polaris/ocean/vertical/__init__.py +++ b/polaris/ocean/vertical/__init__.py @@ -200,40 +200,54 @@ def compute_zint_zmid_from_layer_thickness( n_vert_levels = layer_thickness.sizes['nVertLevels'] - z_bot = -bottom_depth - k = n_vert_levels - mask_bot = np.logical_and(k >= min_level_cell, k - 1 <= max_level_cell) - z_interface_list = [z_bot.where(mask_bot)] - z_mid_list = [] - - for k in range(n_vert_levels - 1, -1, -1): - dz = layer_thickness.isel(nVertLevels=k) - mask_mid = np.logical_and(k >= min_level_cell, k <= max_level_cell) - mask_top = np.logical_and(k >= min_level_cell, k - 1 <= max_level_cell) - dz = dz.where(mask_mid, 0.0) - z_top = z_bot + dz - z_interface_list.append(z_top.where(mask_top)) - z_mid = (z_bot + 0.5 * dz).where(mask_mid) - z_mid_list.append(z_mid) - z_bot = z_top + 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 + ) - dims = list(layer_thickness.dims) - interface_dims = list(dims) + ['nVertLevelsP1'] - interface_dims.remove('nVertLevels') + 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) - z_interface = xr.concat( - reversed(z_interface_list), dim='nVertLevelsP1' - ).transpose(*interface_dims) - z_mid = xr.concat(reversed(z_mid_list), dim='nVertLevels').transpose(*dims) + 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): - 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 + 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..867efa716b 100644 --- a/polaris/ocean/vertical/zlevel.py +++ b/polaris/ocean/vertical/zlevel.py @@ -212,23 +212,21 @@ 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 = refBottomDepth.sizes['nVertLevels'] + z_index = xarray.DataArray( + numpy.arange(n_vert_levels), dims=['nVertLevels'] ) - layerThicknessArray = layerThicknessArray.transpose( - 'nCells', 'nVertLevels' + mask = numpy.logical_and( + z_index >= minLevelCell, z_index <= maxLevelCell + ).transpose('nCells', 'nVertLevels') + + z_top = xarray.where(ssh < -refTopDepth, ssh, -refTopDepth) + z_bot = xarray.where( + -bottomDepth > -refBottomDepth, -bottomDepth, -refBottomDepth ) - return layerThicknessArray + thickness = (z_top - z_bot).transpose('nCells', 'nVertLevels') + + return thickness.where(mask, 0.0) def compute_z_level_resting_thickness( @@ -261,21 +259,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'] - ) - restingThicknessArray = restingThicknessArray.transpose( - 'nCells', 'nVertLevels' + n_vert_levels = layerThickness.sizes['nVertLevels'] + z_index = xarray.DataArray( + numpy.arange(n_vert_levels), dims=['nVertLevels'] ) - return restingThicknessArray + mask = numpy.logical_and( + z_index >= minLevelCell, z_index <= maxLevelCell + ).transpose('nCells', 'nVertLevels') + + layer_stretch = bottomDepth / (ssh + bottomDepth) + resting_thickness = layer_stretch * layerThickness + + return resting_thickness.where(mask, 0.0) diff --git a/polaris/ocean/vertical/zstar.py b/polaris/ocean/vertical/zstar.py index f3c93a3267..c55ad28a49 100644 --- a/polaris/ocean/vertical/zstar.py +++ b/polaris/ocean/vertical/zstar.py @@ -156,23 +156,17 @@ def _compute_z_star_layer_thickness( The thickness of each layer (level) """ - nVertLevels = restingThickness.sizes['nVertLevels'] - layerThickness = [] + n_vert_levels = restingThickness.sizes['nVertLevels'] + z_index = xarray.DataArray( + numpy.arange(n_vert_levels), dims=['nVertLevels'] + ) + mask = numpy.logical_and( + z_index >= minLevelCell, z_index <= maxLevelCell + ).transpose('nCells', 'nVertLevels') - layerStretch = (ssh + bottomDepth) / restingThickness.sum( + layer_stretch = (ssh + bottomDepth) / restingThickness.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 * restingThickness + + return layer_thickness.where(mask, 0.0) From 14e18d99aeac3b8715616ddbc8d204ce1bc2b42f Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Wed, 18 Feb 2026 06:40:08 -0600 Subject: [PATCH 18/23] Switch vert coord parameter names to snake_case --- polaris/ocean/vertical/zlevel.py | 71 +++++++++++++++++++------------- polaris/ocean/vertical/zstar.py | 22 ++++++---- 2 files changed, 55 insertions(+), 38 deletions(-) diff --git a/polaris/ocean/vertical/zlevel.py b/polaris/ocean/vertical/zlevel.py index 867efa716b..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,17 +219,19 @@ def compute_z_level_layer_thickness( The thickness of each layer (level) """ - n_vert_levels = refBottomDepth.sizes['nVertLevels'] + n_vert_levels = ref_bottom_depth.sizes['nVertLevels'] z_index = xarray.DataArray( numpy.arange(n_vert_levels), dims=['nVertLevels'] ) mask = numpy.logical_and( - z_index >= minLevelCell, z_index <= maxLevelCell + z_index >= min_level_cell, z_index <= max_level_cell ).transpose('nCells', 'nVertLevels') - z_top = xarray.where(ssh < -refTopDepth, ssh, -refTopDepth) + z_top = xarray.where(ssh < -ref_top_depth, ssh, -ref_top_depth) z_bot = xarray.where( - -bottomDepth > -refBottomDepth, -bottomDepth, -refBottomDepth + -bottom_depth > -ref_bottom_depth, + -bottom_depth, + -ref_bottom_depth, ) thickness = (z_top - z_bot).transpose('nCells', 'nVertLevels') @@ -230,7 +239,11 @@ def compute_z_level_layer_thickness( 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 @@ -238,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 @@ -259,15 +272,15 @@ def compute_z_level_resting_thickness( The thickness of z-star layers when ssh = 0 """ - n_vert_levels = layerThickness.sizes['nVertLevels'] + n_vert_levels = layer_thickness.sizes['nVertLevels'] z_index = xarray.DataArray( numpy.arange(n_vert_levels), dims=['nVertLevels'] ) mask = numpy.logical_and( - z_index >= minLevelCell, z_index <= maxLevelCell + z_index >= min_level_cell, z_index <= max_level_cell ).transpose('nCells', 'nVertLevels') - layer_stretch = bottomDepth / (ssh + bottomDepth) - resting_thickness = layer_stretch * layerThickness + 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 c55ad28a49..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,17 +160,17 @@ def _compute_z_star_layer_thickness( The thickness of each layer (level) """ - n_vert_levels = restingThickness.sizes['nVertLevels'] + n_vert_levels = resting_thickness.sizes['nVertLevels'] z_index = xarray.DataArray( numpy.arange(n_vert_levels), dims=['nVertLevels'] ) mask = numpy.logical_and( - z_index >= minLevelCell, z_index <= maxLevelCell + z_index >= min_level_cell, z_index <= max_level_cell ).transpose('nCells', 'nVertLevels') - layer_stretch = (ssh + bottomDepth) / restingThickness.sum( + layer_stretch = (ssh + bottom_depth) / resting_thickness.sum( dim='nVertLevels' ) - layer_thickness = layer_stretch * restingThickness + layer_thickness = layer_stretch * resting_thickness return layer_thickness.where(mask, 0.0) From fe160f3c08192bedfefc74ef5dcbcc9dd57d1c27 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Wed, 18 Feb 2026 06:57:53 -0600 Subject: [PATCH 19/23] Improve performance of ztilde module --- polaris/ocean/vertical/ztilde.py | 86 ++++++++++++++++---------------- 1 file changed, 43 insertions(+), 43 deletions(-) diff --git a/polaris/ocean/vertical/ztilde.py b/polaris/ocean/vertical/ztilde.py index ff6db949f5..882f8d3f5b 100644 --- a/polaris/ocean/vertical/ztilde.py +++ b/polaris/ocean/vertical/ztilde.py @@ -119,31 +119,25 @@ def pressure_from_geom_thickness( The pressure at layer midpoints. """ - n_vert_levels = geom_layer_thickness.sizes['nVertLevels'] + dp = Gravity / spec_vol * geom_layer_thickness - p_top = surf_pressure - p_interface_list = [p_top] - p_mid_list = [] + 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'}) - for k in range(n_vert_levels): - dp = ( - Gravity - / spec_vol.isel(nVertLevels=k) - * geom_layer_thickness.isel(nVertLevels=k) - ) - p_bot = p_top + dp - p_interface_list.append(p_bot) - p_mid_list.append(p_top + 0.5 * dp) - p_top = p_bot + 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 = list(dims) + ['nVertLevelsP1'] - interface_dims.remove('nVertLevels') + interface_dims = [dim for dim in dims if dim != 'nVertLevels'] + interface_dims.append('nVertLevelsP1') - p_interface = xr.concat(p_interface_list, dim='nVertLevelsP1').transpose( - *interface_dims - ) - p_mid = xr.concat(p_mid_list, dim='nVertLevels').transpose(*dims) + p_interface = p_interface.transpose(*interface_dims) + p_mid = p_mid.transpose(*dims) return p_interface, p_mid @@ -291,32 +285,38 @@ def geom_height_from_pseudo_height( n_vert_levels = spec_vol.sizes['nVertLevels'] - geom_thickness = spec_vol * h_tilde * rho0 + 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( - n_vert_levels >= min_level_cell, n_vert_levels <= max_level_cell + 1 + z_index_p1 >= min_level_cell, + z_index_p1 <= max_level_cell + 1, ) - geom_z_inter_list: list[xr.DataArray] = [geom_z_bot.where(inter_mask)] - geom_z_mid_list: list[xr.DataArray] = [] - geom_z_prev = geom_z_bot - for z_index in range(n_vert_levels - 1, -1, -1): - mid_mask = np.logical_and( - z_index >= min_level_cell, z_index <= max_level_cell - ) - inter_mask = np.logical_and( - z_index >= min_level_cell, z_index <= max_level_cell + 1 - ) - thickness = xr.where( - mid_mask, geom_thickness.isel(nVertLevels=z_index), 0.0 - ) - geom_z_next = geom_z_prev + thickness - geom_z_mid = geom_z_prev + 0.5 * thickness - geom_z_inter_list.insert(0, geom_z_next.where(inter_mask)) - geom_z_mid_list.insert(0, geom_z_mid.where(mid_mask)) - geom_z_prev = geom_z_next - - geom_z_inter = xr.concat(geom_z_inter_list, dim='nVertLevelsP1') - geom_z_mid = xr.concat(geom_z_mid_list, dim='nVertLevels') + 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 From 82a221dcb775770c2a8f8cd56644731d215c0f0e Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Wed, 18 Feb 2026 08:19:53 -0600 Subject: [PATCH 20/23] Add eos to docs --- docs/developers_guide/ocean/api.md | 15 +++++++++++++++ docs/developers_guide/ocean/framework.md | 22 ++++++++++++++++++++++ 2 files changed, 37 insertions(+) diff --git a/docs/developers_guide/ocean/api.md b/docs/developers_guide/ocean/api.md index fa02a083cd..cd25956c55 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) diff --git a/docs/developers_guide/ocean/framework.md b/docs/developers_guide/ocean/framework.md index 4fc90ac867..5c58ecd9b0 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 From b82890fcf9ea94d2553d6106919ae1202ce4824b Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Wed, 18 Feb 2026 08:24:45 -0600 Subject: [PATCH 21/23] Add z-tilde docs and update vertical coord docs --- docs/developers_guide/ocean/api.md | 7 ++++++ docs/developers_guide/ocean/framework.md | 27 +++++++++++++++++++++--- 2 files changed, 31 insertions(+), 3 deletions(-) diff --git a/docs/developers_guide/ocean/api.md b/docs/developers_guide/ocean/api.md index cd25956c55..97a0ccea26 100644 --- a/docs/developers_guide/ocean/api.md +++ b/docs/developers_guide/ocean/api.md @@ -592,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 @@ -599,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 @@ -607,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 5c58ecd9b0..f07ef035be 100644 --- a/docs/developers_guide/ocean/framework.md +++ b/docs/developers_guide/ocean/framework.md @@ -556,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)= From 6f61b84407363125a3439f034f2bbef371b34c00 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Wed, 18 Feb 2026 10:49:36 -0600 Subject: [PATCH 22/23] Fix eos call in overflow init --- polaris/tasks/ocean/overflow/init.py | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/polaris/tasks/ocean/overflow/init.py b/polaris/tasks/ocean/overflow/init.py index fbbcf71faf..48f05f36e1 100644 --- a/polaris/tasks/ocean/overflow/init.py +++ b/polaris/tasks/ocean/overflow/init.py @@ -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'] = ( From 969f55eabbfb18c8e3973d5c3977b1535ed51906 Mon Sep 17 00:00:00 2001 From: Carolyn Begeman Date: Thu, 5 Mar 2026 11:10:35 -0600 Subject: [PATCH 23/23] Get gravity from PCD --- polaris/ocean/vertical/ztilde.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/polaris/ocean/vertical/ztilde.py b/polaris/ocean/vertical/ztilde.py index 882f8d3f5b..550ce0632b 100644 --- a/polaris/ocean/vertical/ztilde.py +++ b/polaris/ocean/vertical/ztilde.py @@ -15,6 +15,7 @@ import xarray as xr from polaris.config import PolarisConfigParser +from polaris.constants import get_constant from polaris.ocean.eos import compute_specvol __all__ = [ @@ -24,9 +25,7 @@ 'pressure_from_geom_thickness', ] -# TODO: replace with value looked up in GCD YAML file when possible -# Temporarily hard-coded with the GCD/Omega value -Gravity = 9.80665 +Gravity = get_constant('standard_acceleration_of_gravity') def z_tilde_from_pressure(p: xr.DataArray, rho0: float) -> xr.DataArray: