Xarray

convert_ds_dtypes(ds, verbose=False)

Convert all float variables to float32 and all int variables to int32 in an xarray Dataset.

Parameters:

Name	Type	Description	Default
ds	Dataset	Input xarray Dataset.	required
verbose	bool	Whether to print out variables converted	False

Returns:

Name	Type	Description
ds_out	Dataset	Dataset with all float variables converted to float32 and all int variables to int32.

Source code in isca_tools/utils/xarray.py

def convert_ds_dtypes(ds: xr.Dataset, verbose: bool = False) -> xr.Dataset:
    """
    Convert all float variables to float32 and all int variables to int32 in an xarray Dataset.

    Args:
        ds: Input xarray Dataset.
        verbose: Whether to print out variables converted

    Returns:
        ds_out: Dataset with all float variables converted to float32 and all int variables to int32.
    """
    converted = {}
    float_conv = []
    int_conv = []
    for var_name, da in ds.data_vars.items():
        if np.issubdtype(da.dtype, np.floating) and da.dtype != np.float32:
            converted[var_name] = da.astype(np.float32)
            float_conv.append(var_name)
        elif np.issubdtype(da.dtype, np.integer) and da.dtype != np.int32:
            converted[var_name] = da.astype(np.int32)
            int_conv.append(var_name)
        else:
            converted[var_name] = da
    if verbose:
        if len(float_conv) > 0:
            print(f"Converted the following float variables:\n{float_conv}")
        if len(int_conv) > 0:
            print(f"Converted the following integer variables:\n{int_conv}")
    return ds.assign(**converted)

flatten_to_numpy(var, keep_dim=None)

Flattens var to a numpy array with at most 2 dimensions.

Examples:

If var has dims=(lat, lon, lev) and keep_dim=lev, it will return a numpy array of size [n_lat*n_lon, n_lev].

If var has dims=(lat, lon) and keep_dim=None, it will return a numpy array of size [n_lat*n_lon].

Parameters:

Name	Type	Description	Default
var	DataArray	Variable to flatten.	required
keep_dim	Optional[str]	Dimension along which not to flatten.	None

Returns:

Name	Type	Description
var_flatten	ndarray	Numpy array with flattened dimension first, and keep_dim dimension second.

Source code in isca_tools/utils/xarray.py

def flatten_to_numpy(var: xr.DataArray, keep_dim: Optional[str] = None) -> np.ndarray:
    """
    Flattens `var` to a numpy array with at most 2 dimensions.

    Examples:
        If `var` has `dims=(lat, lon, lev)` and `keep_dim=lev`, it will return a numpy array of
            size `[n_lat*n_lon, n_lev]`.

        If `var` has `dims=(lat, lon)` and `keep_dim=None`, it will return a numpy array of
            size `[n_lat*n_lon]`.

    Args:
        var: Variable to flatten.
        keep_dim: Dimension along which not to flatten.

    Returns:
        var_flatten: Numpy array with flattened dimension first, and `keep_dim` dimension second.
    """
    if (keep_dim is not None) and (keep_dim not in var.dims):
        raise ValueError(f"var must have a '{keep_dim}' dimension")

    # dims except vertical
    flatten_dims = [d for d in var.dims if d != keep_dim]

    # stack all flatten_dims into a single "points" axis
    stacked = var.stack(points=flatten_dims)  # dims (..., lev_name) -> (points, lev_name) after transpose
    if keep_dim is not None:
        stacked = stacked.transpose("points", keep_dim)  # ensure order is (points, lev)
    return stacked.values

isel_float(var, ind_float, dim='lev')

Performs da.isel(dim=ind) but for a float index where interpolation is required.

Parameters:

Name	Type	Description	Default
var	DataArray	Variable to find value of at index ind_float along dimension dim.	required
ind_float	DataArray	Fractional index to find value of var along dimension dim.	required
dim		Dimension in which to perform isel.	'lev'

Returns:

Type	Description
DataArray	Value of var at index ind_float along dimension dim.

Source code in isca_tools/utils/xarray.py

