Function broadcast_compat_data

xarray/computation/apply_ufunc.py:648–691

(
    variable: Variable,
    broadcast_dims: tuple[Hashable, ...],
    core_dims: tuple[Hashable, ...],
)

def broadcast_compat_data(
    variable: Variable,
    broadcast_dims: tuple[Hashable, ...],
    core_dims: tuple[Hashable, ...],
) -> Any:
    data = variable.data

    old_dims = variable.dims
    new_dims = broadcast_dims + core_dims

    if new_dims == old_dims:
        # optimize for the typical case
        return data

    set_old_dims = set(old_dims)
    set_new_dims = set(new_dims)
    unexpected_dims = [d for d in old_dims if d not in set_new_dims]

    if unexpected_dims:
        raise ValueError(
            "operand to apply_ufunc encountered unexpected "
            f"dimensions {unexpected_dims!r} on an input variable: these are core "
            "dimensions on other input or output variables"
        )

    # for consistency with numpy, keep broadcast dimensions to the left
    old_broadcast_dims = tuple(d for d in broadcast_dims if d in set_old_dims)
    reordered_dims = old_broadcast_dims + core_dims
    if reordered_dims != old_dims:
        order = tuple(old_dims.index(d) for d in reordered_dims)
        data = duck_array_ops.transpose(data, order)

    if new_dims != reordered_dims:
        key_parts: list[slice | None] = []
        for dim in new_dims:
            if dim in set_old_dims:
                key_parts.append(SLICE_NONE)
            elif key_parts:
                # no need to insert new axes at the beginning that are already
                # handled by broadcasting
                key_parts.append(np.newaxis)
        data = data[tuple(key_parts)]

    return data


def _vectorize(func, signature, output_dtypes, exclude_dims):