From e60d57e47e8e65801b64612187d8b09f8b9b287f Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Mon, 3 Nov 2025 10:54:33 +0100 Subject: [PATCH 1/6] 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 d4386666e54e4dbd73f9129060b91ad9555c4d41 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Mon, 3 Nov 2025 10:56:15 +0100 Subject: [PATCH 2/6] Add TEOS-10 as supported EOS --- polaris/ocean/eos/__init__.py | 37 ++++++++++++++++++-- polaris/ocean/eos/linear.py | 18 +++++++++- polaris/ocean/eos/teos10.py | 63 +++++++++++++++++++++++++++++++++++ 3 files changed, 115 insertions(+), 3 deletions(-) create mode 100644 polaris/ocean/eos/teos10.py diff --git a/polaris/ocean/eos/__init__.py b/polaris/ocean/eos/__init__.py index 9f1206a6e5..60d7d2d19a 100644 --- a/polaris/ocean/eos/__init__.py +++ b/polaris/ocean/eos/__init__.py @@ -1,7 +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. @@ -28,12 +38,26 @@ 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( + sa=salinity, ct=temperature, p=pressure + ) else: raise ValueError(f'Unsupported equation of state type: {eos_type}') 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. @@ -60,6 +84,15 @@ 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( + 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/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 diff --git a/polaris/ocean/eos/teos10.py b/polaris/ocean/eos/teos10.py new file mode 100644 index 0000000000..59a49c11de --- /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 f1ec97f3c8597674f69b76276d4cca1fc37eaa62 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Tue, 4 Nov 2025 08:41:30 -0600 Subject: [PATCH 3/6] Add a module for z-tilde vertical coord. --- polaris/ocean/vertical/__init__.py | 4 +- polaris/ocean/vertical/ztilde.py | 303 +++++++++++++++++++++++++++++ 2 files changed, 305 insertions(+), 2 deletions(-) create mode 100644 polaris/ocean/vertical/ztilde.py 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) diff --git a/polaris/ocean/vertical/ztilde.py b/polaris/ocean/vertical/ztilde.py new file mode 100644 index 0000000000..11c9b92a91 --- /dev/null +++ b/polaris/ocean/vertical/ztilde.py @@ -0,0 +1,303 @@ +""" +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 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', + 'geom_height_from_pseudo_height', +] + + +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 = -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. + """ + + 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') + 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, + 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 42e89664334b7297b2814ea79bd55dda0ce19ad6 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Thu, 6 Nov 2025 12:10:28 -0600 Subject: [PATCH 4/6] Add compute_zint_zmid_from_layer_thickness() This function can compute zMid and zInterface for any supported veritcal coordinate. --- polaris/ocean/vertical/__init__.py | 99 ++++++++++++++++++++---------- 1 file changed, 65 insertions(+), 34 deletions(-) diff --git a/polaris/ocean/vertical/__init__.py b/polaris/ocean/vertical/__init__.py index 52f11ccd0a..ec137faec3 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['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 @@ -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 From 7562c365b9e2d4766a1d5a5eba862b0bccb1f3b6 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Fri, 19 Dec 2025 10:40:10 -0600 Subject: [PATCH 5/6] Update mpas-ocean --> omega map Switch layerThickness --> GeometricLayerThickness Add a few vertical coordinate variables --- polaris/ocean/model/mpaso_to_omega.yaml | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/polaris/ocean/model/mpaso_to_omega.yaml b/polaris/ocean/model/mpaso_to_omega.yaml index 15b5882835..a33a33aead 100644 --- a/polaris/ocean/model/mpaso_to_omega.yaml +++ b/polaris/ocean/model/mpaso_to_omega.yaml @@ -30,7 +30,7 @@ variables: tracer3: Debug3 # state - layerThickness: LayerThickness + layerThickness: GeometricLayerThickness normalVelocity: NormalVelocity # auxiliary state @@ -39,6 +39,12 @@ variables: windStressMeridional: WindStressMeridional relativeVorticity: RelVortVertex + # vertical coordinate + zMid: ZMid + zInterface: ZInterface + pressure: PressureMid + + config: - section: time_management: TimeIntegration From b465575fab1fa04cf5d6bab81c51a6a1750e4ba9 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Fri, 19 Dec 2025 10:41:18 -0600 Subject: [PATCH 6/6] Update Omega CTest meshes The new meshes have as many variables as possible mapped to the new Omega names, and include PseudoThickness (if they previously had layerThickness) as well as GeometricLayerThickness. The PseudoThickness is computed using TEOS-10 for the specific volume. --- utils/omega/ctest/omega_ctest.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/utils/omega/ctest/omega_ctest.py b/utils/omega/ctest/omega_ctest.py index 92b5d7fbb1..c1ee1a981f 100755 --- a/utils/omega/ctest/omega_ctest.py +++ b/utils/omega/ctest/omega_ctest.py @@ -89,9 +89,9 @@ def download_meshes(config): base_url = 'https://web.lcrc.anl.gov/public/e3sm/polaris/' files = [ - 'ocean.QU.240km.151209.nc', - 'PlanarPeriodic48x48.nc', - 'cosine_bell_icos480_initial_state.230220.nc', + 'ocean.QU.240km.omega_vars.251219.nc', + 'PlanarPeriodic48x48.omega_vars.251219.nc', + 'cosine_bell_icos480.omega_vars.250219.nc', ] database_path = 'ocean/omega_ctest'