def isel_float(var: xr.DataArray, ind_float: xr.DataArray, dim='lev') -> xr.DataArray:
    """
    Performs `da.isel(dim=ind)` but for a float index where interpolation is required.

    Args:
        var: Variable to find value of at index `ind_float` along dimension `dim`.
        ind_float: Fractional index to find value of `var` along dimension `dim`.
        dim: Dimension in which to perform `isel`.

    Returns:
        Value of `var` at index `ind_float` along dimension `dim`.
    """
    # Returns lzb pressure in Pa from the kzlbs model index - requires interpolation to non integer index
    # map fractional index k -> physical lev coordinate
    ind_integer = xr.DataArray(
        np.arange(var[dim].size),
        dims=(dim,),
        coords={dim: var[dim]},
    )

    lev_k = xr.apply_ufunc(
        np.interp,
        ind_float,  # x: fractional index
        ind_integer,  # xp: integer indices 0..n_lev-1
        var[dim],  # fp: lev values
        input_core_dims=[[], [dim], [dim]],
        output_core_dims=[[]],
        vectorize=True,
    )

    # now do the "fractional isel" via interp along lev
    return var.interp(lev=lev_k)

periodic_rolling_mean(da, window, dim='dayofyear', *, center=True, min_periods=None)

Compute a rolling mean with periodic boundary conditions.

Pads the selected dimension using wrapped values before applying the rolling mean, then removes the padding. This allows windows near either boundary to include values from the opposite boundary.

Parameters:

Name	Type	Description	Default
da	Union[DataArray, Dataset]	Input DataArray or Dataset.	required
window	int	Number of points in the rolling window. For a centred rolling mean, an odd value is required.	required
dim	str	Dimension over which to calculate the rolling mean.	'dayofyear'
center	bool	Whether to centre the rolling window on each coordinate.	True
min_periods	Optional[int]	Minimum number of valid values required for each result. Defaults to window.	None

Returns:

Type	Description
Union[DataArray, Dataset]	DataArray or Dataset with the same size and coordinates as da,
Union[DataArray, Dataset]	smoothed along dim.

Raises:

Type	Description
ValueError	If window is not a positive integer.
ValueError	If center=True and window is even.

Source code in isca_tools/utils/xarray.py

def periodic_rolling_mean(
    da: Union[xr.DataArray, xr.Dataset],
    window: int,
    dim: str = "dayofyear",
    *,
    center: bool = True,
    min_periods: Optional[int] = None,
) -> Union[xr.DataArray, xr.Dataset]:
    """Compute a rolling mean with periodic boundary conditions.

    Pads the selected dimension using wrapped values before applying the
    rolling mean, then removes the padding. This allows windows near either
    boundary to include values from the opposite boundary.

    Args:
        da: Input DataArray or Dataset.
        window: Number of points in the rolling window. For a centred rolling
            mean, an odd value is required.
        dim: Dimension over which to calculate the rolling mean.
        center: Whether to centre the rolling window on each coordinate.
        min_periods: Minimum number of valid values required for each result.
            Defaults to `window`.

    Returns:
        DataArray or Dataset with the same size and coordinates as `da`,
        smoothed along `dim`.

    Raises:
        ValueError: If `window` is not a positive integer.
        ValueError: If `center=True` and `window` is even.
    """
    if not isinstance(window, int) or window < 1:
        raise ValueError("window must be a positive integer")

    if center and window % 2 == 0:
        raise ValueError(
            "A centred periodic window should be odd; use e.g. 3, 5, 7, ..."
        )

    if min_periods is None:
        min_periods = window

    # For a centred odd window, pad equally at both ends.
    # For an uncentred window, pad on the left because rolling is right-aligned.
    pad_width = window // 2 if center else window - 1
    pad = {dim: (pad_width, pad_width) if center else (pad_width, 0)}

    padded = da.pad(pad, mode="wrap")
    smoothed = padded.rolling(
        {dim: window},
        center=center,
        min_periods=min_periods,
    ).mean()

    # Restore exactly the original coordinate extent.
    return smoothed.isel({dim: slice(pad_width, pad_width + da.sizes[dim])})

print_ds_var_list(ds, phrase=None)

Prints all variables in ds which contain phrase in the variable name or variable long_name.

Parameters:

Name	Type	Description	Default
ds	Dataset	Dataset to investigate variables of.	required
phrase	Optional[str]	Key phrase to search for in variable info.	None

Source code in isca_tools/utils/xarray.py

