expand_dims

Expands the shape of an array by inserting a new axis (dimension) of size one at the position specified by axis.

Parameters:

Name	Type	Description	Default
x	COO	Input COO array.	required
axis	int | tuple[int, ...]	Position in the expanded axes where the new axis is placed.	0

Returns:

Name	Type	Description
result	COO	An expanded output COO array having the same data type as x.

Examples:

>>> import sparse
>>> x = sparse.COO.from_numpy([[1, 0, 0, 0, 2, -3]])
>>> x.shape
(1, 6)
>>> y1 = sparse.expand_dims(x, axis=1)
>>> y1.shape
(1, 1, 6)
>>> y2 = sparse.expand_dims(x, axis=2)
>>> y2.shape
(1, 6, 1)

Source code in sparse/numba_backend/_coo/common.py

def expand_dims(x, axis: int | tuple[int, ...] = 0):
    """
    Expands the shape of an array by inserting a new axis (dimension) of size
    one at the position specified by ``axis``.

    Parameters
    ----------
    x : COO
        Input COO array.
    axis : int | tuple[int, ...]
        Position in the expanded axes where the new axis is placed.

    Returns
    -------
    result : COO
        An expanded output COO array having the same data type as ``x``.

    Examples
    --------
    >>> import sparse
    >>> x = sparse.COO.from_numpy([[1, 0, 0, 0, 2, -3]])
    >>> x.shape
    (1, 6)
    >>> y1 = sparse.expand_dims(x, axis=1)
    >>> y1.shape
    (1, 1, 6)
    >>> y2 = sparse.expand_dims(x, axis=2)
    >>> y2.shape
    (1, 6, 1)

    """
    from .core import COO

    if isinstance(x, np.generic | np.ndarray):
        x = COO.from_numpy(x)

    x = _validate_coo_input(x)

    if isinstance(axis, int):
        axis = (axis,)

    axis = normalize_axis(axis, x.ndim + len(axis))

    iter_curr_coords = iter(x.coords)
    iter_curr_shape = iter(x.shape)
    new_coords = []
    new_shape = []

    for d in range(x.ndim + len(axis)):
        s = next(iter_curr_shape) if d not in axis else 1
        c = next(iter_curr_coords) if d not in axis else np.zeros_like(x.data, dtype=x.coords.dtype)
        new_shape.append(s)
        new_coords.append(c)

    return COO(
        np.stack(new_coords, axis=0) if len(new_coords) > 0 else np.empty((0, x.coords.shape[1]), dtype=x.coords.dtype),
        x.data,
        shape=new_shape,
        fill_value=x.fill_value,
    )