"""Volume and z coordinate operations on data cubes.
Allows for selecting data subsets using certain volume bounds; selecting
depth or height regions; constructing volumetric averages;
"""
from __future__ import annotations
import logging
from typing import Iterable, Literal, Optional, Sequence
import dask.array as da
import iris
import numpy as np
from iris.coords import AuxCoord, CellMeasure
from iris.cube import Cube
from iris.util import broadcast_to_shape
from esmvalcore.iris_helpers import ignore_iris_vague_metadata_warnings
from esmvalcore.preprocessor._shared import (
get_coord_weights,
get_iris_aggregator,
get_normalized_cube,
preserve_float_dtype,
try_adding_calculated_cell_area,
update_weights_kwargs,
)
from esmvalcore.preprocessor._supplementary_vars import (
register_supplementaries,
)
logger = logging.getLogger(__name__)
def calculate_volume(cube: Cube) -> np.ndarray | da.Array:
"""Calculate volume from a cube.
This function is used when the 'ocean_volume' cell measure can't be found.
The output data will be given in cubic meters (m3).
Note
----
It gets the cell_area from the cube if it is available. If not, it
calculates it from the grid. This only works if the grid cell areas can
be calculated (i.e., latitude and longitude are 1D). The depth coordinate
units should be convertible to meters.
Parameters
----------
cube:
input cube.
Returns
-------
np.ndarray | dask.array.Array
Grid volume.
"""
# Load depth field and figure out which dim is which
depth = cube.coord(axis="z")
z_dim = cube.coord_dims(depth)
depth = depth.copy()
# Assert z has length > 0
if not z_dim:
raise ValueError("Cannot compute volume with scalar Z-axis")
# Guess bounds if missing
if not depth.has_bounds():
depth.guess_bounds()
if depth.core_bounds().shape[-1] != 2:
raise ValueError(
f"Z axis bounds shape found {depth.core_bounds().shape}. "
"Bounds should be 2 in the last dimension to compute the "
"thickness."
)
# Convert units to get the thickness in meters
try:
depth.convert_units("m")
except ValueError as err:
raise ValueError(
f"Cannot compute volume using the Z-axis. {err}"
) from err
# Calculate Z-direction thickness
thickness = depth.core_bounds()[..., 1] - depth.core_bounds()[..., 0]
if cube.has_lazy_data():
z_chunks = tuple(cube.lazy_data().chunks[d] for d in z_dim)
if isinstance(thickness, da.Array):
thickness = thickness.rechunk(z_chunks)
else:
thickness = da.asarray(thickness, chunks=z_chunks)
# Get or calculate the horizontal areas of the cube
has_cell_measure = bool(cube.cell_measures("cell_area"))
try_adding_calculated_cell_area(cube)
area = cube.cell_measure("cell_area").copy()
area_dim = cube.cell_measure_dims(area)
area.convert_units("m2")
area_array = area.core_data()
if cube.has_lazy_data():
area_array = da.array(area_array)
# Make sure input cube has not been modified
if not has_cell_measure:
cube.remove_cell_measure("cell_area")
chunks = cube.core_data().chunks if cube.has_lazy_data() else None
area_arr = broadcast_to_shape(
area_array, cube.shape, area_dim, chunks=chunks
)
thickness_arr = broadcast_to_shape(
thickness, cube.shape, z_dim, chunks=chunks
)
grid_volume = area_arr * thickness_arr
if cube.has_lazy_data():
grid_volume = grid_volume.rechunk(chunks)
return grid_volume
def _try_adding_calculated_ocean_volume(cube: Cube) -> None:
"""Try to add calculated cell measure 'ocean_volume' to cube (in-place)."""
if cube.cell_measures("ocean_volume"):
return
logger.debug(
"Found no cell measure 'ocean_volume' in cube %s. Check availability "
"of supplementary variables",
cube.summary(shorten=True),
)
logger.debug("Attempting to calculate grid cell volume")
grid_volume = calculate_volume(cube)
cell_measure = CellMeasure(
grid_volume,
standard_name="ocean_volume",
units="m3",
measure="volume",
)
cube.add_cell_measure(cell_measure, np.arange(cube.ndim))
[docs]
@register_supplementaries(
variables=["volcello", "areacello"],
required="prefer_at_least_one",
)
@preserve_float_dtype
def volume_statistics(
cube: Cube,
operator: str,
normalize: Optional[Literal["subtract", "divide"]] = None,
**operator_kwargs,
) -> Cube:
"""Apply a statistical operation over a volume.
The volume average is weighted according to the cell volume.
Parameters
----------
cube:
Input cube. The input cube should have a
:class:`iris.coords.CellMeasure` named ``'ocean_volume'``, unless it
has a :class:`iris.coords.CellMeasure` named ``'cell_area'`` or
regular 1D latitude and longitude coordinates so the cell areas
can be computed using :func:`iris.analysis.cartography.area_weights`.
The volume will be computed from the area multiplied by the
thickness, computed from the bounds of the vertical coordinate.
In that case, vertical coordinate units should be convertible to
meters.
operator:
The operation. Used to determine the :class:`iris.analysis.Aggregator`
object used to calculate the statistics. Currently, only `mean` is
allowed.
normalize:
If given, do not return the statistics cube itself, but rather, the
input cube, normalized with the statistics cube. Can either be
`subtract` (statistics cube is subtracted from the input cube) or
`divide` (input cube is divided by the statistics cube).
**operator_kwargs:
Optional keyword arguments for the :class:`iris.analysis.Aggregator`
object defined by `operator`.
Note
----
This preprocessor has been designed for oceanic variables, but it might
be applicable to atmospheric data as well.
Returns
-------
iris.cube.Cube
Collapsed cube.
"""
has_cell_measure = bool(cube.cell_measures("ocean_volume"))
# TODO: Test sigma coordinates.
# TODO: Add other operations.
if operator != "mean":
raise ValueError(f"Volume operator {operator} not recognised.")
# get z, y, x coords
z_axis = cube.coord(axis="Z")
y_axis = cube.coord(axis="Y")
x_axis = cube.coord(axis="X")
# assert z axis only uses 1 dimension more than x, y axis
xy_dims = tuple({*cube.coord_dims(y_axis), *cube.coord_dims(x_axis)})
xyz_dims = tuple({*cube.coord_dims(z_axis), *xy_dims})
if len(xyz_dims) > len(xy_dims) + 1:
raise ValueError(
f"X and Y axis coordinates depend on {xy_dims} dimensions, "
f"while X, Y, and Z axis depends on {xyz_dims} dimensions. "
"This may indicate Z axis depending on other dimension than "
"space that could provoke invalid aggregation..."
)
(agg, agg_kwargs) = get_iris_aggregator(operator, **operator_kwargs)
agg_kwargs = update_weights_kwargs(
operator,
agg,
agg_kwargs,
"ocean_volume",
cube,
_try_adding_calculated_ocean_volume,
)
with ignore_iris_vague_metadata_warnings():
result = cube.collapsed([z_axis, y_axis, x_axis], agg, **agg_kwargs)
if normalize is not None:
result = get_normalized_cube(cube, result, normalize)
# Make sure input cube has not been modified
if not has_cell_measure and cube.cell_measures("ocean_volume"):
cube.remove_cell_measure("ocean_volume")
return result
[docs]
@preserve_float_dtype
def axis_statistics(
cube: Cube,
axis: str,
operator: str,
normalize: Optional[Literal["subtract", "divide"]] = None,
**operator_kwargs,
) -> Cube:
"""Perform statistics along a given axis.
Operates over an axis direction.
Note
----
The `mean`, `sum` and `rms` operations are weighted by the corresponding
coordinate bounds by default. For `sum`, the units of the resulting cube
will be multiplied by corresponding coordinate units.
Arguments
---------
cube:
Input cube.
axis:
Direction over where to apply the operator. Possible values are `x`,
`y`, `z`, `t`.
operator:
The operation. Used to determine the :class:`iris.analysis.Aggregator`
object used to calculate the statistics. Allowed options are given in
:ref:`this table <supported_stat_operator>`.
normalize:
If given, do not return the statistics cube itself, but rather, the
input cube, normalized with the statistics cube. Can either be
`subtract` (statistics cube is subtracted from the input cube) or
`divide` (input cube is divided by the statistics cube).
**operator_kwargs:
Optional keyword arguments for the :class:`iris.analysis.Aggregator`
object defined by `operator`.
Returns
-------
iris.cube.Cube
Collapsed cube.
"""
(agg, agg_kwargs) = get_iris_aggregator(operator, **operator_kwargs)
# Check if a coordinate for the desired axis exists
try:
coord = cube.coord(axis=axis)
except iris.exceptions.CoordinateNotFoundError as err:
raise ValueError(
f"Axis {axis} not found in cube {cube.summary(shorten=True)}"
) from err
# Multidimensional coordinates are currently not supported
coord_dims = cube.coord_dims(coord)
if len(coord_dims) > 1:
raise NotImplementedError(
"axis_statistics not implemented for multidimensional coordinates."
)
# For weighted operations, create a dummy weights coordinate using the
# bounds of the original coordinate (this handles units properly, e.g., for
# sums)
agg_kwargs = update_weights_kwargs(
operator,
agg,
agg_kwargs,
"_axis_statistics_weights_",
cube,
_add_axis_stats_weights_coord,
coord=coord,
)
with ignore_iris_vague_metadata_warnings():
result = cube.collapsed(coord, agg, **agg_kwargs)
if normalize is not None:
result = get_normalized_cube(cube, result, normalize)
# Make sure input and output cubes do not have auxiliary coordinate
if cube.coords("_axis_statistics_weights_"):
cube.remove_coord("_axis_statistics_weights_")
if result.coords("_axis_statistics_weights_"):
result.remove_coord("_axis_statistics_weights_")
return result
def _add_axis_stats_weights_coord(cube, coord):
"""Add weights for axis_statistics to cube (in-place)."""
weights = get_coord_weights(cube, coord)
weights_coord = AuxCoord(
weights,
long_name="_axis_statistics_weights_",
units=coord.units,
)
coord_dims = cube.coord_dims(coord)
cube.add_aux_coord(weights_coord, coord_dims)
[docs]
def depth_integration(cube: Cube) -> Cube:
"""Determine the total sum over the vertical component.
Requires a 3D cube. The z-coordinate integration is calculated by taking
the sum in the z direction of the cell contents multiplied by the cell
thickness. The units of the resulting cube are multiplied by the
z-coordinate units.
Arguments
---------
cube:
Input cube.
Returns
-------
iris.cube.Cube
Collapsed cube.
"""
result = axis_statistics(cube, axis="z", operator="sum")
result.rename("Depth_integrated_" + str(cube.name()))
return result