def print_ds_var_list(ds: xr.Dataset, phrase: Optional[str] = None) -> None:
    """
    Prints all variables in `ds` which contain `phrase` in the variable name or variable `long_name`.

    Args:
        ds: Dataset to investigate variables of.
        phrase: Key phrase to search for in variable info.

    """
    # All the exceptions to deal with case when var does not have a long_name
    var_list = list(ds.keys())
    if phrase is None:
        for var in var_list:
            try:
                print(f'{var}: {ds[var].long_name}')
            except AttributeError:
                print(f'{var}')
    else:
        for var in var_list:
            if phrase.lower() in var.lower():
                try:
                    print(f'{var}: {ds[var].long_name}')
                except AttributeError:
                    print(f'{var}')
                continue
            try:
                if phrase.lower() in ds[var].long_name.lower():
                    print(f'{var}: {ds[var].long_name}')
                    continue
            except AttributeError:
                continue
    return None

raise_if_common_dims_not_identical(x, y, name_x='x', name_y='y')

Raise an error if the shared dimensions of two xarray objects differ in order.

This compares only the dimensions that appear in both x and y. It extracts the subsequence of x.dims containing only shared dims and does the same for y.dims, then checks the two subsequences are identical (same dim names in the same order).

Parameters:

Name	Type	Description	Default
x		An xarray object (DataArray or Dataset) with a .dims attribute.	required
y		An xarray object (DataArray or Dataset) with a .dims attribute.	required
name_x		Name to use for x in the error message. Defaults to "x".	'x'
name_y		Name to use for y in the error message. Defaults to "y".	'y'

Raises:

Type	Description
ValueError	If the order (or names) of the dimensions shared by x and y are not identical.

Example

If x.dims == ('time', 'lat', 'lon') and y.dims == ('lat', 'time'), the shared dims are ('time', 'lat') for x and ('lat', 'time') for y, so this function raises.

Source code in isca_tools/utils/xarray.py

def raise_if_common_dims_not_identical(x, y, name_x="x", name_y="y"):
    """Raise an error if the shared dimensions of two xarray objects differ in order.

    This compares only the dimensions that appear in both `x` and `y`. It
    extracts the subsequence of `x.dims` containing only shared dims and does
    the same for `y.dims`, then checks the two subsequences are identical (same
    dim names in the same order).

    Args:
        x: An xarray object (DataArray or Dataset) with a `.dims` attribute.
        y: An xarray object (DataArray or Dataset) with a `.dims` attribute.
        name_x: Name to use for `x` in the error message. Defaults to "x".
        name_y: Name to use for `y` in the error message. Defaults to "y".

    Raises:
        ValueError: If the order (or names) of the dimensions shared by `x` and
            `y` are not identical.

    Example:
        If `x.dims == ('time', 'lat', 'lon')` and `y.dims == ('lat', 'time')`,
        the shared dims are `('time', 'lat')` for `x` and `('lat', 'time')` for
        `y`, so this function raises.
    """
    common = [d for d in x.dims if d in y.dims]  # preserves x order
    x_common = tuple(d for d in x.dims if d in common)
    y_common = tuple(d for d in y.dims if d in common)

    if x_common != y_common:
        raise ValueError(
            f"{name_x} and {name_y} common-dim order differs.\n"
            f"{name_x} common dims = {x_common}\n"
            f"{name_y} common dims = {y_common}\n"
            f"{name_x}.dims = {x.dims}\n"
            f"{name_y}.dims = {y.dims}\n"
            f"Try using `transpose_common_dims_like` function to fix"
        )

select_coord_window(ds, target, dim, n=1)

Select a clipped coordinate window around a target value.

Finds the coordinate value along dim nearest to target and selects that grid cell plus up to n neighbouring cells on either side. The selection is clipped at the dimension boundaries, so it may contain fewer than 2 * n + 1 values near either endpoint.

Parameters:

Name	Type	Description	Default
ds	Dataset	Dataset containing the coordinate and dimension.	required
target	float	Target coordinate value in the same units as ds[dim].	required
dim	str	Name of the one-dimensional coordinate and dimension to select along, for example "lat", "lon", or "lev".	required
n	int	Number of neighbouring grid cells to retain on each side of the closest coordinate value. Defaults to 1.	1

Returns:

Type	Description
	xarray.Dataset: Dataset restricted to the selected coordinate window.

Raises:

Type	Description
ValueError	If n is negative, dim is not a dataset dimension, or the coordinate is not one-dimensional.

Source code in isca_tools/utils/xarray.py

