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..60d7d2d19a --- /dev/null +++ b/polaris/ocean/eos/__init__.py @@ -0,0 +1,98 @@ +import xarray as xr + +from polaris.config import PolarisConfigParser + +from .linear import compute_linear_density +from .teos10 import compute_specvol as compute_teos10_specvol + + +def compute_density( + config: PolarisConfigParser, + temperature: xr.DataArray, + salinity: xr.DataArray, + pressure: xr.DataArray | None = None, +) -> xr.DataArray: + """ + Compute the density of seawater based on the equation of state specified + in the configuration. + + Parameters + ---------- + config : polaris.config.PolarisConfigParser + Configuration object containing ocean parameters. + + temperature : float or xarray.DataArray + Temperature (conservative, potential or in-situ) of the seawater. + + salinity : float or xarray.DataArray + Salinity (practical or absolute) of the seawater. + + pressure : float or xarray.DataArray, optional + Pressure (in-situ or reference) of the seawater. + + Returns + ------- + density : float or xarray.DataArray + Computed density (in-situ or reference) of the seawater. + """ + eos_type = config.get('ocean', 'eos_type') + if eos_type == 'linear': + density = compute_linear_density(config, temperature, salinity) + elif eos_type == 'teos-10': + if pressure is None: + raise ValueError( + 'Pressure must be provided when using the TEOS-10 equation of ' + 'state.' + ) + density = 1.0 / compute_teos10_specvol( + sa=salinity, ct=temperature, p=pressure + ) + else: + raise ValueError(f'Unsupported equation of state type: {eos_type}') + return density + + +def compute_specvol( + config: PolarisConfigParser, + temperature: xr.DataArray, + salinity: xr.DataArray, + pressure: xr.DataArray | None = None, +) -> xr.DataArray: + """ + Compute the specific volume of seawater based on the equation of state + specified in the configuration. + + Parameters + ---------- + config : polaris.config.PolarisConfigParser + Configuration object containing ocean parameters. + + temperature : float or xarray.DataArray + Temperature (conservative, potential or in-situ) of the seawater. + + salinity : float or xarray.DataArray + Salinity (practical or absolute) of the seawater. + + pressure : float or xarray.DataArray, optional + Pressure (in-situ or reference) of the seawater. + + Returns + ------- + specvol : float or xarray.DataArray + Computed specific volume (in-situ or reference) of the seawater. + """ + eos_type = config.get('ocean', 'eos_type') + if eos_type == 'linear': + specvol = 1.0 / compute_linear_density(config, temperature, salinity) + elif eos_type == 'teos-10': + if pressure is None: + raise ValueError( + 'Pressure must be provided when using the TEOS-10 equation of ' + 'state.' + ) + specvol = compute_teos10_specvol( + sa=salinity, ct=temperature, p=pressure + ) + else: + raise ValueError(f'Unsupported equation of state type: {eos_type}') + return specvol diff --git a/polaris/ocean/eos/linear.py b/polaris/ocean/eos/linear.py new file mode 100644 index 0000000000..2ac76bd1d6 --- /dev/null +++ b/polaris/ocean/eos/linear.py @@ -0,0 +1,47 @@ +import xarray as xr + +from polaris.config import PolarisConfigParser + + +def compute_linear_density( + config: PolarisConfigParser, + temperature: xr.DataArray, + salinity: xr.DataArray, +) -> xr.DataArray: + """ + Compute the density of seawater based on the the linear equation of state + with coefficients specified in the configuration. The distinction between + conservative, potential, and in-situ temperature or between absolute and + practical salinity is not relevant for the linear EOS. + + Parameters + ---------- + config : polaris.config.PolarisConfigParser + Configuration object containing ocean parameters. + + temperature : float or xarray.DataArray + Temperature of the seawater. + + salinity : float or xarray.DataArray + Salinity of the seawater. + + Returns + ------- + density : float or xarray.DataArray + Computed density of the seawater. + """ + section = config['ocean'] + alpha = section.getfloat('eos_linear_alpha') + beta = section.getfloat('eos_linear_beta') + rhoref = section.getfloat('eos_linear_rhoref') + Tref = section.getfloat('eos_linear_Tref') + Sref = section.getfloat('eos_linear_Sref') + assert ( + alpha is not None + and beta is not None + and rhoref is not None + and Tref is not None + and Sref is not None + ), 'All linear EOS parameters must be specified in the config options.' + density = rhoref + -alpha * (temperature - Tref) + beta * (salinity - Sref) + return density diff --git a/polaris/ocean/eos/teos10.py b/polaris/ocean/eos/teos10.py new file mode 100644 index 0000000000..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', + } + ) diff --git a/polaris/ocean/model/eos.py b/polaris/ocean/model/eos.py deleted file mode 100644 index 7099c1379a..0000000000 --- a/polaris/ocean/model/eos.py +++ /dev/null @@ -1,16 +0,0 @@ -def compute_density(config, temperature, salinity): - eos_type = config.get('ocean', 'eos_type') - if eos_type == 'linear': - density = _compute_linear_density(config, temperature, salinity) - return density - - -def _compute_linear_density(config, temperature, salinity): - section = config['ocean'] - alpha = section.getfloat('eos_linear_alpha') - beta = section.getfloat('eos_linear_beta') - rhoref = section.getfloat('eos_linear_rhoref') - Tref = section.getfloat('eos_linear_Tref') - Sref = section.getfloat('eos_linear_Sref') - density = rhoref + -alpha * (temperature - Tref) + beta * (salinity - Sref) - return density diff --git a/polaris/ocean/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 diff --git a/polaris/ocean/vertical/__init__.py b/polaris/ocean/vertical/__init__.py index 40f422872e..ec137faec3 100644 --- a/polaris/ocean/vertical/__init__.py +++ b/polaris/ocean/vertical/__init__.py @@ -85,7 +85,7 @@ def init_vertical_coord(config, ds): if coord_type == 'z-level': init_z_level_vertical_coord(config, ds) - elif coord_type == 'z-star': + elif coord_type == 'z-star' or coord_type == 'z-tilde': init_z_star_vertical_coord(config, ds) elif coord_type == 'sigma': init_sigma_vertical_coord(config, ds) @@ -110,8 +110,11 @@ def init_vertical_coord(config, ds): dim='Time', axis=0 ) - ds['zMid'] = _compute_zmid_from_layer_thickness( - ds.layerThickness, ds.ssh, ds.cellMask + ds['zInterface'], ds['zMid'] = compute_zint_zmid_from_layer_thickness( + layer_thickness=ds.layerThickness, + bottom_depth=ds.bottomDepth, + min_level_cell=ds.minLevelCell, + max_level_cell=ds.maxLevelCell, ) # fortran 1-based indexing @@ -148,7 +151,7 @@ def update_layer_thickness(config, ds): if coord_type == 'z-level': update_z_level_layer_thickness(config, ds) - elif coord_type == 'z-star': + elif coord_type == 'z-star' or coord_type == 'z-tilde': update_z_star_layer_thickness(config, ds) elif coord_type == 'sigma': update_sigma_layer_thickness(config, ds) @@ -162,47 +165,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 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 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 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'