Add a "depth_average" function

[mx-moth]

This can either be an example, or added to the depth operations, as appropriate.

There are uses for computing an average value of a variable across a depth variable, weighting the values with the size of each depth layer. This is useful for tracking a tracer variable dispersing across a 2D domain, but the variable is 3D or 4D.

import numpy
import xarray
from emsarray.types import DataArrayOrName
from emsarray.utils import name_to_data_array


def depth_average(
    dataset: xarray.Dataset,
    data_array: DataArrayOrName,
    *,
    layer_sizes: xarray.DataArray | None = None
) -> xarray.DataArray:
    data_array = name_to_data_array(dataset, data_array)
    depth_variable = dataset.ems.get_depth_coordinate_for_data_array(data_array)
    depth_dimension = depth_variable.dims[0]

    if layer_sizes is None:
        if 'bounds' in depth_variable.attrs:
            depth_bounds = dataset[depth_variable.attrs['bounds']]
            assert depth_bounds.shape == (depth_variable.shape[0], 2)
            layer_sizes = bounds_to_size(depth_bounds)
        else:
            raise ValueError(
                f"Depth variable {depth_variable.name!r} does not have a bounds attribute. "
                "The layer_sizes parameter is required if the depth variable has no bounds.")

    total_depth = (numpy.isfinite(data_array) * layer_sizes).sum(depth_dimension)

    weighted_sum = (data_array * layer_sizes).sum(depth_dimension)
    weighted_average = weighted_sum / total_depth

    return weighted_average


def bounds_to_size(bounds):
    return bounds[:, 0] - bounds[:, 1]

This is complicated by many depth variable lacking proper bounds attributes. Some examples of how to estimate the bounds for this should be included. One untested option:

def estimate_bounds_from_centres(
    centres,
    *,
    clamp_lower: float | None = None,
    clamp_upper: float | None = None,
):
    if clamp_lower is None:
        clamp_lower = centres[0] - (centres[1] - centres[0])
    if clamp_upper is None:
        clamp_upper = centres[-1] + (centres[-1] - centres[-2])
    inner_bounds = centres[1:] - centres[:-1]

    return numpy.c_([
        [clamp_lower] + inner_bounds,
        inner_bounds + [clamp_upper],
    ])

For COMPAS all files, the Mesh2_layerfaces variable hold some of the useful information:

# Compute the thickness of each layer.
layerfaces = ds['Mesh2_layerfaces'].values
layerfaces[-1] = 0  # This is usually very high in the air and spoils the weighting
layer_sizes = xarray.DataArray(layerfaces[1:] - layerfaces[:-1], dims=('Mesh2_layers',))

Example use:

import emsarray
import xarray
from depth_average import depth_average

ds = emsarray.open_dataset("/home/hea211/example-datasets/setas_compas_all_2018-05.nc")
temp = ds['temp']

# Compute the thickness of each layer.
layerfaces = ds['Mesh2_layerfaces'].values
layerfaces[-1] = 0  # This is usually very high in the air and spoils the weighting
layer_sizes = xarray.DataArray(layerfaces[1:] - layerfaces[:-1], dims=('Mesh2_layers',))

# Find the depth averaged temperature
avg_temp = depth_average(ds, temp, layer_sizes=layer_sizes)