Function apply_dataarray_vfunc

xarray/computation/apply_ufunc.py:276–338

Apply a variable level function over DataArray, Variable and/or ndarray objects.

(
    func,
    *args,
    signature: _UFuncSignature,
    join: JoinOptions = "inner",
    exclude_dims=frozenset(),
    keep_attrs="override",
)

def apply_dataarray_vfunc(
    func,
    *args,
    signature: _UFuncSignature,
    join: JoinOptions = "inner",
    exclude_dims=frozenset(),
    keep_attrs="override",
) -> tuple[DataArray, ...] | DataArray:
    """Apply a variable level function over DataArray, Variable and/or ndarray
    objects.
    """
    from xarray.core.dataarray import DataArray

    if len(args) > 1:
        args = tuple(
            deep_align(
                args,
                join=join,
                copy=False,
                exclude=exclude_dims,
                raise_on_invalid=False,
            )
        )

    objs = _all_of_type(args, DataArray)

    if keep_attrs == "drop":
        name = result_name(args)
    else:
        first_obj = _first_of_type(args, DataArray)
        name = first_obj.name
    result_coords, result_indexes = build_output_coords_and_indexes(
        args, signature, exclude_dims, combine_attrs=keep_attrs
    )

    data_vars = [getattr(a, "variable", a) for a in args]
    result_var = func(*data_vars)

    out: tuple[DataArray, ...] | DataArray
    if signature.num_outputs > 1:
        out = tuple(
            DataArray(
                variable, coords=coords, indexes=indexes, name=name, fastpath=True
            )
            for variable, coords, indexes in zip(
                result_var, result_coords, result_indexes, strict=True
            )
        )
    else:
        (coords,) = result_coords
        (indexes,) = result_indexes
        out = DataArray(
            result_var, coords=coords, indexes=indexes, name=name, fastpath=True
        )

    attrs = merge_attrs([x.attrs for x in objs], combine_attrs=keep_attrs)
    if isinstance(out, tuple):
        for da in out: