Source code for esmvalcore.preprocessor._supplementary_vars

"""Preprocessor functions for ancillary variables and cell measures."""

import logging
from collections.abc import Callable
from typing import Iterable, Literal

import iris.coords
from iris.cube import Cube

logger = logging.getLogger(__name__)

PREPROCESSOR_SUPPLEMENTARIES = {}


def register_supplementaries(
    variables: list[str],
    required: Literal["require_at_least_one", "prefer_at_least_one"],
) -> Callable:
    """Register supplementary variables required for a preprocessor function.

    Parameters
    ----------
    variables:
        List of variable names.
    required:
        How strong the requirement is. Can be 'require_at_least_one' if at
        least one variable must be available or 'prefer_at_least_one' if it is
        preferred that at least one variable is available, but not strictly
        necessary.
    """
    valid = ("require_at_least_one", "prefer_at_least_one")
    if required not in valid:
        raise NotImplementedError(f"`required` should be one of {valid}")
    supplementaries = {
        "variables": variables,
        "required": required,
    }

    def wrapper(func):
        PREPROCESSOR_SUPPLEMENTARIES[func.__name__] = supplementaries
        return func

    return wrapper


def add_cell_measure(
    cube: Cube,
    cell_measure_cube: Cube,
    measure: Literal["area", "volume"],
) -> None:
    """Add cell measure to cube (in-place).

    Note
    ----
    This assumes that the cell measure spans the rightmost dimensions of the
    cube.

    Parameters
    ----------
    cube:
        Iris cube with input data.
    cell_measure_cube:
        Iris cube with cell measure data.
    measure:
        Name of the measure, can be 'area' or 'volume'.

    Returns
    -------
    iris.cube.Cube
        Cube with added cell measure

    Raises
    ------
    ValueError
        If measure name is not 'area' or 'volume'.
    """
    if measure not in ["area", "volume"]:
        raise ValueError(
            f"measure name must be 'area' or 'volume', got {measure} instead"
        )
    coord_dims = tuple(
        range(cube.ndim - len(cell_measure_cube.shape), cube.ndim)
    )
    cell_measure_data = cell_measure_cube.core_data()
    if cell_measure_cube.has_lazy_data():
        cube_chunks = tuple(cube.lazy_data().chunks[d] for d in coord_dims)
        cell_measure_data = cell_measure_data.rechunk(cube_chunks)
    cell_measure = iris.coords.CellMeasure(
        cell_measure_data,
        standard_name=cell_measure_cube.standard_name,
        units=cell_measure_cube.units,
        measure=measure,
        var_name=cell_measure_cube.var_name,
        attributes=cell_measure_cube.attributes,
    )
    cube.add_cell_measure(cell_measure, coord_dims)
    logger.debug(
        "Added %s as cell measure in cube of %s.",
        cell_measure_cube.var_name,
        cube.var_name,
    )


def add_ancillary_variable(cube: Cube, ancillary_cube: Cube) -> None:
    """Add ancillary variable to cube (in-place).

    Note
    ----
    This assumes that the ancillary variable spans the rightmost dimensions of
    the cube.

    Parameters
    ----------
    cube:
        Iris cube with input data.
    ancillary_cube:
        Iris cube with ancillary data.

    Returns
    -------
    iris.cube.Cube
        Cube with added ancillary variables
    """
    coord_dims = tuple(range(cube.ndim - len(ancillary_cube.shape), cube.ndim))
    ancillary_data = ancillary_cube.core_data()
    if ancillary_cube.has_lazy_data():
        cube_chunks = tuple(cube.lazy_data().chunks[d] for d in coord_dims)
        ancillary_data = ancillary_data.rechunk(cube_chunks)
    ancillary_var = iris.coords.AncillaryVariable(
        ancillary_data,
        standard_name=ancillary_cube.standard_name,
        units=ancillary_cube.units,
        var_name=ancillary_cube.var_name,
        attributes=ancillary_cube.attributes,
    )
    cube.add_ancillary_variable(ancillary_var, coord_dims)
    logger.debug(
        "Added %s as ancillary variable in cube of %s.",
        ancillary_cube.var_name,
        cube.var_name,
    )


[docs] def add_supplementary_variables( cube: Cube, supplementary_cubes: Iterable[Cube], ) -> Cube: """Add ancillary variables and/or cell measures to cube (in-place). Parameters ---------- cube: Cube to add to. supplementary_cubes: Iterable of cubes containing the supplementary variables. Returns ------- iris.cube.Cube Cube with added ancillary variables and/or cell measures. """ measure_names: dict[str, Literal["area", "volume"]] = { "areacella": "area", "areacello": "area", "volcello": "volume", } for supplementary_cube in supplementary_cubes: if supplementary_cube.var_name in measure_names: measure_name = measure_names[supplementary_cube.var_name] add_cell_measure(cube, supplementary_cube, measure_name) else: add_ancillary_variable(cube, supplementary_cube) return cube
[docs] def remove_supplementary_variables(cube: Cube) -> Cube: """Remove supplementary variables from cube (in-place). Strip cell measures or ancillary variables from the cube. Parameters ---------- cube: Iris cube with data and cell measures or ancillary variables. Returns ------- iris.cube.Cube Cube without cell measures or ancillary variables. """ if cube.cell_measures(): for measure in cube.cell_measures(): cube.remove_cell_measure(measure) if cube.ancillary_variables(): for variable in cube.ancillary_variables(): cube.remove_ancillary_variable(variable) return cube