对角线

从COO数组中提取对角线。等同于numpy.diagonal。

参数

名称	类型	描述	默认值
a	COO	要进行操作的数组。	必需
offset	int	对角线相对于主对角线的偏移量。默认为主对角线 (0)。	0
axis1	int	应提取对角线的第一个轴。默认为第一个轴 (0)。	0
axis2	int	应提取对角线的第二个轴。默认为第二个轴 (1)。	1

示例

>>> 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]])

返回值

名称	类型	描述
out	COO	操作的结果。

引发

类型	描述
ValueError	If a.shape[axis1] != a.shape[axis2]

另请参阅

numpy.diagonal : NumPy 等效函数

源代码位于 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)