Function vectorized_indexers

xarray/testing/strategies.py:638–733

Generate vectorized (fancy) indexers where all arrays are broadcastable. In vectorized indexing, all array indexers must have compatible shapes that can be broadcast together, and the result shape is determined by broadcasting the indexer arrays. Parameters ---------- draw

(
    draw,
    /,
    *,
    sizes: dict[Hashable, int],
    min_dims: int = 2,
    max_dims: int | None = None,
    min_ndim: int = 1,
    max_ndim: int = 3,
    min_size: int = 1,
    max_size: int = 5,
)

@st.composite
def vectorized_indexers(
    draw,
    /,
    *,
    sizes: dict[Hashable, int],
    min_dims: int = 2,
    max_dims: int | None = None,
    min_ndim: int = 1,
    max_ndim: int = 3,
    min_size: int = 1,
    max_size: int = 5,
) -> dict[Hashable, xr.DataArray]:
    """Generate vectorized (fancy) indexers where all arrays are broadcastable.

    In vectorized indexing, all array indexers must have compatible shapes
    that can be broadcast together, and the result shape is determined by
    broadcasting the indexer arrays.

    Parameters
    ----------
    draw : callable
        The Hypothesis draw function (automatically provided by @st.composite).
    sizes : dict[Hashable, int]
        Dictionary mapping dimension names to their sizes.
    min_dims : int, optional
        Minimum number of dimensions to index. Default is 2, so that we always have a "trajectory".
        Use ``outer_array_indexers`` for the ``min_dims==1`` case.
    max_dims : int or None, optional
        Maximum number of dimensions to index.
    min_ndim : int, optional
        Minimum number of dimensions for the result arrays.
    max_ndim : int, optional
        Maximum number of dimensions for the result arrays.
    min_size : int, optional
        Minimum size for each dimension in the result arrays.
    max_size : int, optional
        Maximum size for each dimension in the result arrays.

    Returns
    -------
    sizes : mapping of hashable to DataArray or Variable
        Indexers as a dict with keys randomly selected from sizes.keys().
        Values are DataArrays of integer indices that are all broadcastable
        to a common shape.

    See Also
    --------
    hypothesis.extra.numpy.arrays
    """
    selected_dims = draw(unique_subset_of(sizes, min_size=min_dims, max_size=max_dims))

    # Generate a common broadcast shape for all arrays
    # Use min_ndim to max_ndim dimensions for the result shape
    result_shape = draw(
        st.lists(
            st.integers(min_value=min_size, max_value=max_size),
            min_size=min_ndim,
            max_size=max_ndim,