diagonal

Extract diagonal from a COO array. The equivalent of numpy.diagonal.

Parameters:

Name	Type	Description	Default
a	COO	The array to perform the operation on.	required
offset	int	Offset of the diagonal from the main diagonal. Defaults to main diagonal (0).	0
axis1	int	First axis from which the diagonals should be taken. Defaults to first axis (0).	0
axis2	int	Second axis from which the diagonals should be taken. Defaults to second axis (1).	1

Examples:

>>> import sparse
>>> x = sparse.as_coo(np.arange(9).reshape(3, 3))
>>> sparse.diagonal(x).todense()
array([0, 4, 8])
>>> sparse.diagonal(x, offset=1).todense()
array([1, 5])

>>> x = sparse.as_coo(np.arange(12).reshape((2, 3, 2)))
>>> x_diag = sparse.diagonal(x, axis1=0, axis2=2)
>>> x_diag.shape
(3, 2)
>>> x_diag.todense()
array([[ 0,  7],
       [ 2,  9],
       [ 4, 11]])

Returns:

Name	Type	Description
out	COO	The result of the operation.

Raises:

Type	Description
ValueError	If a.shape[axis1] != a.shape[axis2]

See Also

numpy.diagonal : NumPy equivalent function

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

def diagonal(a, offset=0, axis1=0, axis2=1):
    """
    Extract diagonal from a COO array. The equivalent of [`numpy.diagonal`][].

    Parameters
    ----------
    a : COO
        The array to perform the operation on.
    offset : int, optional
        Offset of the diagonal from the main diagonal. Defaults to main diagonal (0).
    axis1 : int, optional
        First axis from which the diagonals should be taken.
        Defaults to first axis (0).
    axis2 : int, optional
        Second axis from which the diagonals should be taken.
        Defaults to second axis (1).

    Examples
    --------
    >>> import sparse
    >>> x = sparse.as_coo(np.arange(9).reshape(3, 3))
    >>> sparse.diagonal(x).todense()
    array([0, 4, 8])
    >>> sparse.diagonal(x, offset=1).todense()
    array([1, 5])

    >>> x = sparse.as_coo(np.arange(12).reshape((2, 3, 2)))
    >>> x_diag = sparse.diagonal(x, axis1=0, axis2=2)
    >>> x_diag.shape
    (3, 2)
    >>> x_diag.todense()
    array([[ 0,  7],
           [ 2,  9],
           [ 4, 11]])

    Returns
    -------
    out: COO
        The result of the operation.

    Raises
    ------
    ValueError
        If a.shape[axis1] != a.shape[axis2]

    See Also
    --------
    [`numpy.diagonal`][] : NumPy equivalent function
    """
    from .core import COO

    if a.shape[axis1] != a.shape[axis2]:
        raise ValueError("a.shape[axis1] != a.shape[axis2]")

    diag_axes = [axis for axis in range(len(a.shape)) if axis != axis1 and axis != axis2] + [axis1]
    diag_shape = [a.shape[axis] for axis in diag_axes]
    diag_shape[-1] -= abs(offset)

    diag_idx = _diagonal_idx(a.coords, axis1, axis2, offset)

    diag_coords = [a.coords[axis][diag_idx] for axis in diag_axes]
    diag_data = a.data[diag_idx]

    return COO(diag_coords, diag_data, diag_shape)