def get_diagonal(A, norm_eq=False, inv=False):
    """Return the diagonal or inverse of diagonal for A, (A.H A) or (A A.H).

    Parameters
    ----------
    A   : {dense or sparse matrix}
        e.g. array, matrix, csr_matrix, ...
    norm_eq : {0, 1, 2}
        0 ==> D = diag(A)
        1 ==> D = diag(A.H A)
        2 ==> D = diag(A A.H)
    inv : {True, False}
        If True, D = 1.0/D

    Returns
    -------
    diagonal, D, of appropriate system

    Notes
    -----
    This function is especially useful for its fast methods
    of obtaining diag(A A.H) and diag(A.H A).  Dinv is zero
    wherever D is zero

    Examples
    --------
    >>> from pyamg.util.utils import get_diagonal
    >>> from pyamg.gallery import poisson
    >>> A = poisson( (5,), format='csr' )
    >>> D = get_diagonal(A)
    >>> print D
    [ 2.  2.  2.  2.  2.]
    >>> D = get_diagonal(A, norm_eq=1, inv=True)
    >>> print D
    [ 0.2         0.16666667  0.16666667  0.16666667  0.2       ]

    """
    # if not isspmatrix(A):
    if not (isspmatrix_csr(A) or isspmatrix_csc(A) or isspmatrix_bsr(A)):
        warn('Implicit conversion to sparse matrix')
        A = csr_matrix(A)

    # critical to sort the indices of A
    A.sort_indices()
    if norm_eq == 1:
        # This transpose involves almost no work, use csr data structures as
        # csc, or vice versa
        At = A.T
        D = (At.multiply(At.conjugate()))*np.ones((At.shape[0],))
    elif norm_eq == 2:
        D = (A.multiply(A.conjugate()))*np.ones((A.shape[0],))
    else:
        D = A.diagonal()

    if inv:
        Dinv = np.zeros_like(D)
        mask = (D != 0.0)
        Dinv[mask] = 1.0 / D[mask]
        return Dinv
    else:
        return D
Return the diagonal or inverse of diagonal for A, (A.H A) or (A A.H).

    Parameters
    ----------
    A   : {dense or sparse matrix}
        e.g. array, matrix, csr_matrix, ...
    norm_eq : {0, 1, 2}
        0 ==> D = diag(A)
        1 ==> D = diag(A.H A)
        2 ==> D = diag(A A.H)
    inv : {True, False}
        If True, D = 1.0/D

    Returns
    -------
    diagonal, D, of appropriate system

    Notes
    -----
    This function is especially useful for its fast methods
    of obtaining diag(A A.H) and diag(A.H A).  Dinv is zero
    wherever D is zero

    Examples
    --------
    >>> from pyamg.util.utils import get_diagonal
    >>> from pyamg.gallery import poisson
    >>> A = poisson( (5,), format='csr' )
    >>> D = get_diagonal(A)
    >>> print D
    [ 2.  2.  2.  2.  2.]
    >>> D = get_diagonal(A, norm_eq=1, inv=True)
    >>> print D
    [ 0.2         0.16666667  0.16666667  0.16666667  0.2       ]