Function apply_dict_of_variables_vfunc

xarray/computation/apply_ufunc.py:433–467

Apply a variable level function over dicts of DataArray, DataArray, Variable and ndarray objects.

(
    func,
    *args,
    signature: _UFuncSignature,
    join="inner",
    fill_value=None,
    on_missing_core_dim: MissingCoreDimOptions = "raise",
)

def apply_dict_of_variables_vfunc(
    func,
    *args,
    signature: _UFuncSignature,
    join="inner",
    fill_value=None,
    on_missing_core_dim: MissingCoreDimOptions = "raise",
):
    """Apply a variable level function over dicts of DataArray, DataArray,
    Variable and ndarray objects.
    """
    args = tuple(_as_variables_or_variable(arg) for arg in args)
    names = join_dict_keys(args, how=join)
    grouped_by_name = collect_dict_values(args, names, fill_value)

    result_vars = {}
    for name, variable_args in zip(names, grouped_by_name, strict=True):
        core_dim_present = _check_core_dims(signature, variable_args, name)
        if core_dim_present is True:
            result_vars[name] = func(*variable_args)
        elif on_missing_core_dim == "raise":
            raise ValueError(core_dim_present)
        elif on_missing_core_dim == "copy":
            result_vars[name] = variable_args[0]
        elif on_missing_core_dim == "drop":
            pass
        else:
            raise ValueError(
                f"Invalid value for `on_missing_core_dim`: {on_missing_core_dim!r}"
            )

    if signature.num_outputs > 1:
        return _unpack_dict_tuples(result_vars, signature.num_outputs)
    else:
        return result_vars


def _fast_dataset(