Function apply_variable_ufunc

xarray/computation/apply_ufunc.py:707–870

Apply an ndarray level function over Variable and/or ndarray objects.

(
    func,
    *args,
    signature: _UFuncSignature,
    exclude_dims=frozenset(),
    dask="forbidden",
    output_dtypes=None,
    vectorize=False,
    keep_attrs="override",
    dask_gufunc_kwargs=None,
)

def apply_variable_ufunc(
    func,
    *args,
    signature: _UFuncSignature,
    exclude_dims=frozenset(),
    dask="forbidden",
    output_dtypes=None,
    vectorize=False,
    keep_attrs="override",
    dask_gufunc_kwargs=None,
) -> Variable | tuple[Variable, ...]:
    """Apply an ndarray level function over Variable and/or ndarray objects."""
    from xarray.core.formatting import short_array_repr
    from xarray.core.variable import as_compatible_data

    dim_sizes = unified_dim_sizes(
        (a for a in args if hasattr(a, "dims")), exclude_dims=exclude_dims
    )
    broadcast_dims = tuple(
        dim for dim in dim_sizes if dim not in signature.all_core_dims
    )
    output_dims = [broadcast_dims + out for out in signature.output_core_dims]

    input_data = [
        (
            broadcast_compat_data(arg, broadcast_dims, core_dims)
            if isinstance(arg, Variable)
            else arg
        )
        for arg, core_dims in zip(args, signature.input_core_dims, strict=True)
    ]

    if any(is_chunked_array(array) for array in input_data):
        if dask == "forbidden":
            raise ValueError(
                "apply_ufunc encountered a chunked array on an "
                "argument, but handling for chunked arrays has not "
                "been enabled. Either set the ``dask`` argument "
                "or load your data into memory first with "
                "``.load()`` or ``.compute()``"
            )
        elif dask == "parallelized":
            chunkmanager = get_chunked_array_type(*input_data)

            numpy_func = func

            if dask_gufunc_kwargs is None:
                dask_gufunc_kwargs = {}
            else:
                dask_gufunc_kwargs = dask_gufunc_kwargs.copy()

            allow_rechunk = dask_gufunc_kwargs.get("allow_rechunk", None)
            if allow_rechunk is None:
                for n, (data, core_dims) in enumerate(
                    zip(input_data, signature.input_core_dims, strict=True)
                ):
                    if is_chunked_array(data):
                        # core dimensions cannot span multiple chunks