(A, blocksize, inv_flag=True)
return block_diag
Return the block diagonal of A, in array form.

    Parameters
    ----------
    A : csr_matrix
        assumed to be square
    blocksize : int
        square block size for the diagonal
    inv_flag : bool
        if True, return the inverse of the block diagonal

    Returns
    -------
    block_diag : array
        block diagonal of A in array form,
        array size is (A.shape[0]/blocksize, blocksize, blocksize)

    Examples
    --------
    >>> from numpy import arange
    >>> from scipy.sparse import csr_matrix
    >>> from pyamg.util.utils import get_block_diag
    >>> A = csr_matrix(arange(36).reshape(6,6))
    >>> block_diag_inv = get_block_diag(A, blocksize=2, inv_flag=False)
    >>> print(block_diag_inv)
    [[[ 0.  1.]
      [ 6.  7.]]
    <BLANKLINE>
     [[14. 15.]
      [20. 21.]]
    <BLANKLINE>
     [[28. 29.]
      [34. 35.]]]
    >>> block_diag_inv = get_block_diag(A, blocksize=2, inv_flag=True)
Return the block diagonal of A, in array form.
Return the block diagonal of A in array form .
def get_block_diag(A, blocksize, inv_flag=True):
    """Return the block diagonal of A, in array form.

    Parameters
    ----------
    A : csr_matrix
        assumed to be square
    blocksize : int
        square block size for the diagonal
    inv_flag : bool
        if True, return the inverse of the block diagonal

    Returns
    -------
    block_diag : array
        block diagonal of A in array form,
        array size is (A.shape[0]/blocksize, blocksize, blocksize)

    Examples
    --------
    >>> from numpy import arange
    >>> from scipy.sparse import csr_matrix
    >>> from pyamg.util.utils import get_block_diag
    >>> A = csr_matrix(arange(36).reshape(6,6))
    >>> block_diag_inv = get_block_diag(A, blocksize=2, inv_flag=False)
    >>> print(block_diag_inv)
    [[[ 0.  1.]
      [ 6.  7.]]
    <BLANKLINE>
     [[14. 15.]
      [20. 21.]]
    <BLANKLINE>
     [[28. 29.]
      [34. 35.]]]
    >>> block_diag_inv = get_block_diag(A, blocksize=2, inv_flag=True)

    """
    if not isspmatrix(A):
        raise TypeError('Expected sparse matrix')
    if A.shape[0] != A.shape[1]:
        raise ValueError('Expected square matrix')
    if np.mod(A.shape[0], blocksize) != 0:
        raise ValueError('blocksize and A.shape must be compatible')

    # If the block diagonal of A already exists, return that
    if hasattr(A, 'block_D_inv') and inv_flag:
        if (A.block_D_inv.shape[1] == blocksize) and\
           (A.block_D_inv.shape[2] == blocksize) and \
           (A.block_D_inv.shape[0] == int(A.shape[0]/blocksize)):
            return A.block_D_inv
    elif hasattr(A, 'block_D') and (not inv_flag):
        if (A.block_D.shape[1] == blocksize) and\
           (A.block_D.shape[2] == blocksize) and \
           (A.block_D.shape[0] == int(A.shape[0]/blocksize)):
            return A.block_D

    # Convert to BSR
    if not isspmatrix_bsr(A):
        A = bsr_matrix(A, blocksize=(blocksize, blocksize))
    if A.blocksize != (blocksize, blocksize):
        A = A.tobsr(blocksize=(blocksize, blocksize))

    # Peel off block diagonal by extracting block entries from the now BSR
    # matrix A
    A = A.asfptype()
    block_diag = np.zeros((int(A.shape[0]/blocksize), blocksize, blocksize),
                          dtype=A.dtype)

    AAIJ = (np.arange(1, A.indices.shape[0]+1), A.indices, A.indptr)
    shape = (int(A.shape[0]/blocksize), int(A.shape[0]/blocksize))
    diag_entries = csr_matrix(AAIJ, shape=shape).diagonal()
    diag_entries -= 1
    nonzero_mask = (diag_entries != -1)
    diag_entries = diag_entries[nonzero_mask]
    if diag_entries.shape != (0,):
        block_diag[nonzero_mask, :, :] = A.data[diag_entries, :, :]

    if inv_flag:
        # Invert each block
        if block_diag.shape[1] < 7:
            # This specialized routine lacks robustness for large matrices
            amg_core.pinv_array(block_diag.ravel(), block_diag.shape[0],
                                block_diag.shape[1], 'T')
        else:
            linalg.pinv_array(block_diag)
        A.block_D_inv = block_diag
    else:
        A.block_D = block_diag

    return block_diag
def get_block_diag ( A , blocksize , inv_flag = True ) : if not isspmatrix ( A ) : raise TypeError ( 'Expected sparse matrix' ) if A . shape [ 0 ] != A . shape [ 1 ] : raise ValueError ( 'Expected square matrix' ) if np . mod ( A . shape [ 0 ] , blocksize ) != 0 : raise ValueError ( 'blocksize and A.shape must be compatible' ) # If the block diagonal of A already exists, return that if hasattr ( A , 'block_D_inv' ) and inv_flag : if ( A . block_D_inv . shape [ 1 ] == blocksize ) and ( A . block_D_inv . shape [ 2 ] == blocksize ) and ( A . block_D_inv . shape [ 0 ] == int ( A . shape [ 0 ] / blocksize ) ) : return A . block_D_inv elif hasattr ( A , 'block_D' ) and ( not inv_flag ) : if ( A . block_D . shape [ 1 ] == blocksize ) and ( A . block_D . shape [ 2 ] == blocksize ) and ( A . block_D . shape [ 0 ] == int ( A . shape [ 0 ] / blocksize ) ) : return A . block_D # Convert to BSR if not isspmatrix_bsr ( A ) : A = bsr_matrix ( A , blocksize = ( blocksize , blocksize ) ) if A . blocksize != ( blocksize , blocksize ) : A = A . tobsr ( blocksize = ( blocksize , blocksize ) ) # Peel off block diagonal by extracting block entries from the now BSR # matrix A A = A . asfptype ( ) block_diag = np . zeros ( ( int ( A . shape [ 0 ] / blocksize ) , blocksize , blocksize ) , dtype = A . dtype ) AAIJ = ( np . arange ( 1 , A . indices . shape [ 0 ] + 1 ) , A . indices , A . indptr ) shape = ( int ( A . shape [ 0 ] / blocksize ) , int ( A . shape [ 0 ] / blocksize ) ) diag_entries = csr_matrix ( AAIJ , shape = shape ) . diagonal ( ) diag_entries -= 1 nonzero_mask = ( diag_entries != - 1 ) diag_entries = diag_entries [ nonzero_mask ] if diag_entries . shape != ( 0 , ) : block_diag [ nonzero_mask , : , : ] = A . data [ diag_entries , : , : ] if inv_flag : # Invert each block if block_diag . shape [ 1 ] < 7 : # This specialized routine lacks robustness for large matrices amg_core . pinv_array ( block_diag . ravel ( ) , block_diag . shape [ 0 ] , block_diag . shape [ 1 ] , 'T' ) else : linalg . pinv_array ( block_diag ) A . block_D_inv = block_diag else : A . block_D = block_diag return block_diag