some new features

.venv/lib/python3.12/site-packages/scipy/sparse/_extract.py

@@ -0,0 +1,178 @@
+"""Functions to extract parts of sparse matrices
+"""
+
+__docformat__ = "restructuredtext en"
+
+__all__ = ['find', 'tril', 'triu']
+
+
+from ._coo import coo_matrix, coo_array
+from ._base import sparray
+
+
+def find(A):
+    """Return the indices and values of the nonzero elements of a matrix
+
+    Parameters
+    ----------
+    A : dense or sparse array or matrix
+        Matrix whose nonzero elements are desired.
+
+    Returns
+    -------
+    (I,J,V) : tuple of arrays
+        I,J, and V contain the row indices, column indices, and values
+        of the nonzero entries.
+
+
+    Examples
+    --------
+    >>> from scipy.sparse import csr_array, find
+    >>> A = csr_array([[7.0, 8.0, 0],[0, 0, 9.0]])
+    >>> find(A)
+    (array([0, 0, 1], dtype=int32),
+     array([0, 1, 2], dtype=int32),
+     array([ 7.,  8.,  9.]))
+
+    """
+
+    A = coo_array(A, copy=True)
+    A.sum_duplicates()
+    # remove explicit zeros
+    nz_mask = A.data != 0
+    return A.row[nz_mask], A.col[nz_mask], A.data[nz_mask]
+
+
+def tril(A, k=0, format=None):
+    """Return the lower triangular portion of a sparse array or matrix
+
+    Returns the elements on or below the k-th diagonal of A.
+        - k = 0 corresponds to the main diagonal
+        - k > 0 is above the main diagonal
+        - k < 0 is below the main diagonal
+
+    Parameters
+    ----------
+    A : dense or sparse array or matrix
+        Matrix whose lower trianglar portion is desired.
+    k : integer : optional
+        The top-most diagonal of the lower triangle.
+    format : string
+        Sparse format of the result, e.g. format="csr", etc.
+
+    Returns
+    -------
+    L : sparse matrix
+        Lower triangular portion of A in sparse format.
+
+    See Also
+    --------
+    triu : upper triangle in sparse format
+
+    Examples
+    --------
+    >>> from scipy.sparse import csr_array, tril
+    >>> A = csr_array([[1, 2, 0, 0, 3], [4, 5, 0, 6, 7], [0, 0, 8, 9, 0]],
+    ...               dtype='int32')
+    >>> A.toarray()
+    array([[1, 2, 0, 0, 3],
+           [4, 5, 0, 6, 7],
+           [0, 0, 8, 9, 0]])
+    >>> tril(A).toarray()
+    array([[1, 0, 0, 0, 0],
+           [4, 5, 0, 0, 0],
+           [0, 0, 8, 0, 0]])
+    >>> tril(A).nnz
+    4
+    >>> tril(A, k=1).toarray()
+    array([[1, 2, 0, 0, 0],
+           [4, 5, 0, 0, 0],
+           [0, 0, 8, 9, 0]])
+    >>> tril(A, k=-1).toarray()
+    array([[0, 0, 0, 0, 0],
+           [4, 0, 0, 0, 0],
+           [0, 0, 0, 0, 0]])
+    >>> tril(A, format='csc')
+    <Compressed Sparse Column sparse array of dtype 'int32'
+        with 4 stored elements and shape (3, 5)>
+
+    """
+    coo_sparse = coo_array if isinstance(A, sparray) else coo_matrix
+
+    # convert to COOrdinate format where things are easy
+    A = coo_sparse(A, copy=False)
+    mask = A.row + k >= A.col
+
+    row = A.row[mask]
+    col = A.col[mask]
+    data = A.data[mask]
+    new_coo = coo_sparse((data, (row, col)), shape=A.shape, dtype=A.dtype)
+    return new_coo.asformat(format)
+
+
+def triu(A, k=0, format=None):
+    """Return the upper triangular portion of a sparse array or matrix
+
+    Returns the elements on or above the k-th diagonal of A.
+        - k = 0 corresponds to the main diagonal
+        - k > 0 is above the main diagonal
+        - k < 0 is below the main diagonal
+
+    Parameters
+    ----------
+    A : dense or sparse array or matrix
+        Matrix whose upper trianglar portion is desired.
+    k : integer : optional
+        The bottom-most diagonal of the upper triangle.
+    format : string
+        Sparse format of the result, e.g. format="csr", etc.
+
+    Returns
+    -------
+    L : sparse array or matrix 
+        Upper triangular portion of A in sparse format.
+        Sparse array if A is a sparse array, otherwise matrix.
+
+    See Also
+    --------
+    tril : lower triangle in sparse format
+
+    Examples
+    --------
+    >>> from scipy.sparse import csr_array, triu
+    >>> A = csr_array([[1, 2, 0, 0, 3], [4, 5, 0, 6, 7], [0, 0, 8, 9, 0]],
+    ...                dtype='int32')
+    >>> A.toarray()
+    array([[1, 2, 0, 0, 3],
+           [4, 5, 0, 6, 7],
+           [0, 0, 8, 9, 0]])
+    >>> triu(A).toarray()
+    array([[1, 2, 0, 0, 3],
+           [0, 5, 0, 6, 7],
+           [0, 0, 8, 9, 0]])
+    >>> triu(A).nnz
+    8
+    >>> triu(A, k=1).toarray()
+    array([[0, 2, 0, 0, 3],
+           [0, 0, 0, 6, 7],
+           [0, 0, 0, 9, 0]])
+    >>> triu(A, k=-1).toarray()
+    array([[1, 2, 0, 0, 3],
+           [4, 5, 0, 6, 7],
+           [0, 0, 8, 9, 0]])
+    >>> triu(A, format='csc')
+    <Compressed Sparse Column sparse array of dtype 'int32'
+        with 8 stored elements and shape (3, 5)>
+
+    """
+    coo_sparse = coo_array if isinstance(A, sparray) else coo_matrix
+
+    # convert to COOrdinate format where things are easy
+    A = coo_sparse(A, copy=False)
+    mask = A.row + k <= A.col
+
+    row = A.row[mask]
+    col = A.col[mask]
+    data = A.data[mask]
+    new_coo = coo_sparse((data, (row, col)), shape=A.shape, dtype=A.dtype)
+    return new_coo.asformat(format)