Function apply_dataset_vfunc

xarray/computation/apply_ufunc.py:486–556

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

(
    func,
    *args,
    signature: _UFuncSignature,
    join="inner",
    dataset_join="exact",
    fill_value=_NO_FILL_VALUE,
    exclude_dims=frozenset(),
    keep_attrs="override",
    on_missing_core_dim: MissingCoreDimOptions = "raise",
)

def apply_dataset_vfunc(
    func,
    *args,
    signature: _UFuncSignature,
    join="inner",
    dataset_join="exact",
    fill_value=_NO_FILL_VALUE,
    exclude_dims=frozenset(),
    keep_attrs="override",
    on_missing_core_dim: MissingCoreDimOptions = "raise",
) -> Dataset | tuple[Dataset, ...]:
    """Apply a variable level function over Dataset, dict of DataArray,
    DataArray, Variable and/or ndarray objects.
    """
    from xarray.core.dataset import Dataset

    if dataset_join not in _JOINS_WITHOUT_FILL_VALUES and fill_value is _NO_FILL_VALUE:
        raise TypeError(
            "to apply an operation to datasets with different "
            "data variables with apply_ufunc, you must supply the "
            "dataset_fill_value argument."
        )

    objs = _all_of_type(args, Dataset)

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

    list_of_coords, list_of_indexes = build_output_coords_and_indexes(
        args, signature, exclude_dims, combine_attrs=keep_attrs
    )
    args = tuple(getattr(arg, "data_vars", arg) for arg in args)

    result_vars = apply_dict_of_variables_vfunc(
        func,
        *args,
        signature=signature,
        join=dataset_join,
        fill_value=fill_value,
        on_missing_core_dim=on_missing_core_dim,
    )

    out: Dataset | tuple[Dataset, ...]
    if signature.num_outputs > 1:
        out = tuple(
            itertools.starmap(
                _fast_dataset,
                zip(result_vars, list_of_coords, list_of_indexes, strict=True),
            )
        )