def select_coord_window(
    ds: Union[xr.DataArray, xr.Dataset],
    target: float,
    dim: str,
    n: int = 1,
):
    """Select a clipped coordinate window around a target value.

    Finds the coordinate value along ``dim`` nearest to ``target`` and selects
    that grid cell plus up to ``n`` neighbouring cells on either side. The
    selection is clipped at the dimension boundaries, so it may contain fewer
    than ``2 * n + 1`` values near either endpoint.

    Args:
        ds (xarray.Dataset): Dataset containing the coordinate and dimension.
        target (float): Target coordinate value in the same units as
            ``ds[dim]``.
        dim (str): Name of the one-dimensional coordinate and dimension to
            select along, for example ``"lat"``, ``"lon"``, or ``"lev"``.
        n (int, optional): Number of neighbouring grid cells to retain on each
            side of the closest coordinate value. Defaults to 1.

    Returns:
        xarray.Dataset: Dataset restricted to the selected coordinate window.

    Raises:
        ValueError: If ``n`` is negative, ``dim`` is not a dataset dimension,
            or the coordinate is not one-dimensional.
    """
    if n < 0:
        raise ValueError("n must be non-negative")
    if dim not in ds.dims:
        raise ValueError(f"{dim!r} is not a dimension in ds")
    if dim not in ds.coords or ds[dim].ndim != 1:
        raise ValueError(f"{dim!r} must be a one-dimensional coordinate")

    i_dim = abs(ds[dim] - target).argmin().item()
    start = max(0, i_dim - n)
    stop = min(ds.sizes[dim], i_dim + n + 1)

    return ds.isel({dim: slice(start, stop)})

set_attrs(var, overwrite=True, **kwargs)

Set attributes of a given variable.

Examples:

set_attrs(ds.plev, long_name='pressure', units='Pa')

Parameters:

Name	Type	Description	Default
var	DataArray	Variable to set attributes of.	required
overwrite	bool	If True, overwrite existing attributes, otherwise leave unchanged.	True
**kwargs	str	Attributes to set. Common ones include long_name and units	{}

Returns:

Type	Description
DataArray	var with attributes set.

Source code in isca_tools/utils/xarray.py

def set_attrs(var: xr.DataArray, overwrite: bool = True, **kwargs: str) -> xr.DataArray:
    """
    Set attributes of a given variable.

    Examples:
        `set_attrs(ds.plev, long_name='pressure', units='Pa')`

    Args:
        var: Variable to set attributes of.
        overwrite: If `True`, overwrite existing attributes, otherwise leave unchanged.
        **kwargs: Attributes to set. Common ones include `long_name` and `units`

    Returns:
        `var` with attributes set.
    """
    # Function to set main attributes of given variable
    for key in kwargs:
        if (key in var.attrs) and not overwrite:
            continue
        var.attrs[key] = kwargs[key]
    return var

transpose_common_dims_like(y, x)

Reorder y so any dims shared with x follow x's dim order.

Dims that are not in x.dims are left in their original relative order and appended after the reordered common dims.

Parameters:

Name	Type	Description	Default
y	DataArray	Target xarray object (DataArray or Dataset) to reorder.	required
x	DataArray	Reference xarray object (DataArray or Dataset). Only x.dims is used.	required

Returns:

Type	Description
DataArray	y transposed so that its common dims with x match the order in x.

Raises:

Type	Description
ValueError	If y is missing any dim required to complete the transpose order (should only happen if y.dims changes unexpectedly).

Source code in isca_tools/utils/xarray.py

def transpose_common_dims_like(y: xr.DataArray, x: xr.DataArray) -> xr.DataArray:
    """Reorder y so any dims shared with x follow x's dim order.

    Dims that are not in `x.dims` are left in their original relative order and
    appended after the reordered common dims.

    Args:
        y: Target xarray object (DataArray or Dataset) to reorder.
        x: Reference xarray object (DataArray or Dataset). Only `x.dims` is used.

    Returns:
        `y` transposed so that its common dims with `x` match the order in `x`.

    Raises:
        ValueError: If `y` is missing any dim required to complete the transpose
            order (should only happen if `y.dims` changes unexpectedly).
    """
    common_order = [d for d in x.dims if d in y.dims]
    remaining = [d for d in y.dims if d not in common_order]
    order = tuple(common_order + remaining)
    return y.transpose(*order)

unflatten_from_numpy(arr, var, keep_dim=None)

Reconstructs an xarray.DataArray from a flattened NumPy array created by flatten_to_numpy.

Examples:

If var had dims=(lat, lon, lev) and keep_dim='lev', and arr has shape (n_lat*n_lon, n_lev), this will return a DataArray with dims (lat, lon, lev).

If var had dims=(lat, lon)andkeep_dim=None, andarr` has shape (n_lat*n_lon), this will return a DataArray with dims (lat, lon).

Parameters:

Name	Type	Description	Default
arr	ndarray	Flattened NumPy array from flatten_to_numpy.	required
var	DataArray	The original DataArray used to determine dimension order, shape, and coordinates.	required
keep_dim	Optional[str]	Dimension that was kept unflattened in flatten_to_numpy.	None

Returns:

Type	Description
DataArray	xr.DataArray: DataArray with the original dimensions and coordinates restored.

Source code in isca_tools/utils/xarray.py

def unflatten_from_numpy(arr: np.ndarray, var: xr.DataArray, keep_dim: Optional[str] = None) -> xr.DataArray:
    """
    Reconstructs an xarray.DataArray from a flattened NumPy array created by `flatten_to_numpy`.

    Examples:
        If `var` had dims=(lat, lon, lev) and `keep_dim='lev'`, and `arr` has shape (n_lat*n_lon, n_lev),
        this will return a DataArray with dims (lat, lon, lev).

        If `var` had dims=(lat, lon)` and `keep_dim=None`, and `arr` has shape (n_lat*n_lon),
        this will return a DataArray with dims (lat, lon).

    Args:
        arr: Flattened NumPy array from `flatten_to_numpy`.
        var: The original DataArray used to determine dimension order, shape, and coordinates.
        keep_dim: Dimension that was kept unflattened in `flatten_to_numpy`.

    Returns:
        xr.DataArray: DataArray with the original dimensions and coordinates restored.
    """
    # Validate keep_dim
    if (keep_dim is not None) and (keep_dim not in var.dims):
        raise ValueError(f"var must have a '{keep_dim}' dimension")

    # Identify flattened dims
    flatten_dims = [d for d in var.dims if d != keep_dim]

    # Compute target shape
    target_shape = [var.sizes[d] for d in flatten_dims]
    if keep_dim is not None:
        target_shape.append(var.sizes[keep_dim])

    # Reshape numpy array
    reshaped = arr.reshape(target_shape)

    # Reconstruct DataArray with original dimension order
    if keep_dim is not None:
        dims = flatten_dims + [keep_dim]
    else:
        dims = flatten_dims

    # Unstack the flattened dims
    da_flat = xr.DataArray(reshaped, dims=dims, coords={d: var.coords[d] for d in dims if d in var.coords},
                           attrs=var.attrs)

    # Reverse the stacking by unstacking the combined "points" dimension
    if keep_dim is not None:
        da_flat = da_flat.transpose(*var.dims)
    else:
        da_flat = da_flat.transpose(*var.dims)

    return da_flat

update_dim_slice(obj, dim, dim_val, var, var_name=None)

Update values at one coordinate along dim for a Dataset or DataArray.

If obj is a Dataset, updates obj[var_name] at dim=dim_val. If obj is a DataArray, updates obj itself at dim=dim_val and ignores var_name (unless you want to rename the returned DataArray).

Parameters:

Name	Type	Description	Default
obj	Union[DataArray, Dataset]	xarray Dataset or DataArray to update.	required
dim	str	Dimension name to index (e.g. "fit_method").	required
dim_val		Coordinate value along dim to update (e.g. "simulated").	required
var	Union[DataArray, ndarray, float]	New values (DataArray/ndarray/scalar) broadcastable to the target slice.	required
var_name	Optional[str]	Variable name to update if obj is a Dataset. If None, tries to infer from var.name when var is a DataArray.	None

Returns:

Type	Description
	Updated object (same type as input). Note: this mutates data in place.

Raises:

Type	Description
ValueError	If obj is a Dataset and var_name cannot be inferred.
TypeError	If obj is neither Dataset nor DataArray.

Source code in isca_tools/utils/xarray.py

def update_dim_slice(obj: Union[xr.DataArray, xr.Dataset], dim: str, dim_val,
                     var: Union[xr.DataArray, np.ndarray, float], var_name: Optional[str] = None):
    """Update values at one coordinate along `dim` for a Dataset or DataArray.

    If `obj` is a Dataset, updates `obj[var_name]` at `dim=dim_val`.
    If `obj` is a DataArray, updates `obj` itself at `dim=dim_val` and ignores
    `var_name` (unless you want to rename the returned DataArray).

    Args:
        obj: xarray Dataset or DataArray to update.
        dim: Dimension name to index (e.g. "fit_method").
        dim_val: Coordinate value along `dim` to update (e.g. "simulated").
        var: New values (DataArray/ndarray/scalar) broadcastable to the target slice.
        var_name: Variable name to update if `obj` is a Dataset. If None, tries
            to infer from `var.name` when `var` is a DataArray.

    Returns:
        Updated object (same type as input). Note: this mutates data in place.

    Raises:
        ValueError: If `obj` is a Dataset and `var_name` cannot be inferred.
        TypeError: If `obj` is neither Dataset nor DataArray.
    """
    if isinstance(obj, xr.Dataset):
        if var_name is None:
            if isinstance(var, xr.DataArray) and (var.name is not None):
                var_name = var.name
            else:
                raise ValueError("Provide var_name (or pass a DataArray with var.name set).")

        try:
            obj[var_name].loc[{dim: dim_val}] = var
        except ValueError as e:
            raise ValueError(f"Encountered ValueError: {e}\nSolution is probably ds=ds.copy(deep=True) so can update ds.")
        return obj

    if isinstance(obj, xr.DataArray):
        try:
            obj.loc[{dim: dim_val}] = var
        except ValueError as e:
            raise ValueError(f"Encountered ValueError: {e}\nSolution is probably ds=ds.copy(deep=True) so can update ds.")
        return obj

    raise TypeError(f"Expected xarray Dataset or DataArray, got {type(obj)}")

wrap_with_apply_ufunc(func, input_core_dims=None, output_core_dims=None, dask='parallelized', vectorize=True, **ufunc_kwargs)

Wrap a function for use with xarray.apply_ufunc. Usage: func_xr = wrap_with_apply_ufunc(func_numpy, input_core_dims=[[], []]) # only 2 args, rest are kwarg var_xr = func_xr(arg1, arg2, kwarg1=5, kwarg2=3) # where kwarg are key word arguments of the function

Parameters:

Name	Type	Description	Default
func	Callable	The function to wrap.	required
input_core_dims	Optional[Sequence[Sequence[str]]]	Core dimensions for input parameters by the function. Defaults to [[]] for a single scalar-like output.	None
output_core_dims	Optional[Sequence[Sequence[str]]]	Core dimensions produced by the function. Defaults to [[]] for a single scalar-like output.	None
dask	Literal['forbidden', 'allowed', 'parallelized']	Dask handling mode.	'parallelized'
vectorize	bool	Whether to auto-vectorize over non-core dims.	True
**ufunc_kwargs	Any	Extra apply_ufunc keyword arguments.	{}

Returns:

Type	Description
Callable	A callable that behaves like func but operates on xarray objects.

Source code in isca_tools/utils/xarray.py

def wrap_with_apply_ufunc(
    func: Callable,
    input_core_dims: Optional[Sequence[Sequence[str]]] = None,
    output_core_dims: Optional[Sequence[Sequence[str]]] = None,
    dask: Literal["forbidden", "allowed", "parallelized"] = "parallelized",
    vectorize: bool = True,
    **ufunc_kwargs: Any,
) -> Callable:
    """Wrap a function for use with xarray.apply_ufunc.
    Usage:
    `func_xr = wrap_with_apply_ufunc(func_numpy, input_core_dims=[[], []])` # only 2 args, rest are kwarg
    `var_xr = func_xr(arg1, arg2, kwarg1=5, kwarg2=3)` # where kwarg are key word arguments of the function

    Args:
        func: The function to wrap.
        input_core_dims: Core dimensions for input parameters by the function. Defaults
            to ``[[]]`` for a single scalar-like output.
        output_core_dims: Core dimensions produced by the function. Defaults
            to ``[[]]`` for a single scalar-like output.
        dask: Dask handling mode.
        vectorize: Whether to auto-vectorize over non-core dims.
        **ufunc_kwargs: Extra apply_ufunc keyword arguments.

    Returns:
        A callable that behaves like ``func`` but operates on xarray objects.
    """

    def wrapped(*args: Any, **kwargs: Any):
        # Default: one empty core-dim list for each input argument
        if input_core_dims is None:
            input_core = [[] for _ in args]
        else:
            input_core = input_core_dims

        # Default: one empty core-dim list for the function output
        output_core = output_core_dims if output_core_dims is not None else [[]]

        return xr.apply_ufunc(
            func,
            *args,
            input_core_dims=input_core,
            output_core_dims=output_core,
            dask=dask,
            vectorize=vectorize,
            kwargs=kwargs,
            **ufunc_kwargs,
        )

    return wrapped