eye

Return a 2-D array in the specified format with ones on the diagonal and zeros elsewhere.

Parameters:

Name	Type	Description	Default
N	int	Number of rows in the output.	required
M	int	Number of columns in the output. If None, defaults to N.	None
k	int	Index of the diagonal: 0 (the default) refers to the main diagonal, a positive value refers to an upper diagonal, and a negative value to a lower diagonal.	0
dtype	data - type	Data-type of the returned array.	float
format	str	A format string.	'coo'

Returns:

Name	Type	Description
I	SparseArray of shape (N, M)	An array where all elements are equal to zero, except for the k-th diagonal, whose values are equal to one.

Examples:

>>> eye(2, dtype=int).todense()
array([[1, 0],
       [0, 1]])
>>> eye(3, k=1).todense()
array([[0., 1., 0.],
       [0., 0., 1.],
       [0., 0., 0.]])

Source code in sparse/numba_backend/_common.py

@_check_device
def eye(N, M=None, k=0, dtype=float, format="coo", *, device=None, **kwargs):
    """Return a 2-D array in the specified format with ones on the diagonal and zeros elsewhere.

    Parameters
    ----------
    N : int
        Number of rows in the output.
    M : int, optional
        Number of columns in the output. If None, defaults to `N`.
    k : int, optional
        Index of the diagonal: 0 (the default) refers to the main diagonal,
        a positive value refers to an upper diagonal, and a negative value
        to a lower diagonal.
    dtype : data-type, optional
        Data-type of the returned array.
    format : str, optional
        A format string.

    Returns
    -------
    I : SparseArray of shape (N, M)
        An array where all elements are equal to zero, except for the `k`-th
        diagonal, whose values are equal to one.

    Examples
    --------
    >>> eye(2, dtype=int).todense()  # doctest: +NORMALIZE_WHITESPACE
    array([[1, 0],
           [0, 1]])
    >>> eye(3, k=1).todense()  # doctest: +SKIP
    array([[0., 1., 0.],
           [0., 0., 1.],
           [0., 0., 0.]])
    """
    from ._coo import COO

    if M is None:
        M = N

    N = int(N)
    M = int(M)
    k = int(k)

    data_length = builtins.min(N, M)
    if k > 0:
        data_length = builtins.max(builtins.min(data_length, M - k), 0)
    elif k < 0:
        data_length = builtins.max(builtins.min(data_length, N + k), 0)

    if data_length == 0:
        return zeros((N, M), dtype=dtype, format=format, device=device)

    if k > 0:
        n_coords = np.arange(data_length, dtype=np.intp)
        m_coords = n_coords + k
    elif k < 0:
        m_coords = np.arange(data_length, dtype=np.intp)
        n_coords = m_coords - k
    else:
        n_coords = m_coords = np.arange(data_length, dtype=np.intp)

    coords = np.stack([n_coords, m_coords])
    data = np.array(1, dtype=dtype)

    return COO(coords, data=data, shape=(N, M), has_duplicates=False, sorted=True).asformat(format, **kwargs)