class MultilevelSolver:
    """Stores multigrid hierarchy and implements the multigrid cycle.

    The class constructs the cycling process and points to the methods for
    coarse grid solves.  A MultilevelSolver object is typically returned from a
    particular AMG method (see ruge_stuben_solver or smoothed_aggregation_solver
    for example).  A call to MultilevelSolver.solve() is a typical access
    point.  The class also defines methods for constructing operator, cycle, and
    grid complexities.

    Attributes
    ----------
    levels : level array
        Array of level objects that contain A, R, and P.
    coarse_solver : string
        String passed to coarse_grid_solver indicating the solve type

    Methods
    -------
    aspreconditioner()
        Create a preconditioner using this multigrid cycle
    cycle_complexity()
        A measure of the cost of a single multigrid cycle.
    grid_complexity()
        A measure of the rate of coarsening.
    operator_complexity()
        A measure of the size of the multigrid hierarchy.
    solve()
        Iteratively solves a linear system for the right hand side.

    """

    class Level:
        """Stores one level of the multigrid hierarchy.

        All level objects will have an 'A' attribute referencing the matrix
        of that level.  All levels, except for the coarsest level, will
        also have 'P' and 'R' attributes referencing the prolongation and
        restriction operators that act between each level and the next
        coarser level.

        Attributes
        ----------
        A : csr_matrix
            Problem matrix for Ax=b
        R : csr_matrix
            Restriction matrix between levels (often R = P.T)
        P : csr_matrix
            Prolongation or Interpolation matrix.

        Notes
        -----
        The functionality of this class is a struct

        """

        def __init__(self):
            """Level construct (empty)."""
            self.A = None

    class level(Level):  # noqa: N801
        """Deprecated level class."""

        def __init__(self):
            """Raise deprecation warning on use, not import."""
            super().__init__()
            warn('level() is deprectated.  use Level()',
                 category=DeprecationWarning, stacklevel=2)

    def __init__(self, levels, coarse_solver='pinv'):
        """Class constructor to initialize the cycle and ensure list of levels is complete.

        Parameters
        ----------
        levels : level array
            Array of level objects that contain A, R, and P.
        coarse_solver: string, callable, tuple
            The solver method is either (1) a string such as 'splu' or 'pinv'
            of a callable object which receives only parameters (A, b) and
            returns an (approximate or exact) solution to the linear system Ax
            = b, or (2) a callable object that takes parameters (A,b) and
            returns an (approximate or exact) solution to Ax = b, or (3) a
            tuple of the form (string|callable, args), where args is a
            dictionary of arguments to be passed to the function denoted by
            string or callable.

            Sparse direct methods:

            * splu         : sparse LU solver

            Sparse iterative methods:

            * any method in scipy.sparse.linalg or pyamg.krylov (e.g. 'cg').
            * Methods in pyamg.krylov take precedence.
            * relaxation method, such as 'gauss_seidel' or 'jacobi',

            Dense methods:

            * pinv     : pseudoinverse (SVD)
            * lu       : LU factorization
            * cholesky : Cholesky factorization

        Notes
        -----
        If not defined, the R attribute on each level is set to
        the transpose of P.

        Examples
        --------
        >>> # manual construction of a two-level AMG hierarchy
        >>> from pyamg.gallery import poisson
        >>> from pyamg.multilevel import MultilevelSolver
        >>> from pyamg.strength import classical_strength_of_connection
        >>> from pyamg.classical.interpolate import direct_interpolation
        >>> from pyamg.classical.split import RS
        >>> # compute necessary operators
        >>> A = poisson((100, 100), format='csr')
        >>> C = classical_strength_of_connection(A)
        >>> splitting = RS(A)
        >>> P = direct_interpolation(A, C, splitting)
        >>> R = P.T
        >>> # store first level data
        >>> levels = []
        >>> levels.append(MultilevelSolver.Level())
        >>> levels.append(MultilevelSolver.Level())
        >>> levels[0].A = A
        >>> levels[0].C = C
        >>> levels[0].splitting = splitting
        >>> levels[0].P = P
        >>> levels[0].R = R
        >>> # store second level data
        >>> levels[1].A = R @ A @ P                      # coarse-level matrix
        >>> # create MultilevelSolver
        >>> ml = MultilevelSolver(levels, coarse_solver='splu')
        >>> print(ml)
        MultilevelSolver
        Number of Levels:     2
        Operator Complexity:  1.891
        Grid Complexity:      1.500
        Coarse Solver:        'splu'
          level   unknowns     nonzeros
             0       10000        49600 [52.88%]
             1        5000        44202 [47.12%]
        <BLANKLINE>

        """
        self.symmetric_smoothing = False  # force change_smoothers to set to True
        self.levels = levels

        self.coarse_solver = coarse_grid_solver(coarse_solver)

        for level in levels[:-1]:
            if not hasattr(level, 'R'):
                level.R = level.P.H

    def __repr__(self):
        """Print basic statistics about the multigrid hierarchy."""
        output = 'MultilevelSolver\n'
        output += f'Number of Levels:     {len(self.levels)}\n'
        output += f'Operator Complexity: {self.operator_complexity():6.3f}\n'
        output += f'Grid Complexity:     {self.grid_complexity():6.3f}\n'
        output += f'Coarse Solver:        {self.coarse_solver.name()}\n'

        total_nnz = sum(level.A.nnz for level in self.levels)

        #          123456712345678901 123456789012 123456789
        #               0       10000        49600 [52.88%]
        output += '  level   unknowns     nonzeros\n'
        for n, level in enumerate(self.levels):
            A = level.A
            ratio = 100 * A.nnz / total_nnz
            output += f'{n:>6} {A.shape[1]:>11} {A.nnz:>12} [{ratio:2.2f}%]\n'

        return output

    def cycle_complexity(self, cycle='V'):
        """Cycle complexity of V, W, AMLI, and F(1,1) cycle with simple relaxation.

        Cycle complexity is an approximate measure of the number of
        floating point operations (FLOPs) required to perform a single
        multigrid cycle relative to the cost a single smoothing operation.

        Parameters
        ----------
        cycle : {'V','W','F','AMLI'}
            Type of multigrid cycle to perform in each iteration.

        Returns
        -------
        cc : float
            Defined as F_sum / F_0, where
            F_sum is the total number of nonzeros in the matrix on all
            levels encountered during a cycle and F_0 is the number of
            nonzeros in the matrix on the finest level.

        Notes
        -----
        This is only a rough estimate of the true cycle complexity. The
        estimate assumes that the cost of pre and post-smoothing are
        (each) equal to the number of nonzeros in the matrix on that level.
        This assumption holds for smoothers like Jacobi and Gauss-Seidel.
        However, the true cycle complexity of cycle using more expensive
        methods, like block Gauss-Seidel will be underestimated.

        Additionally, if the cycle used in practice isn't a (1,1)-cycle,
        then this cost estimate will be off.

        """
        cycle = str(cycle).upper()

        nnz = [level.A.nnz for level in self.levels]

        def V(level):
            if len(self.levels) == 1:
                return nnz[0]

            if level == len(self.levels) - 2:
                return 2 * nnz[level] + nnz[level + 1]

            return 2 * nnz[level] + V(level + 1)

        def W(level):
            if len(self.levels) == 1:
                return nnz[0]

            if level == len(self.levels) - 2:
                return 2 * nnz[level] + nnz[level + 1]

            return 2 * nnz[level] + 2 * W(level + 1)

        def F(level):
            if len(self.levels) == 1:
                return nnz[0]

            if level == len(self.levels) - 2:
                return 2 * nnz[level] + nnz[level + 1]

            return 2 * nnz[level] + F(level + 1) + V(level + 1)

        if cycle == 'V':
            flops = V(0)
        elif cycle in ('W', 'AMLI'):
            flops = W(0)
        elif cycle == 'F':
            flops = F(0)
        else:
            raise TypeError(f'Unrecognized cycle type ({cycle})')

        return float(flops) / float(nnz[0])

    def operator_complexity(self):
        """Operator complexity of this multigrid hierarchy.

        Defined as:
            Number of nonzeros in the matrix on all levels /
            Number of nonzeros in the matrix on the finest level

        """
        return sum(level.A.nnz for level in self.levels) /\
            float(self.levels[0].A.nnz)

    def grid_complexity(self):
        """Grid complexity of this multigrid hierarchy.

        Defined as:
            Number of unknowns on all levels /
            Number of unknowns on the finest level

        """
        return sum(level.A.shape[0] for level in self.levels) /\
            float(self.levels[0].A.shape[0])

    def psolve(self, b):
        """Lagacy solve interface."""
        return self.solve(b, maxiter=1)

    def aspreconditioner(self, cycle='V'):
        """Create a preconditioner using this multigrid cycle.

        Parameters
        ----------
        cycle : {'V','W','F','AMLI'}
            Type of multigrid cycle to perform in each iteration.

        Returns
        -------
        precond : LinearOperator
            Preconditioner suitable for the iterative solvers in defined in
            the scipy.sparse.linalg module (e.g. cg, gmres) and any other
            solver that uses the LinearOperator interface.  Refer to the
            LinearOperator documentation in scipy.sparse.linalg

        See Also
        --------
        MultilevelSolver.solve, scipy.sparse.linalg.LinearOperator

        Examples
        --------
        >>> from pyamg.aggregation import smoothed_aggregation_solver
        >>> from pyamg.gallery import poisson
        >>> from scipy.sparse.linalg import cg
        >>> import scipy as sp
        >>> A = poisson((100, 100), format='csr')          # matrix
        >>> b = np.random.rand(A.shape[0])                 # random RHS
        >>> ml = smoothed_aggregation_solver(A)            # AMG solver
        >>> M = ml.aspreconditioner(cycle='V')             # preconditioner
        >>> x, info = cg(A, b, tol=1e-8, maxiter=30, M=M)  # solve with CG

        """
        shape = self.levels[0].A.shape
        dtype = self.levels[0].A.dtype

        def matvec(b):
            return self.solve(b, maxiter=1, cycle=cycle, tol=1e-12)

        return LinearOperator(shape, matvec, dtype=dtype)

    def solve(self, b, x0=None, tol=1e-5, maxiter=100, cycle='V', accel=None,
              callback=None, residuals=None, return_info=False):
        """Execute multigrid cycling.

        Parameters
        ----------
        b : array
            Right hand side.
        x0 : array
            Initial guess.
        tol : float
            Stopping criteria: relative residual r[k]/||b|| tolerance.
            If `accel` is used, the stopping criteria is set by the Krylov method.
        maxiter : int
            Stopping criteria: maximum number of allowable iterations.
        cycle : {'V','W','F','AMLI'}
            Type of multigrid cycle to perform in each iteration.
        accel : string, function
            Defines acceleration method.  Can be a string such as 'cg'
            or 'gmres' which is the name of an iterative solver in
            pyamg.krylov (preferred) or scipy.sparse.linalg.
            If accel is not a string, it will be treated like a function
            with the same interface provided by the iterative solvers in SciPy.
        callback : function
            User-defined function called after each iteration.  It is
            called as callback(xk) where xk is the k-th iterate vector.
        residuals : list
            List to contain residual norms at each iteration.  The residuals
            will be the residuals from the Krylov iteration -- see the `accel`
            method to see verify whether this ||r|| or ||Mr|| (as in the case of
            GMRES).
        return_info : bool
            If true, will return (x, info)
            If false, will return x (default)

        Returns
        -------
        x : array
            Approximate solution to Ax=b after k iterations

        info : string
            Halting status

            ==  =======================================
            0   successful exit
            >0  convergence to tolerance not achieved,
                return iteration count instead.
            ==  =======================================

        See Also
        --------
        aspreconditioner

        Examples
        --------
        >>> from numpy import ones
        >>> from pyamg import ruge_stuben_solver
        >>> from pyamg.gallery import poisson
        >>> A = poisson((100, 100), format='csr')
        >>> b = A * ones(A.shape[0])
        >>> ml = ruge_stuben_solver(A, max_coarse=10)
        >>> residuals = []
        >>> x = ml.solve(b, tol=1e-12, residuals=residuals) # standalone solver

        """
        if x0 is None:
            x = np.zeros_like(b)
        else:
            x = np.array(x0)  # copy

        A = self.levels[0].A

        cycle = str(cycle).upper()

        # AMLI cycles require hermitian matrix
        if (cycle == 'AMLI') and hasattr(A, 'symmetry'):
            if A.symmetry != 'hermitian':
                raise ValueError('AMLI cycles require \
                    symmetry to be hermitian')

        if accel is not None:

            # Check for symmetric smoothing scheme when using CG
            if (accel == 'cg') and (not self.symmetric_smoothing):
                warn('Incompatible non-symmetric multigrid preconditioner '
                     'detected, due to presmoother/postsmoother combination. '
                     'CG requires SPD preconditioner, not just SPD matrix.')

            # Check for AMLI compatability
            if (accel != 'fgmres') and (cycle == 'AMLI'):
                raise ValueError('AMLI cycles require acceleration (accel) '
                                 'to be fgmres, or no acceleration')

            # Acceleration is being used
            kwargs = {}
            if isinstance(accel, str):
                kwargs = {}
                if hasattr(krylov, accel):
                    accel = getattr(krylov, accel)
                else:
                    accel = getattr(sla, accel)
                    kwargs['atol'] = 'legacy'

            M = self.aspreconditioner(cycle=cycle)

            try:  # try PyAMG style interface which has a residuals parameter
                x, info = accel(A, b, x0=x0, tol=tol, maxiter=maxiter, M=M,
                                callback=callback, residuals=residuals, **kwargs)
                if return_info:
                    return x, info
                return x
            except TypeError:
                # try the scipy.sparse.linalg style interface,
                # which requires a callback function if a residual
                # history is desired

                if residuals is not None:
                    residuals[:] = [np.linalg.norm(b - A @ x)]

                    def callback_wrapper(x):
                        if np.isscalar(x):
                            residuals.append(x)
                        else:
                            residuals.append(np.linalg.norm(b - A @ x))
                        if callback is not None:
                            callback(x)
                else:
                    callback_wrapper = callback

                x, info = accel(A, b, x0=x0, tol=tol, maxiter=maxiter, M=M,
                                callback=callback_wrapper, **kwargs)
                if return_info:
                    return x, info
                return x

        else:
            # Scale tol by normb
            # Don't scale tol earlier. The accel routine should also scale tol
            normb = np.linalg.norm(b)
            if normb == 0.0:
                normb = 1.0  # set so that we have an absolute tolerance

        # Start cycling (no acceleration)
        normr = np.linalg.norm(b - A @ x)
        if residuals is not None:
            residuals[:] = [normr]  # initial residual

        # Create uniform types for A, x and b
        # Clearly, this logic doesn't handle the case of real A and complex b
        tp = upcast(b.dtype, x.dtype, A.dtype)
        [b, x] = to_type(tp, [b, x])
        b = np.ravel(b)
        x = np.ravel(x)

        it = 0

        while True:  # it <= maxiter and normr >= tol:
            if len(self.levels) == 1:
                # hierarchy has only 1 level
                x = self.coarse_solver(A, b)
            else:
                self.__solve(0, x, b, cycle)

            it += 1

            normr = np.linalg.norm(b - A @ x)
            if residuals is not None:
                residuals.append(normr)

            if callback is not None:
                callback(x)

            if normr < tol * normb:
                if return_info:
                    return x, 0
                return x

            if it == maxiter:
                if return_info:
                    return x, it
                return x

    def __solve(self, lvl, x, b, cycle):
        """Multigrid cycling.

        Parameters
        ----------
        lvl : int
            Solve problem on level `lvl`
        x : numpy array
            Initial guess `x` and return correction
        b : numpy array
            Right-hand side for Ax=b
        cycle : {'V','W','F','AMLI'}
            Recursively called cycling function.  The
            Defines the cycling used:
            cycle = 'V',    V-cycle
            cycle = 'W',    W-cycle
            cycle = 'F',    F-cycle
            cycle = 'AMLI', AMLI-cycle

        """
        A = self.levels[lvl].A

        self.levels[lvl].presmoother(A, x, b)

        residual = b - A @ x

        coarse_b = self.levels[lvl].R * residual
        coarse_x = np.zeros_like(coarse_b)

        if lvl == len(self.levels) - 2:
            coarse_x[:] = self.coarse_solver(self.levels[-1].A, coarse_b)
        else:
            if cycle == 'V':
                self.__solve(lvl + 1, coarse_x, coarse_b, 'V')
            elif cycle == 'W':
                self.__solve(lvl + 1, coarse_x, coarse_b, cycle)
                self.__solve(lvl + 1, coarse_x, coarse_b, cycle)
            elif cycle == 'F':
                self.__solve(lvl + 1, coarse_x, coarse_b, cycle)
                self.__solve(lvl + 1, coarse_x, coarse_b, 'V')
            elif cycle == 'AMLI':
                # Run nAMLI AMLI cycles, which compute "optimal" corrections by
                # orthogonalizing the coarse-grid corrections in the A-norm
                nAMLI = 2
                Ac = self.levels[lvl + 1].A
                p = np.zeros((nAMLI, coarse_b.shape[0]), dtype=coarse_b.dtype)
                beta = np.zeros((nAMLI, nAMLI), dtype=coarse_b.dtype)
                for k in range(nAMLI):
                    # New search direction --> M^{-1}*residual
                    p[k, :] = 1
                    self.__solve(lvl + 1, p[k, :].reshape(coarse_b.shape),
                                 coarse_b, cycle)

                    # Orthogonalize new search direction to old directions
                    for j in range(k):  # loops from j = 0...(k-1)
                        beta[k, j] = np.inner(p[j, :].conj(), Ac * p[k, :]) /\
                            np.inner(p[j, :].conj(), Ac * p[j, :])
                        p[k, :] -= beta[k, j] * p[j, :]

                    # Compute step size
                    Ap = Ac * p[k, :]
                    alpha = np.inner(p[k, :].conj(), np.ravel(coarse_b)) /\
                        np.inner(p[k, :].conj(), Ap)

                    # Update solution
                    coarse_x += alpha * p[k, :].reshape(coarse_x.shape)

                    # Update residual
                    coarse_b -= alpha * Ap.reshape(coarse_b.shape)
            else:
                raise TypeError(f'Unrecognized cycle type ({cycle})')

        x += self.levels[lvl].P * coarse_x   # coarse grid correction

        self.levels[lvl].postsmoother(A, x, b) def change_smoothers(ml, presmoother, postsmoother):
    """Initialize pre and post smoothers.

    Initialize pre- and post- smoothers throughout a MultilevelSolver, with
    the option of having different smoothers at different levels

    For each level of the MultilevelSolver 'ml' (except the coarsest level),
    initialize the .presmoother() and .postsmoother() methods used in the
    multigrid cycle.

    Parameters
    ----------
    ml : pyamg multilevel hierarchy
        Data structure that stores the multigrid hierarchy.
    presmoother : None, string, tuple, list
        presmoother can be (1) the name of a supported smoother, e.g.
        "gauss_seidel", (2) a tuple of the form ('method','opts') where
        'method' is the name of a supported smoother and 'opts' a dict of
        keyword arguments to the smoother, or (3) a list of instances of
        options 1 or 2.  See the Examples section for illustrations of the
        format.

        If presmoother is a list, presmoother[i] determines the smoothing
        strategy for level i.  Else, presmoother defines the same strategy
        for all levels.

        If len(presmoother) < len(ml.levels), then
        presmoother[-1] is used for all remaining levels

        If len(presmoother) > len(ml.levels), then
        the remaining smoothing strategies are ignored

    postsmoother : string, tuple, list
        Defines postsmoother in identical fashion to presmoother

    Returns
    -------
    ml changed in place
    ml.levels[i].presmoother   <===  presmoother[i]
    ml.levels[i].postsmoother  <===  postsmoother[i]
    ml.symmetric_smoothing is marked True/False depending on whether
        the smoothing scheme is symmetric.

    Notes
    -----
    - Parameter 'omega' of the Jacobi, Richardson, and jacobi_ne
      methods is scaled by the spectral radius of the matrix on
      each level.  Therefore 'omega' should be in the interval (0,2).
    - Parameter 'withrho' (default: True) controls whether the omega is
      rescaled by the spectral radius in jacobi, block_jacobi, and jacobi_ne
    - By initializing the smoothers after the hierarchy has been setup, allows
      for "algebraically" directed relaxation, such as strength_based_schwarz,
      which uses only the strong connections of a degree-of-freedom to define
      overlapping regions
    - Available smoother methods::

        gauss_seidel
        block_gauss_seidel
        jacobi
        block_jacobi
        richardson
        sor
        chebyshev
        gauss_seidel_nr
        gauss_seidel_ne
        jacobi_ne
        cg
        gmres
        cgne
        cgnr
        schwarz
        strength_based_schwarz
        None

    Examples
    --------
    >>> from pyamg.gallery import poisson
    >>> from pyamg.aggregation import smoothed_aggregation_solver
    >>> from pyamg.relaxation.smoothing import change_smoothers
    >>> from pyamg.util.linalg import norm
    >>> import numpy as np
    >>> A = poisson((10,10), format='csr')
    >>> b = np.random.rand(A.shape[0],)
    >>> ml = smoothed_aggregation_solver(A, max_coarse=10)
    >>> # Set all levels to use gauss_seidel's defaults
    >>> smoothers = 'gauss_seidel'
    >>> change_smoothers(ml, presmoother=smoothers, postsmoother=smoothers)
    >>> residuals=[]
    >>> x = ml.solve(b, tol=1e-8, residuals=residuals)
    >>> # Set all levels to use three iterations of gauss_seidel's defaults
    >>> smoothers = ('gauss_seidel', {'iterations' : 3})
    >>> change_smoothers(ml, presmoother=smoothers, postsmoother=None)
    >>> residuals=[]
    >>> x = ml.solve(b, tol=1e-8, residuals=residuals)
    >>> # Set level 0 to use gauss_seidel's defaults, and all
    >>> # subsequent levels to use 5 iterations of cgnr
    >>> smoothers = ['gauss_seidel', ('cgnr', {'maxiter' : 5})]
    >>> change_smoothers(ml, presmoother=smoothers, postsmoother=smoothers)
    >>> residuals=[]
    >>> x = ml.solve(b, tol=1e-8, residuals=residuals)

    """
    ml.symmetric_smoothing = True

    # interpret arguments into list
    if isinstance(presmoother, (str, tuple)) or (presmoother is None):
        presmoother = [presmoother]
    elif not isinstance(presmoother, list):
        raise ValueError('Unrecognized presmoother')

    if isinstance(postsmoother, (str, tuple)) or (postsmoother is None):
        postsmoother = [postsmoother]
    elif not isinstance(postsmoother, list):
        raise ValueError('Unrecognized postsmoother')

    # set ml.levels[i].presmoother = presmoother[i],
    #     ml.levels[i].postsmoother = postsmoother[i]
    fn1 = None      # Predefine to keep scope beyond first loop
    fn2 = None
    kwargs1 = {}
    kwargs2 = {}
    min_len = min(len(presmoother), len(postsmoother), len(ml.levels[:-1]))
    # same = (len(presmoother) == len(postsmoother))
    for i in range(0, min_len):
        # unpack presmoother[i]
        fn1, kwargs1 = _unpack_arg(presmoother[i])
        # get function handle
        try:
            setup_presmoother = _setup_call(str(fn1).lower())
        except NameError as e:
            raise NameError(f'Invalid presmoother method: {fn1}') from e

        ml.levels[i].presmoother = setup_presmoother(ml.levels[i], **kwargs1)

        # unpack postsmoother[i]
        fn2, kwargs2 = _unpack_arg(postsmoother[i])
        # get function handle
        try:
            setup_postsmoother = _setup_call(str(fn2).lower())
        except NameError as e:
            raise NameError(f'Invalid postsmoother method: {fn2}') from e
        ml.levels[i].postsmoother = setup_postsmoother(ml.levels[i], **kwargs2)

        # Check if symmetric smoothing scheme
        if 'iterations' in kwargs1:
            it1 = kwargs1['iterations']
        else:
            it1 = DEFAULT_NITER

        if 'iterations' in kwargs2:
            it2 = kwargs2['iterations']
        else:
            it2 = DEFAULT_NITER

        if (fn1 != fn2) or (it1 != it2):
            ml.symmetric_smoothing = False
        elif fn1 not in SYMMETRIC_RELAXATION:
            if 'sweep' in kwargs1:
                sweep1 = kwargs1['sweep']
            else:
                sweep1 = DEFAULT_SWEEP

            if 'sweep' in kwargs2:
                sweep2 = kwargs2['sweep']
            else:
                sweep2 = DEFAULT_SWEEP

            if (sweep1, sweep2) not in (('forward', 'backward'),
                                        ('backward', 'forward'),
                                        ('symmetric', 'symmetric')):
                ml.symmetric_smoothing = False

    if len(presmoother) < len(postsmoother):
        mid_len = min(len(postsmoother), len(ml.levels[:-1]))
        for i in range(min_len, mid_len):
            # Set up presmoother
            ml.levels[i].presmoother = setup_presmoother(ml.levels[i], **kwargs1)

            # unpack postsmoother[i]
            fn2, kwargs2 = _unpack_arg(postsmoother[i])
            # get function handle
            try:
                setup_postsmoother = _setup_call(str(fn2).lower())
            except NameError as e:
                raise NameError(f'Invalid postsmoother method: {fn2}') from e

            ml.levels[i].postsmoother = setup_postsmoother(ml.levels[i], **kwargs2)

            # Check if symmetric smoothing scheme
            if 'iterations' in kwargs1:
                it1 = kwargs1['iterations']
            else:
                it1 = DEFAULT_NITER

            if 'iterations' in kwargs2:
                it2 = kwargs2['iterations']
            else:
                it2 = DEFAULT_NITER

            if (fn1 != fn2) or (it1 != it2):
                ml.symmetric_smoothing = False
            elif fn1 not in SYMMETRIC_RELAXATION:
                if 'sweep' in kwargs1:
                    sweep1 = kwargs1['sweep']
                else:
                    sweep1 = DEFAULT_SWEEP

                if 'sweep' in kwargs2:
                    sweep2 = kwargs2['sweep']
                else:
                    sweep2 = DEFAULT_SWEEP

                if (sweep1, sweep2) not in (('forward', 'backward'),
                                            ('backward', 'forward'),
                                            ('symmetric', 'symmetric')):
                    ml.symmetric_smoothing = False

    elif len(presmoother) > len(postsmoother):
        mid_len = min(len(presmoother), len(ml.levels[:-1]))
        for i in range(min_len, mid_len):
            # unpack presmoother[i]
            fn1, kwargs1 = _unpack_arg(presmoother[i])
            # get function handle
            try:
                setup_presmoother = _setup_call(str(fn1).lower())
            except NameError as e:
                raise NameError(f'Invalid presmoother method: {fn1}') from e

            ml.levels[i].presmoother = setup_presmoother(ml.levels[i], **kwargs1)

            # Set up postsmoother
            ml.levels[i].postsmoother = setup_postsmoother(ml.levels[i], **kwargs2)

            # Check if symmetric smoothing scheme
            if 'iterations' in kwargs1:
                it1 = kwargs1['iterations']
            else:
                it1 = DEFAULT_NITER

            if 'iterations' in kwargs2:
                it2 = kwargs2['iterations']
            else:
                it2 = DEFAULT_NITER

            if (fn1 != fn2) or (it1 != it2):
                ml.symmetric_smoothing = False
            elif fn1 not in SYMMETRIC_RELAXATION:
                if 'sweep' in kwargs1:
                    sweep1 = kwargs1['sweep']
                else:
                    sweep1 = DEFAULT_SWEEP

                if 'sweep' in kwargs2:
                    sweep2 = kwargs2['sweep']
                else:
                    sweep2 = DEFAULT_SWEEP

                if (sweep1, sweep2) not in (('forward', 'backward'),
                                            ('backward', 'forward'),
                                            ('symmetric', 'symmetric')):
                    ml.symmetric_smoothing = False

    else:
        mid_len = min_len

    # Fill in remaining levels
    for i in range(mid_len, len(ml.levels[:-1])):
        ml.levels[i].presmoother = setup_presmoother(ml.levels[i], **kwargs1)
        ml.levels[i].postsmoother = setup_postsmoother(ml.levels[i], **kwargs2) def eliminate_diag_dom_nodes(A, C, theta=1.02):
    r"""Eliminate diagonally dominance.

    Helper function that eliminates diagonally dominant rows and cols from A
    in the separate matrix C.  This is useful because it eliminates nodes in C
    which we don't want coarsened.  These eliminated nodes in C just become
    the rows and columns of the identity.

    Parameters
    ----------
    A : {csr_matrix, bsr_matrix}
        Sparse NxN matrix
    C : {csr_matrix}
        Sparse MxM matrix, where M is the number of nodes in A.  M=N if A
        is CSR or is BSR with blocksize 1.  Otherwise M = N/blocksize.
    theta : {float}
        determines diagonal dominance threshhold

    Returns
    -------
    C : {csr_matrix}
        C updated such that the rows and columns corresponding to diagonally
        dominant rows in A have been eliminated and replaced with rows and
        columns of the identity.

    Notes
    -----
    Diagonal dominance is defined as
    :math:`\| (e_i, A) - a_{ii} \|_1  <  \\theta a_{ii}`
    that is, the 1-norm of the off diagonal elements in row i must be less than
    theta times the diagonal element.

    Examples
    --------
    >>> from pyamg.gallery import poisson
    >>> from pyamg.util.utils import eliminate_diag_dom_nodes
    >>> A = poisson( (4,), format='csr' )
    >>> C = eliminate_diag_dom_nodes(A, A.copy(), 1.1)
    >>> C.toarray()
    array([[ 1.,  0.,  0.,  0.],
           [ 0.,  2., -1.,  0.],
           [ 0., -1.,  2.,  0.],
           [ 0.,  0.,  0.,  1.]])

    """
    # Find the diagonally dominant rows in A.
    A_abs = A.copy()
    A_abs.data = np.abs(A_abs.data)
    D_abs = get_diagonal(A_abs, norm_eq=0, inv=False)
    diag_dom_rows = (D_abs > (theta*(A_abs*np.ones((A_abs.shape[0],),
                                                   dtype=A_abs) - D_abs)))

    # Account for BSR matrices and translate diag_dom_rows from dofs to nodes
    bsize = get_blocksize(A_abs)
    if bsize > 1:
        diag_dom_rows = np.array(diag_dom_rows, dtype=int)
        diag_dom_rows = diag_dom_rows.reshape(-1, bsize)
        diag_dom_rows = np.sum(diag_dom_rows, axis=1)
        diag_dom_rows = (diag_dom_rows == bsize)

    # Replace these rows/cols in # C with rows/cols of the identity.
    Id = eye(C.shape[0], C.shape[1], format='csr')
    Id.data[diag_dom_rows] = 0.0
    C = Id * C * Id
    Id.data[diag_dom_rows] = 1.0
    Id.data[np.where(diag_dom_rows == 0)[0]] = 0.0
    C = C + Id

    del A_abs
    return C def get_blocksize(A):
    """Return the block size of a matrix."""
    if isspmatrix_bsr(A):
        return A.blocksize[0]
    return 1 def levelize_strength_or_aggregation(to_levelize, max_levels, max_coarse):
    """Turn parameter into a list per level.

    Helper function to preprocess the strength and aggregation parameters
    passed to smoothed_aggregation_solver and rootnode_solver.

    Parameters
    ----------
    to_levelize : {string, tuple, list}
        Parameter to preprocess, i.e., levelize and convert to a level-by-level
        list such that entry i specifies the parameter at level i
    max_levels : int
        Defines the maximum number of levels considered
    max_coarse : int
        Defines the maximum coarse grid size allowed

    Returns
    -------
    (max_levels, max_coarse, to_levelize) : tuple
        New max_levels and max_coarse values and then the parameter list
        to_levelize, such that entry i specifies the parameter choice at level
        i.  max_levels and max_coarse are returned, because they may be updated
        if strength or aggregation set a predefined coarsening and possibly
        change these values.

    Notes
    -----
    This routine is needed because the user will pass in a parameter option
    such as smooth='jacobi', or smooth=['jacobi', None], and this option must
    be "levelized", or converted to a list of length max_levels such that entry
    [i] in that list is the parameter choice for level i.

    The parameter choice in to_levelize can be a string, tuple or list.  If
    it is a string or tuple, then that option is assumed to be the
    parameter setting at every level.  If to_levelize is inititally a list,
    if the length of the list is less than max_levels, the last entry in the
    list defines that parameter for all subsequent levels.


    Examples
    --------
    >>> from pyamg.util.utils import levelize_strength_or_aggregation
    >>> strength = ['evolution', 'classical']
    >>> levelize_strength_or_aggregation(strength, 4, 10)
    (4, 10, ['evolution', 'classical', 'classical'])

    """
    if isinstance(to_levelize, tuple):
        if to_levelize[0] == 'predefined':
            to_levelize = [to_levelize]
            max_levels = 2
            max_coarse = 0
        else:
            to_levelize = [to_levelize for i in range(max_levels-1)]

    elif isinstance(to_levelize, str):
        if to_levelize == 'predefined':
            raise ValueError('predefined to_levelize requires a user-provided '
                             'CSR matrix representing strength or aggregation '
                             'i.e., ("predefined", {"C" : CSR_MAT}).')  # noqa: FS003
        to_levelize = [to_levelize for i in range(max_levels-1)]

    elif isinstance(to_levelize, list):
        if isinstance(to_levelize[-1], tuple) and\
           (to_levelize[-1][0] == 'predefined'):
            # to_levelize is a list that ends with a predefined operator
            max_levels = len(to_levelize) + 1
            max_coarse = 0
        else:
            # to_levelize a list that __doesn't__ end with 'predefined'
            if len(to_levelize) < max_levels-1:
                mlz = max_levels - 1 - len(to_levelize)
                toext = [to_levelize[-1] for i in range(mlz)]
                to_levelize.extend(toext)

    elif to_levelize is None:
        to_levelize = [(None, {}) for i in range(max_levels-1)]
    else:
        raise ValueError('invalid to_levelize')

    return max_levels, max_coarse, to_levelize def levelize_smooth_or_improve_candidates(to_levelize, max_levels):
    """Turn parameter in to a list per level.

    Helper function to preprocess the smooth and improve_candidates
    parameters passed to smoothed_aggregation_solver and rootnode_solver.

    Parameters
    ----------
    to_levelize : {string, tuple, list}
        Parameter to preprocess, i.e., levelize and convert to a level-by-level
        list such that entry i specifies the parameter at level i
    max_levels : int
        Defines the maximum number of levels considered

    Returns
    -------
    to_levelize : list
        The parameter list such that entry i specifies the parameter choice
        at level i.

    Notes
    -----
    This routine is needed because the user will pass in a parameter option
    such as smooth='jacobi', or smooth=['jacobi', None], and this option must
    be "levelized", or converted to a list of length max_levels such that entry
    [i] in that list is the parameter choice for level i.

    The parameter choice in to_levelize can be a string, tuple or list.  If
    it is a string or tuple, then that option is assumed to be the
    parameter setting at every level.  If to_levelize is inititally a list,
    if the length of the list is less than max_levels, the last entry in the
    list defines that parameter for all subsequent levels.

    Examples
    --------
    >>> from pyamg.util.utils import levelize_smooth_or_improve_candidates
    >>> improve_candidates = ['gauss_seidel', None]
    >>> levelize_smooth_or_improve_candidates(improve_candidates, 4)
    ['gauss_seidel', None, None, None]

    """
    # handle default value (mutable)
    # improve_candidates=(('block_gauss_seidel',
    #                      {'sweep': 'symmetric', 'iterations': 4}),
    #                     None)
    # -> make it a list
    if isinstance(to_levelize, tuple):
        if isinstance(to_levelize[0], tuple):
            to_levelize = list(to_levelize)

    if isinstance(to_levelize, (str, tuple)):
        to_levelize = [to_levelize for i in range(max_levels)]
    elif isinstance(to_levelize, list):
        if len(to_levelize) < max_levels:
            mlz = max_levels - len(to_levelize)
            toext = [to_levelize[-1] for i in range(mlz)]
            to_levelize.extend(toext)
    elif to_levelize is None:
        to_levelize = [(None, {}) for i in range(max_levels)]

    return to_levelize def classical_strength_of_connection(A, theta=0.0, norm='abs'):
    """Classical Strength Measure.

    Return a strength of connection matrix using the classical AMG measure
    An off-diagonal entry A[i,j] is a strong connection iff::

             |A[i,j]| >= theta * max(|A[i,k]|), where k != i     (norm='abs')
             -A[i,j]  >= theta * max(-A[i,k]),  where k != i     (norm='min')

    Parameters
    ----------
    A : csr_matrix or bsr_matrix
        Square, sparse matrix in CSR or BSR format
    theta : float
        Threshold parameter in [0,1].
    norm: 'string'
        'abs' : to use the absolute value,
        'min' : to use the negative value (see above)

    Returns
    -------
    S : csr_matrix
        Matrix graph defining strong connections.  S[i,j]=1 if vertex i
        is strongly influenced by vertex j.

    See Also
    --------
    symmetric_strength_of_connection : symmetric measure used in SA
    evolution_strength_of_connection : relaxation based strength measure

    Notes
    -----
    - A symmetric A does not necessarily yield a symmetric strength matrix S
    - Calls C++ function classical_strength_of_connection
    - The version as implemented is designed form M-matrices.  Trottenberg et
      al. use max A[i,k] over all negative entries, which is the same.  A
      positive edge weight never indicates a strong connection.
    - See [2000BrHeMc]_ and [2001bTrOoSc]_

    References
    ----------
    .. [2000BrHeMc] Briggs, W. L., Henson, V. E., McCormick, S. F., "A multigrid
        tutorial", Second edition. Society for Industrial and Applied
        Mathematics (SIAM), Philadelphia, PA, 2000. xii+193 pp.

    .. [2001bTrOoSc] Trottenberg, U., Oosterlee, C. W., Schuller, A., "Multigrid",
        Academic Press, Inc., San Diego, CA, 2001. xvi+631 pp.

    Examples
    --------
    >>> import numpy as np
    >>> from pyamg.gallery import stencil_grid
    >>> from pyamg.strength import classical_strength_of_connection
    >>> n=3
    >>> stencil = np.array([[-1.0,-1.0,-1.0],
    ...                        [-1.0, 8.0,-1.0],
    ...                        [-1.0,-1.0,-1.0]])
    >>> A = stencil_grid(stencil, (n,n), format='csr')
    >>> S = classical_strength_of_connection(A, 0.0)

    """
    if sparse.isspmatrix_bsr(A):
        blocksize = A.blocksize[0]
    else:
        blocksize = 1

    if not sparse.isspmatrix_csr(A):
        warn('Implicit conversion of A to csr', sparse.SparseEfficiencyWarning)
        A = sparse.csr_matrix(A)

    if (theta < 0 or theta > 1):
        raise ValueError('expected theta in [0,1]')

    Sp = np.empty_like(A.indptr)
    Sj = np.empty_like(A.indices)
    Sx = np.empty_like(A.data)

    if norm not in ('abs', 'min'):
        raise ValueError('Unknown norm')

    if norm == 'abs':
        amg_core.classical_strength_of_connection_abs(
            A.shape[0], theta, A.indptr, A.indices, A.data, Sp, Sj, Sx)
    elif norm == 'min':
        amg_core.classical_strength_of_connection_min(
            A.shape[0], theta, A.indptr, A.indices, A.data, Sp, Sj, Sx)

    S = sparse.csr_matrix((Sx, Sj, Sp), shape=A.shape)

    if blocksize > 1:
        S = amalgamate(S, blocksize)

    # Strength represents "distance", so take the magnitude
    S.data = np.abs(S.data)

    # Scale S by the largest magnitude entry in each row
    S = scale_rows_by_largest_entry(S)

    return S def symmetric_strength_of_connection(A, theta=0):
    """Symmetric Strength Measure.

    Compute strength of connection matrix using the standard symmetric measure

    An off-diagonal connection A[i,j] is strong iff::

        abs(A[i,j]) >= theta * sqrt( abs(A[i,i]) * abs(A[j,j]) )

    Parameters
    ----------
    A : csr_matrix
        Matrix graph defined in sparse format.  Entry A[i,j] describes the
        strength of edge [i,j]
    theta : float
        Threshold parameter (positive).

    Returns
    -------
    S : csr_matrix
        Matrix graph defining strong connections.  S[i,j]=1 if vertex i
        is strongly influenced by vertex j.

    See Also
    --------
    symmetric_strength_of_connection : symmetric measure used in SA
    evolution_strength_of_connection : relaxation based strength measure

    Notes
    -----
        - For vector problems, standard strength measures may produce
          undesirable aggregates.  A "block approach" from Vanek et al. is used
          to replace vertex comparisons with block-type comparisons.  A
          connection between nodes i and j in the block case is strong if::

          ||AB[i,j]|| >= theta * sqrt( ||AB[i,i]||*||AB[j,j]|| ) where AB[k,l]

          is the matrix block (degrees of freedom) associated with nodes k and
          l and ||.|| is a matrix norm, such a Frobenius.

        - See [1996bVaMaBr]_ for more details.

    References
    ----------
    .. [1996bVaMaBr] Vanek, P. and Mandel, J. and Brezina, M.,
       "Algebraic Multigrid by Smoothed Aggregation for
       Second and Fourth Order Elliptic Problems",
       Computing, vol. 56, no. 3, pp. 179--196, 1996.
       http://citeseer.ist.psu.edu/vanek96algebraic.html

    Examples
    --------
    >>> import numpy as np
    >>> from pyamg.gallery import stencil_grid
    >>> from pyamg.strength import symmetric_strength_of_connection
    >>> n=3
    >>> stencil = np.array([[-1.0,-1.0,-1.0],
    ...                        [-1.0, 8.0,-1.0],
    ...                        [-1.0,-1.0,-1.0]])
    >>> A = stencil_grid(stencil, (n,n), format='csr')
    >>> S = symmetric_strength_of_connection(A, 0.0)

    """
    if theta < 0:
        raise ValueError('expected a positive theta')

    if not sparse.isspmatrix_csr(A) and not sparse.isspmatrix_bsr(A):
        raise TypeError('expected csr_matrix or bsr_matrix')

    if sparse.isspmatrix_csr(A):
        # if theta == 0:
        #     return A

        Sp = np.empty_like(A.indptr)
        Sj = np.empty_like(A.indices)
        Sx = np.empty_like(A.data)

        fn = amg_core.symmetric_strength_of_connection
        fn(A.shape[0], theta, A.indptr, A.indices, A.data, Sp, Sj, Sx)

        S = sparse.csr_matrix((Sx, Sj, Sp), shape=A.shape)

    elif sparse.isspmatrix_bsr(A):
        M, N = A.shape
        R, C = A.blocksize

        if R != C:
            raise ValueError('matrix must have square blocks')

        if theta == 0:
            data = np.ones(len(A.indices), dtype=A.dtype)
            S = sparse.csr_matrix((data, A.indices.copy(), A.indptr.copy()),
                                  shape=(int(M / R), int(N / C)))
        else:
            # the strength of connection matrix is based on the
            # Frobenius norms of the blocks
            data = (np.conjugate(A.data) * A.data).reshape(-1, R * C)
            data = data.sum(axis=1)
            A = sparse.csr_matrix((data, A.indices, A.indptr),
                                  shape=(int(M / R), int(N / C)))
            return symmetric_strength_of_connection(A, theta)

    # Strength represents "distance", so take the magnitude
    S.data = np.abs(S.data)

    # Scale S by the largest magnitude entry in each row
    S = scale_rows_by_largest_entry(S)

    return S def evolution_strength_of_connection(A, B=None, epsilon=4.0, k=2,
                                     proj_type='l2', block_flag=False,
                                     symmetrize_measure=True):
    """Evolution Strength Measure.

    Construct strength of connection matrix using an Evolution-based measure

    Parameters
    ----------
    A : csr_matrix, bsr_matrix
        Sparse NxN matrix
    B : string, array
        If B=None, then the near nullspace vector used is all ones.  If B is
        an (NxK) array, then B is taken to be the near nullspace vectors.
    epsilon : scalar
        Drop tolerance
    k : integer
        ODE num time steps, step size is assumed to be 1/rho(DinvA)
    proj_type : {'l2','D_A'}
        Define norm for constrained min prob, i.e. define projection
    block_flag : boolean
        If True, use a block D inverse as preconditioner for A during
        weighted-Jacobi

    Returns
    -------
    Atilde : csr_matrix
        Sparse matrix of strength values

    See [2008OlScTu]_ for more details.

    References
    ----------
    .. [2008OlScTu] Olson, L. N., Schroder, J., Tuminaro, R. S.,
       "A New Perspective on Strength Measures in Algebraic Multigrid",
       submitted, June, 2008.

    Examples
    --------
    >>> import numpy as np
    >>> from pyamg.gallery import stencil_grid
    >>> from pyamg.strength import evolution_strength_of_connection
    >>> n=3
    >>> stencil =  np.array([[-1.0,-1.0,-1.0],
    ...                        [-1.0, 8.0,-1.0],
    ...                        [-1.0,-1.0,-1.0]])
    >>> A = stencil_grid(stencil, (n,n), format='csr')
    >>> S = evolution_strength_of_connection(A,  np.ones((A.shape[0],1)))

    """
    # ====================================================================
    # Check inputs
    if epsilon < 1.0:
        raise ValueError('expected epsilon > 1.0')
    if k <= 0:
        raise ValueError('number of time steps must be > 0')
    if proj_type not in ['l2', 'D_A']:
        raise ValueError('proj_type must be "l2" or "D_A"')
    if (not sparse.isspmatrix_csr(A)) and (not sparse.isspmatrix_bsr(A)):
        raise TypeError('expected csr_matrix or bsr_matrix')

    # ====================================================================
    # Format A and B correctly.
    # B must be in mat format, this isn't a deep copy
    if B is None:
        Bmat = np.ones((A.shape[0], 1), dtype=A.dtype)
    else:
        Bmat = np.asarray(B)

    # Pre-process A.  We need A in CSR, to be devoid of explicit 0's and have
    # sorted indices
    if not sparse.isspmatrix_csr(A):
        csrflag = False
        numPDEs = A.blocksize[0]
        D = A.diagonal()
        # Calculate Dinv*A
        if block_flag:
            Dinv = get_block_diag(A, blocksize=numPDEs, inv_flag=True)
            Dinv = sparse.bsr_matrix((Dinv, np.arange(Dinv.shape[0]),
                                      np.arange(Dinv.shape[0] + 1)),
                                     shape=A.shape)
            Dinv_A = (Dinv * A).tocsr()
        else:
            Dinv = np.zeros_like(D)
            mask = (D != 0.0)
            Dinv[mask] = 1.0 / D[mask]
            Dinv[D == 0] = 1.0
            Dinv_A = scale_rows(A, Dinv, copy=True)
        A = A.tocsr()
    else:
        csrflag = True
        numPDEs = 1
        D = A.diagonal()
        Dinv = np.zeros_like(D)
        mask = (D != 0.0)
        Dinv[mask] = 1.0 / D[mask]
        Dinv[D == 0] = 1.0
        Dinv_A = scale_rows(A, Dinv, copy=True)

    A.eliminate_zeros()
    A.sort_indices()

    # Handle preliminaries for the algorithm
    dimen = A.shape[1]
    NullDim = Bmat.shape[1]

    # Get spectral radius of Dinv*A, this will be used to scale the time step
    # size for the ODE
    rho_DinvA = approximate_spectral_radius(Dinv_A)

    # Calculate D_A for later use in the minimization problem
    if proj_type == 'D_A':
        D_A = sparse.spdiags([D], [0], dimen, dimen, format='csr')
    else:
        D_A = sparse.eye(dimen, dimen, format='csr', dtype=A.dtype)

    # Calculate (I - delta_t Dinv A)^k
    #      In order to later access columns, we calculate the transpose in
    #      CSR format so that columns will be accessed efficiently
    # Calculate the number of time steps that can be done by squaring, and
    # the number of time steps that must be done incrementally
    nsquare = int(np.log2(k))
    ninc = k - 2**nsquare

    # Calculate one time step
    Id = sparse.eye(dimen, dimen, format='csr', dtype=A.dtype)
    Atilde = (Id - (1.0 / rho_DinvA) * Dinv_A)
    Atilde = Atilde.T.tocsr()

    # Construct a sparsity mask for Atilde that will restrict Atilde^T to the
    # nonzero pattern of A, with the added constraint that row i of Atilde^T
    # retains only the nonzeros that are also in the same PDE as i.
    mask = A.copy()

    # Restrict to same PDE
    if numPDEs > 1:
        row_length = np.diff(mask.indptr)
        my_pde = np.mod(np.arange(dimen), numPDEs)
        my_pde = np.repeat(my_pde, row_length)
        mask.data[np.mod(mask.indices, numPDEs) != my_pde] = 0.0
        del row_length, my_pde
        mask.eliminate_zeros()

    # If the total number of time steps is a power of two, then there is
    # a very efficient computational short-cut.  Otherwise, we support
    # other numbers of time steps, through an inefficient algorithm.
    if ninc > 0:
        warn('The most efficient time stepping for the Evolution Strength '
             f'Method is done in powers of two.\nYou have chosen {k} time steps.')

        # Calculate (Atilde^nsquare)^T = (Atilde^T)^nsquare
        for _i in range(nsquare):
            Atilde = Atilde * Atilde

        JacobiStep = (Id - (1.0 / rho_DinvA) * Dinv_A).T.tocsr()
        for _i in range(ninc):
            Atilde = Atilde * JacobiStep
        del JacobiStep

        # Apply mask to Atilde, zeros in mask have already been eliminated at
        # start of routine.
        mask.data[:] = 1.0
        Atilde = Atilde.multiply(mask)
        Atilde.eliminate_zeros()
        Atilde.sort_indices()

    elif nsquare == 0:
        if numPDEs > 1:
            # Apply mask to Atilde, zeros in mask have already been eliminated
            # at start of routine.
            mask.data[:] = 1.0
            Atilde = Atilde.multiply(mask)
            Atilde.eliminate_zeros()
            Atilde.sort_indices()

    else:
        # Use computational short-cut for case (ninc == 0) and (nsquare > 0)
        # Calculate Atilde^k only at the sparsity pattern of mask.
        for _i in range(nsquare - 1):
            Atilde = Atilde * Atilde

        # Call incomplete mat-mat mult
        AtildeCSC = Atilde.tocsc()
        AtildeCSC.sort_indices()
        mask.sort_indices()
        Atilde.sort_indices()
        amg_core.incomplete_mat_mult_csr(Atilde.indptr, Atilde.indices,
                                         Atilde.data, AtildeCSC.indptr,
                                         AtildeCSC.indices, AtildeCSC.data,
                                         mask.indptr, mask.indices, mask.data,
                                         dimen)

        del AtildeCSC, Atilde
        Atilde = mask
        Atilde.eliminate_zeros()
        Atilde.sort_indices()

    del Dinv, Dinv_A, mask

    # Calculate strength based on constrained min problem of
    # min( z - B*x ), such that
    # (B*x)|_i = z|_i, i.e. they are equal at point i
    # z = (I - (t/k) Dinv A)^k delta_i
    #
    # Strength is defined as the relative point-wise approx. error between
    # B*x and z.  We don't use the full z in this problem, only that part of
    # z that is in the sparsity pattern of A.
    #
    # Can use either the D-norm, and inner product, or l2-norm and inner-prod
    # to solve the constrained min problem.  Using D gives scale invariance.
    #
    # This is a quadratic minimization problem with a linear constraint, so
    # we can build a linear system and solve it to find the critical point,
    # i.e. minimum.
    #
    # We exploit a known shortcut for the case of NullDim = 1.  The shortcut is
    # mathematically equivalent to the longer constrained min. problem

    if NullDim == 1:
        # Use shortcut to solve constrained min problem if B is only a vector
        # Strength(i,j) = | 1 - (z(i)/b(j))/(z(j)/b(i)) |
        # These ratios can be calculated by diagonal row and column scalings

        # Create necessary vectors for scaling Atilde
        #   Its not clear what to do where B == 0.  This is an
        #   an easy programming solution, that may make sense.
        Bmat_forscaling = np.ravel(Bmat)
        Bmat_forscaling[Bmat_forscaling == 0] = 1.0
        DAtilde = Atilde.diagonal()
        DAtildeDivB = np.ravel(DAtilde) / Bmat_forscaling

        # Calculate best approximation, z_tilde, in span(B)
        #   Importantly, scale_rows and scale_columns leave zero entries
        #   in the matrix.  For previous implementations this was useful
        #   because we assume data and Atilde.data are the same length below
        data = Atilde.data.copy()
        Atilde.data[:] = 1.0
        Atilde = scale_rows(Atilde, DAtildeDivB)
        Atilde = scale_columns(Atilde, np.ravel(Bmat_forscaling))

        # If angle in the complex plane between z and z_tilde is
        # greater than 90 degrees, then weak.  We can just look at the
        # dot product to determine if angle is greater than 90 degrees.
        angle = np.multiply(np.real(Atilde.data), np.real(data)) +\
            np.multiply(np.imag(Atilde.data), np.imag(data))
        angle = angle < 0.0
        angle = np.array(angle, dtype=bool)

        # Calculate Approximation ratio
        Atilde.data = Atilde.data / data

        # If approximation ratio is less than tol, then weak connection
        weak_ratio = (np.abs(Atilde.data) < 1e-4)

        # Calculate Approximation error
        Atilde.data = abs(1.0 - Atilde.data)

        # Set small ratios and large angles to weak
        Atilde.data[weak_ratio] = 0.0
        Atilde.data[angle] = 0.0

        # Set near perfect connections to 1e-4
        Atilde.eliminate_zeros()
        Atilde.data[Atilde.data < np.sqrt(np.finfo(float).eps)] = 1e-4

        del data, weak_ratio, angle

    else:
        # For use in computing local B_i^H*B, precompute the element-wise
        # multiply of each column of B with each other column.  We also scale
        # by 2.0 to account for BDB's eventual use in a constrained
        # minimization problem
        BDBCols = int(np.sum(np.arange(NullDim + 1)))
        BDB = np.zeros((dimen, BDBCols), dtype=A.dtype)
        counter = 0
        for i in range(NullDim):
            for j in range(i, NullDim):
                BDB[:, counter] = 2.0 *\
                    (np.conjugate(np.ravel(Bmat[:, i])) * np.ravel(D_A * Bmat[:, j]))
                counter = counter + 1

        # Choose tolerance for dropping "numerically zero" values later
        tol = set_tol(Atilde.dtype)

        # Use constrained min problem to define strength
        amg_core.evolution_strength_helper(Atilde.data,
                                           Atilde.indptr,
                                           Atilde.indices,
                                           Atilde.shape[0],
                                           np.ravel(Bmat),
                                           np.ravel((D_A * B.conj()).T),
                                           np.ravel(BDB),
                                           BDBCols, NullDim, tol)

        Atilde.eliminate_zeros()

    # All of the strength values are real by this point, so ditch the complex
    # part
    Atilde.data = np.array(np.real(Atilde.data), dtype=float)

    # Apply drop tolerance
    if epsilon != np.inf:
        amg_core.apply_distance_filter(dimen, epsilon, Atilde.indptr,
                                       Atilde.indices, Atilde.data)
        Atilde.eliminate_zeros()

    # Symmetrize
    if symmetrize_measure:
        Atilde = 0.5 * (Atilde + Atilde.T)

    # Set diagonal to 1.0, as each point is strongly connected to itself.
    Id = sparse.eye(dimen, dimen, format='csr')
    Id.data -= Atilde.diagonal()
    Atilde = Atilde + Id

    # If converted BSR to CSR, convert back and return amalgamated matrix,
    #   i.e. the sparsity structure of the blocks of Atilde
    if not csrflag:
        Atilde = Atilde.tobsr(blocksize=(numPDEs, numPDEs))

        n_blocks = Atilde.indices.shape[0]
        blocksize = Atilde.blocksize[0] * Atilde.blocksize[1]
        CSRdata = np.zeros((n_blocks,))
        amg_core.min_blocks(n_blocks, blocksize,
                            np.ravel(np.asarray(Atilde.data)), CSRdata)
        # Atilde = sparse.csr_matrix((data, row, col), shape=(*,*))
        Atilde = sparse.csr_matrix((CSRdata, Atilde.indices, Atilde.indptr),
                                   shape=(int(Atilde.shape[0] / numPDEs),
                                          int(Atilde.shape[1] / numPDEs)))

    # Standardized strength values require small values be weak and large
    # values be strong.  So, we invert the algebraic distances computed here
    Atilde.data = 1.0 / Atilde.data

    # Scale C by the largest magnitude entry in each row
    Atilde = scale_rows_by_largest_entry(Atilde)

    return Atilde def energy_based_strength_of_connection(A, theta=0.0, k=2):
    """Energy Strength Measure.

    Compute a strength of connection matrix using an energy-based measure.

    Parameters
    ----------
    A : sparse-matrix
        matrix from which to generate strength of connection information
    theta : float
        Threshold parameter in [0,1]
    k : int
        Number of relaxation steps used to generate strength information

    Returns
    -------
    S : csr_matrix
        Matrix graph defining strong connections.  The sparsity pattern
        of S matches that of A.  For BSR matrices, S is a reduced strength
        of connection matrix that describes connections between supernodes.

    Notes
    -----
    This method relaxes with weighted-Jacobi in order to approximate the
    matrix inverse.  A normalized change of energy is then used to define
    point-wise strength of connection values.  Specifically, let v be the
    approximation to the i-th column of the inverse, then

    (S_ij)^2 = <v_j, v_j>_A / <v, v>_A,

    where v_j = v, such that entry j in v has been zeroed out.  As is common,
    larger values imply a stronger connection.

    Current implementation is a very slow pure-python implementation for
    experimental purposes, only.

    See [2006BrBrMaMaMc]_ for more details.

    References
    ----------
    .. [2006BrBrMaMaMc] Brannick, Brezina, MacLachlan, Manteuffel, McCormick.
       "An Energy-Based AMG Coarsening Strategy",
       Numerical Linear Algebra with Applications,
       vol. 13, pp. 133-148, 2006.

    Examples
    --------
    >>> import numpy as np
    >>> from pyamg.gallery import stencil_grid
    >>> from pyamg.strength import energy_based_strength_of_connection
    >>> n=3
    >>> stencil =  np.array([[-1.0,-1.0,-1.0],
    ...                      [-1.0, 8.0,-1.0],
    ...                      [-1.0,-1.0,-1.0]])
    >>> A = stencil_grid(stencil, (n,n), format='csr')
    >>> S = energy_based_strength_of_connection(A, 0.0)

    """
    if theta < 0:
        raise ValueError('expected a positive theta')
    if not sparse.isspmatrix(A):
        raise ValueError('expected sparse matrix')
    if k < 0:
        raise ValueError('expected positive number of steps')
    if not isinstance(k, int):
        raise ValueError('expected integer')

    if sparse.isspmatrix_bsr(A):
        bsr_flag = True
        numPDEs = A.blocksize[0]
        if A.blocksize[0] != A.blocksize[1]:
            raise ValueError('expected square blocks in BSR matrix A')
    else:
        bsr_flag = False

    # Convert A to csc and Atilde to csr
    if sparse.isspmatrix_csr(A):
        Atilde = A.copy()
        A = A.tocsc()
    else:
        A = A.tocsc()
        Atilde = A.copy()
        Atilde = Atilde.tocsr()

    # Calculate the weighted-Jacobi parameter
    D = A.diagonal()
    Dinv = 1.0 / D
    Dinv[D == 0] = 0.0
    Dinv = sparse.csc_matrix((Dinv, (np.arange(A.shape[0]),
                                     np.arange(A.shape[1]))), shape=A.shape)
    DinvA = Dinv * A
    omega = 1.0 / approximate_spectral_radius(DinvA)
    del DinvA

    # Approximate A-inverse with k steps of w-Jacobi and a zero initial guess
    S = sparse.csc_matrix(A.shape, dtype=A.dtype)  # empty matrix
    Id = sparse.eye(A.shape[0], A.shape[1], format='csc')
    for _i in range(k + 1):
        S = S + omega * (Dinv * (Id - A * S))

    # Calculate the strength entries in S column-wise, but only strength
    # values at the sparsity pattern of A
    for i in range(Atilde.shape[0]):
        v = S[:, i].toarray()
        v = v.ravel()
        Av = A @ v
        denom = np.sqrt(np.inner(v.conj(), Av))
        # replace entries in row i with strength values
        for j in range(Atilde.indptr[i], Atilde.indptr[i + 1]):
            col = Atilde.indices[j]
            vj = v[col].copy()
            v[col] = 0.0
            #   =  (||v_j||_A - ||v||_A) / ||v||_A
            val = np.sqrt(np.inner(v.conj(), A @ v)) / denom - 1.0

            # Negative values generally imply a weak connection
            if val > -0.01:
                Atilde.data[j] = abs(val)
            else:
                Atilde.data[j] = 0.0

            v[col] = vj

    # Apply drop tolerance
    Atilde = classical_strength_of_connection(Atilde, theta=theta)
    Atilde.eliminate_zeros()

    # Put ones on the diagonal
    Atilde = Atilde + Id.tocsr()
    Atilde.sort_indices()

    # Amalgamate Atilde for the BSR case, using ones for all strong connections
    if bsr_flag:
        Atilde = Atilde.tobsr(blocksize=(numPDEs, numPDEs))
        nblocks = Atilde.indices.shape[0]
        uone = np.ones((nblocks,))
        Atilde = sparse.csr_matrix((uone, Atilde.indices, Atilde.indptr),
                                   shape=(
                                       int(Atilde.shape[0] / numPDEs),
                                       int(Atilde.shape[1] / numPDEs)))

    # Scale C by the largest magnitude entry in each row
    Atilde = scale_rows_by_largest_entry(Atilde)

    return Atilde def distance_strength_of_connection(A, V, theta=2.0, relative_drop=True):
    """Distance based strength-of-connection.

    Parameters
    ----------
    A : csr_matrix or bsr_matrix
        Square, sparse matrix in CSR or BSR format
    V : array
        Coordinates of the vertices of the graph of A
    relative_drop : bool
        If false, then a connection must be within a distance of theta
        from a point to be strongly connected.
        If true, then the closest connection is always strong, and other points
        must be within theta times the smallest distance to be strong

    Returns
    -------
    C : csr_matrix
        C(i,j) = distance(point_i, point_j)
        Strength of connection matrix where strength values are
        distances, i.e. the smaller the value, the stronger the connection.
        Sparsity pattern of C is copied from A.

    Notes
    -----
    - theta is a drop tolerance that is applied row-wise
    - If a BSR matrix given, then the return matrix is still CSR.  The strength
      is given between super nodes based on the BSR block size.

    Examples
    --------
    >>> from pyamg.gallery import load_example
    >>> from pyamg.strength import distance_strength_of_connection
    >>> data = load_example('airfoil')
    >>> A = data['A'].tocsr()
    >>> S = distance_strength_of_connection(data['A'], data['vertices'])

    """
    # Amalgamate for the supernode case
    if sparse.isspmatrix_bsr(A):
        sn = int(A.shape[0] / A.blocksize[0])
        u = np.ones((A.data.shape[0],))
        A = sparse.csr_matrix((u, A.indices, A.indptr), shape=(sn, sn))

    if not sparse.isspmatrix_csr(A):
        warn('Implicit conversion of A to csr', sparse.SparseEfficiencyWarning)
        A = sparse.csr_matrix(A)

    dim = V.shape[1]

    # Create two arrays for differencing the different coordinates such
    # that C(i,j) = distance(point_i, point_j)
    cols = A.indices
    rows = np.repeat(np.arange(A.shape[0]), A.indptr[1:] - A.indptr[0:-1])

    # Insert difference for each coordinate into C
    C = (V[rows, 0] - V[cols, 0])**2
    for d in range(1, dim):
        C += (V[rows, d] - V[cols, d])**2
    C = np.sqrt(C)
    C[C < 1e-6] = 1e-6

    C = sparse.csr_matrix((C, A.indices.copy(), A.indptr.copy()),
                          shape=A.shape)

    # Apply drop tolerance
    if relative_drop is True:
        if theta != np.inf:
            amg_core.apply_distance_filter(C.shape[0], theta, C.indptr,
                                           C.indices, C.data)
    else:
        amg_core.apply_absolute_distance_filter(C.shape[0], theta, C.indptr,
                                                C.indices, C.data)
    C.eliminate_zeros()

    C = C + sparse.eye(C.shape[0], C.shape[1], format='csr')

    # Standardized strength values require small values be weak and large
    # values be strong.  So, we invert the distances.
    C.data = 1.0 / C.data

    # Scale C by the largest magnitude entry in each row
    C = scale_rows_by_largest_entry(C)

    return C def algebraic_distance(A, alpha=0.5, R=5, k=20, epsilon=2.0, p=2):
    """Algebraic Distance Strength Measure.

    Parameters
    ----------
    A : csr_matrix
        Sparse NxN matrix
    alpha : scalar
        Weight for Jacobi
    R : integer
        Number of random vectors
    k : integer
        Number of relaxation passes
    epsilon : scalar
        Drop tolerance
    p : scalar or inf
        p-norm of the measure

    Returns
    -------
    C : csr_matrix
        Sparse matrix of strength values

    References
    ----------
    .. [SaSaSc] Ilya Safro, Peter Sanders, and Christian Schulz,
        "Advanced Coarsening Schemes for Graph Partitioning"

    Notes
    -----
    No unit testing yet.

    Does not handle BSR matrices yet.

    See [SaSaSc]_ for more details.

    """
    if not sparse.isspmatrix_csr(A):
        A = sparse.csr_matrix(A)

    if alpha < 0:
        raise ValueError('expected alpha>0')

    if R <= 0 or not isinstance(R, int):
        raise ValueError('expected integer R>0')

    if k <= 0 or not isinstance(k, int):
        raise ValueError('expected integer k>0')

    if epsilon < 1:
        raise ValueError('expected epsilon>1.0')

    if p < 1:
        raise ValueError('expected p>1 or equal to numpy.inf')

    def distance(x):
        (rows, cols) = A.nonzero()
        if p != np.inf:
            avg = np.sum(np.abs(x[rows] - x[cols])**p, axis=1) / R
            return (avg)**(1.0 / p)

        return np.abs(x[rows] - x[cols]).max(axis=1)

    return distance_measure_common(A, distance, alpha, R, k, epsilon) def affinity_distance(A, alpha=0.5, R=5, k=20, epsilon=4.0):
    """Affinity Distance Strength Measure.

    Parameters
    ----------
    A : csr_matrix
        Sparse NxN matrix
    alpha : scalar
        Weight for Jacobi
    R : integer
        Number of random vectors
    k : integer
        Number of relaxation passes
    epsilon : scalar
        Drop tolerance

    Returns
    -------
    C : csr_matrix
        Sparse matrix of strength values

    References
    ----------
    .. [LiBr] Oren E. Livne and Achi Brandt, "Lean Algebraic Multigrid
        (LAMG): Fast Graph Laplacian Linear Solver"

    Notes
    -----
    No unit testing yet.

    Does not handle BSR matrices yet.

    See [LiBr]_ for more details.

    """
    if not sparse.isspmatrix_csr(A):
        A = sparse.csr_matrix(A)

    if alpha < 0:
        raise ValueError('expected alpha>0')

    if R <= 0 or not isinstance(R, int):
        raise ValueError('expected integer R>0')

    if k <= 0 or not isinstance(k, int):
        raise ValueError('expected integer k>0')

    if epsilon < 1:
        raise ValueError('expected epsilon>1.0')

    def distance(x):
        (rows, cols) = A.nonzero()
        return 1 - np.sum(x[rows] * x[cols], axis=1)**2 / \
            (np.sum(x[rows]**2, axis=1) * np.sum(x[cols]**2, axis=1))

    return distance_measure_common(A, distance, alpha, R, k, epsilon) def standard_aggregation(C):
    """Compute the sparsity pattern of the tentative prolongator.

    Parameters
    ----------
    C : csr_matrix
        strength of connection matrix

    Returns
    -------
    AggOp : csr_matrix
        aggregation operator which determines the sparsity pattern
        of the tentative prolongator
    Cpts : array
        array of Cpts, i.e., Cpts[i] = root node of aggregate i

    Examples
    --------
    >>> from scipy.sparse import csr_matrix
    >>> from pyamg.gallery import poisson
    >>> from pyamg.aggregation.aggregate import standard_aggregation
    >>> A = poisson((4,), format='csr')   # 1D mesh with 4 vertices
    >>> A.toarray()
    array([[ 2., -1.,  0.,  0.],
           [-1.,  2., -1.,  0.],
           [ 0., -1.,  2., -1.],
           [ 0.,  0., -1.,  2.]])
    >>> standard_aggregation(A)[0].toarray() # two aggregates
    array([[1, 0],
           [1, 0],
           [0, 1],
           [0, 1]], dtype=int8)
    >>> A = csr_matrix([[1,0,0],[0,1,1],[0,1,1]])
    >>> A.toarray()                      # first vertex is isolated
    array([[1, 0, 0],
           [0, 1, 1],
           [0, 1, 1]])
    >>> standard_aggregation(A)[0].toarray() # one aggregate
    array([[0],
           [1],
           [1]], dtype=int8)

    See Also
    --------
    amg_core.standard_aggregation

    """
    if not sparse.isspmatrix_csr(C):
        raise TypeError('expected csr_matrix')

    if C.shape[0] != C.shape[1]:
        raise ValueError('expected square matrix')

    index_type = C.indptr.dtype
    num_rows = C.shape[0]

    Tj = np.empty(num_rows, dtype=index_type)  # stores the aggregate #s
    Cpts = np.empty(num_rows, dtype=index_type)  # stores the Cpts

    fn = amg_core.standard_aggregation

    num_aggregates = fn(num_rows, C.indptr, C.indices, Tj, Cpts)
    Cpts = Cpts[:num_aggregates]

    # no nodes aggregated
    if num_aggregates == 0:
        # return all zero matrix and no Cpts
        return sparse.csr_matrix((num_rows, 1), dtype='int8'),\
            np.array([], dtype=index_type)

    shape = (num_rows, num_aggregates)

    # some nodes not aggregated
    if Tj.min() == -1:
        mask = Tj != -1
        row = np.arange(num_rows, dtype=index_type)[mask]
        col = Tj[mask]
        data = np.ones(len(col), dtype='int8')
        return sparse.coo_matrix((data, (row, col)), shape=shape).tocsr(), Cpts

    # all nodes aggregated
    Tp = np.arange(num_rows+1, dtype=index_type)
    Tx = np.ones(len(Tj), dtype='int8')
    return sparse.csr_matrix((Tx, Tj, Tp), shape=shape), Cpts def naive_aggregation(C):
    """Compute the sparsity pattern of the tentative prolongator.

    Parameters
    ----------
    C : csr_matrix
        strength of connection matrix

    Returns
    -------
    AggOp : csr_matrix
        aggregation operator which determines the sparsity pattern
        of the tentative prolongator
    Cpts : array
        array of Cpts, i.e., Cpts[i] = root node of aggregate i

    Examples
    --------
    >>> from scipy.sparse import csr_matrix
    >>> from pyamg.gallery import poisson
    >>> from pyamg.aggregation.aggregate import naive_aggregation
    >>> A = poisson((4,), format='csr')   # 1D mesh with 4 vertices
    >>> A.toarray()
    array([[ 2., -1.,  0.,  0.],
           [-1.,  2., -1.,  0.],
           [ 0., -1.,  2., -1.],
           [ 0.,  0., -1.,  2.]])
    >>> naive_aggregation(A)[0].toarray() # two aggregates
    array([[1, 0],
           [1, 0],
           [0, 1],
           [0, 1]], dtype=int8)
    >>> A = csr_matrix([[1,0,0],[0,1,1],[0,1,1]])
    >>> A.toarray()                      # first vertex is isolated
    array([[1, 0, 0],
           [0, 1, 1],
           [0, 1, 1]])
    >>> naive_aggregation(A)[0].toarray() # two aggregates
    array([[1, 0],
           [0, 1],
           [0, 1]], dtype=int8)

    See Also
    --------
    amg_core.naive_aggregation

    Notes
    -----
    Differs from standard aggregation.  Each dof is considered.  If it has been
    aggregated, skip over.  Otherwise, put dof and any unaggregated neighbors
    in an aggregate.  Results in possibly much higher complexities than
    standard aggregation.

    """
    if not sparse.isspmatrix_csr(C):
        raise TypeError('expected csr_matrix')

    if C.shape[0] != C.shape[1]:
        raise ValueError('expected square matrix')

    index_type = C.indptr.dtype
    num_rows = C.shape[0]

    Tj = np.empty(num_rows, dtype=index_type)  # stores the aggregate #s
    Cpts = np.empty(num_rows, dtype=index_type)  # stores the Cpts

    fn = amg_core.naive_aggregation

    num_aggregates = fn(num_rows, C.indptr, C.indices, Tj, Cpts)
    Cpts = Cpts[:num_aggregates]
    Tj = Tj - 1

    if num_aggregates == 0:
        # all zero matrix
        return sparse.csr_matrix((num_rows, 1), dtype='int8'), Cpts

    shape = (num_rows, num_aggregates)
    # all nodes aggregated
    Tp = np.arange(num_rows+1, dtype=index_type)
    Tx = np.ones(len(Tj), dtype='int8')
    return sparse.csr_matrix((Tx, Tj, Tp), shape=shape), Cpts def lloyd_aggregation(C, ratio=0.03, distance='unit', maxiter=10):
    """Aggregate nodes using Lloyd Clustering.

    Parameters
    ----------
    C : csr_matrix
        strength of connection matrix
    ratio : scalar
        Fraction of the nodes which will be seeds.
    distance : ['unit','abs','inv',None]
        Distance assigned to each edge of the graph G used in Lloyd clustering

        For each nonzero value C[i,j]:

        =======  ===========================
        'unit'   G[i,j] = 1
        'abs'    G[i,j] = abs(C[i,j])
        'inv'    G[i,j] = 1.0/abs(C[i,j])
        'same'   G[i,j] = C[i,j]
        'sub'    G[i,j] = C[i,j] - min(C)
        =======  ===========================

    maxiter : int
        Maximum number of iterations to perform

    Returns
    -------
    AggOp : csr_matrix
        aggregation operator which determines the sparsity pattern
        of the tentative prolongator
    seeds : array
        array of Cpts, i.e., Cpts[i] = root node of aggregate i

    See Also
    --------
    amg_core.standard_aggregation

    Examples
    --------
    >>> from scipy.sparse import csr_matrix
    >>> from pyamg.gallery import poisson
    >>> from pyamg.aggregation.aggregate import lloyd_aggregation
    >>> A = poisson((4,), format='csr')   # 1D mesh with 4 vertices
    >>> A.toarray()
    array([[ 2., -1.,  0.,  0.],
           [-1.,  2., -1.,  0.],
           [ 0., -1.,  2., -1.],
           [ 0.,  0., -1.,  2.]])
    >>> lloyd_aggregation(A)[0].toarray() # one aggregate
    array([[1],
           [1],
           [1],
           [1]], dtype=int8)
    >>> # more seeding for two aggregates
    >>> Agg = lloyd_aggregation(A,ratio=0.5)[0].toarray()

    """
    if ratio <= 0 or ratio > 1:
        raise ValueError('ratio must be > 0.0 and <= 1.0')

    if not (sparse.isspmatrix_csr(C) or sparse.isspmatrix_csc(C)):
        raise TypeError('expected csr_matrix or csc_matrix')

    if distance == 'unit':
        data = np.ones_like(C.data).astype(float)
    elif distance == 'abs':
        data = abs(C.data)
    elif distance == 'inv':
        data = 1.0/abs(C.data)
    elif distance == 'same':
        data = C.data
    elif distance == 'min':
        data = C.data - C.data.min()
    else:
        raise ValueError(f'Unrecognized value distance={distance}')

    if C.dtype == complex:
        data = np.real(data)

    assert data.min() >= 0

    G = C.__class__((data, C.indices, C.indptr), shape=C.shape)

    num_seeds = int(min(max(ratio * G.shape[0], 1), G.shape[0]))

    _, clusters, seeds = lloyd_cluster(G, num_seeds, maxiter=maxiter)

    row = (clusters >= 0).nonzero()[0]
    col = clusters[row]
    data = np.ones(len(row), dtype='int8')
    AggOp = sparse.coo_matrix((data, (row, col)),
                              shape=(G.shape[0], num_seeds)).tocsr()
    return AggOp, seeds def fit_candidates(AggOp, B, tol=1e-10):
    """Fit near-nullspace candidates to form the tentative prolongator.

    Parameters
    ----------
    AggOp : csr_matrix
        Describes the sparsity pattern of the tentative prolongator.
        Has dimension (#blocks, #aggregates)
    B : array
        The near-nullspace candidates stored in column-wise fashion.
        Has dimension (#blocks * blocksize, #candidates)
    tol : scalar
        Threshold for eliminating local basis functions.
        If after orthogonalization a local basis function Q[:, j] is small,
        i.e. ||Q[:, j]|| < tol, then Q[:, j] is set to zero.

    Returns
    -------
    (Q, R) : (bsr_matrix, array)
        The tentative prolongator Q is a sparse block matrix with dimensions
        (#blocks * blocksize, #aggregates * #candidates) formed by dense blocks
        of size (blocksize, #candidates).  The coarse level candidates are
        stored in R which has dimensions (#aggregates * #candidates,
        #candidates).

    See Also
    --------
    amg_core.fit_candidates

    Notes
    -----
        Assuming that each row of AggOp contains exactly one non-zero entry,
        i.e. all unknowns belong to an aggregate, then Q and R satisfy the
        relationship B = Q*R.  In other words, the near-nullspace candidates
        are represented exactly by the tentative prolongator.

        If AggOp contains rows with no non-zero entries, then the range of the
        tentative prolongator will not include those degrees of freedom. This
        situation is illustrated in the examples below.

    References
    ----------
    .. [1] Vanek, P. and Mandel, J. and Brezina, M.,
       "Algebraic Multigrid by Smoothed Aggregation for
       Second and Fourth Order Elliptic Problems",
       Computing, vol. 56, no. 3, pp. 179--196, 1996.
       http://citeseer.ist.psu.edu/vanek96algebraic.html


    Examples
    --------
    >>> from scipy.sparse import csr_matrix
    >>> from pyamg.aggregation.tentative import fit_candidates
    >>> # four nodes divided into two aggregates
    ... AggOp = csr_matrix( [[1, 0],
    ...                      [1, 0],
    ...                      [0, 1],
    ...                      [0, 1]] )
    >>> # B contains one candidate, the constant vector
    ... B = [[1],
    ...      [1],
    ...      [1],
    ...      [1]]
    >>> Q, R = fit_candidates(AggOp, B)
    >>> Q.toarray()
    array([[0.70710678, 0.        ],
           [0.70710678, 0.        ],
           [0.        , 0.70710678],
           [0.        , 0.70710678]])
    >>> R
    array([[1.41421356],
           [1.41421356]])
    >>> # Two candidates, the constant vector and a linear function
    ... B = [[1, 0],
    ...      [1, 1],
    ...      [1, 2],
    ...      [1, 3]]
    >>> Q, R = fit_candidates(AggOp, B)
    >>> Q.toarray()
    array([[ 0.70710678, -0.70710678,  0.        ,  0.        ],
           [ 0.70710678,  0.70710678,  0.        ,  0.        ],
           [ 0.        ,  0.        ,  0.70710678, -0.70710678],
           [ 0.        ,  0.        ,  0.70710678,  0.70710678]])
    >>> R
    array([[1.41421356, 0.70710678],
           [0.        , 0.70710678],
           [1.41421356, 3.53553391],
           [0.        , 0.70710678]])
    >>> # aggregation excludes the third node
    ... AggOp = csr_matrix( [[1, 0],
    ...                      [1, 0],
    ...                      [0, 0],
    ...                      [0, 1]] )
    >>> B = [[1],
    ...      [1],
    ...      [1],
    ...      [1]]
    >>> Q, R = fit_candidates(AggOp, B)
    >>> Q.toarray()
    array([[0.70710678, 0.        ],
           [0.70710678, 0.        ],
           [0.        , 0.        ],
           [0.        , 1.        ]])
    >>> R
    array([[1.41421356],
           [1.        ]])

    """
    if not isspmatrix_csr(AggOp):
        raise TypeError('expected csr_matrix for argument AggOp')

    B = np.asarray(B)
    if B.dtype not in ['float32', 'float64', 'complex64', 'complex128']:
        B = np.asarray(B, dtype='float64')

    if len(B.shape) != 2:
        raise ValueError('expected 2d array for argument B')

    if B.shape[0] % AggOp.shape[0] != 0:
        raise ValueError(f'Dimensions of AggOp {AggOp.shape} and B {B.shape} '
                         'are incompatible')

    N_fine, N_coarse = AggOp.shape

    K1 = int(B.shape[0] / N_fine)  # dof per supernode (e.g. 3 for 3d vectors)
    K2 = B.shape[1]                # candidates

    # the first two dimensions of R and Qx are collapsed later
    R = np.empty((N_coarse, K2, K2), dtype=B.dtype)    # coarse candidates
    Qx = np.empty((AggOp.nnz, K1, K2), dtype=B.dtype)  # BSR data array

    AggOp_csc = AggOp.tocsc()

    fn = amg_core.fit_candidates
    fn(N_fine, N_coarse, K1, K2,
       AggOp_csc.indptr, AggOp_csc.indices, Qx.ravel(),
       B.ravel(), R.ravel(), tol)

    Q = bsr_matrix((Qx.swapaxes(1, 2).copy(), AggOp_csc.indices,
                    AggOp_csc.indptr), shape=(K2*N_coarse, K1*N_fine))
    Q = Q.T.tobsr()
    R = R.reshape(-1, K2)

    return Q, R def jacobi_prolongation_smoother(S, T, C, B, omega=4.0/3.0, degree=1,
                                 filter_entries=False, weighting='diagonal'):
    """Jacobi prolongation smoother.

    Parameters
    ----------
    S : csr_matrix, bsr_matrix
        Sparse NxN matrix used for smoothing.  Typically, A.
    T : csr_matrix, bsr_matrix
        Tentative prolongator
    C : csr_matrix, bsr_matrix
        Strength-of-connection matrix
    B : array
        Near nullspace modes for the coarse grid such that T*B
        exactly reproduces the fine grid near nullspace modes
    omega : scalar
        Damping parameter
    filter_entries : boolean
        If true, filter S before smoothing T.  This option can greatly control
        complexity.
    weighting : string
        'block', 'diagonal' or 'local' weighting for constructing the Jacobi D
        'local' Uses a local row-wise weight based on the Gershgorin estimate.
        Avoids any potential under-damping due to inaccurate spectral radius
        estimates.
        'block' uses a block diagonal inverse of A if A is BSR
        'diagonal' uses classic Jacobi with D = diagonal(A)

    Returns
    -------
    P : csr_matrix, bsr_matrix
        Smoothed (final) prolongator defined by P = (I - omega/rho(K) K) * T
        where K = diag(S)^-1 * S and rho(K) is an approximation to the
        spectral radius of K.

    Notes
    -----
    If weighting is not 'local', then results using Jacobi prolongation
    smoother are not precisely reproducible due to a random initial guess used
    for the spectral radius approximation.  For precise reproducibility,
    set numpy.random.seed(..) to the same value before each test.

    Examples
    --------
    >>> from pyamg.aggregation import jacobi_prolongation_smoother
    >>> from pyamg.gallery import poisson
    >>> from scipy.sparse import coo_matrix
    >>> import numpy as np
    >>> data = np.ones((6,))
    >>> row = np.arange(0,6)
    >>> col = np.kron([0,1],np.ones((3,)))
    >>> T = coo_matrix((data,(row,col)),shape=(6,2)).tocsr()
    >>> T.toarray()
    array([[1., 0.],
           [1., 0.],
           [1., 0.],
           [0., 1.],
           [0., 1.],
           [0., 1.]])
    >>> A = poisson((6,),format='csr')
    >>> P = jacobi_prolongation_smoother(A,T,A,np.ones((2,1)))
    >>> P.toarray()
    array([[0.64930164, 0.        ],
           [1.        , 0.        ],
           [0.64930164, 0.35069836],
           [0.35069836, 0.64930164],
           [0.        , 1.        ],
           [0.        , 0.64930164]])

    """
    # preprocess weighting
    if weighting == 'block':
        if sparse.isspmatrix_csr(S):
            weighting = 'diagonal'
        elif sparse.isspmatrix_bsr(S):
            if S.blocksize[0] == 1:
                weighting = 'diagonal'

    if filter_entries:
        # Implement filtered prolongation smoothing for the general case by
        # utilizing satisfy constraints

        if sparse.isspmatrix_bsr(S):
            numPDEs = S.blocksize[0]
        else:
            numPDEs = 1

        # Create a filtered S with entries dropped that aren't in C
        C = unamal(C, numPDEs, numPDEs)
        S = S.multiply(C)
        S.eliminate_zeros()

    if weighting == 'diagonal':
        # Use diagonal of S
        D_inv = get_diagonal(S, inv=True)
        D_inv_S = scale_rows(S, D_inv, copy=True)
        D_inv_S = (omega/approximate_spectral_radius(D_inv_S))*D_inv_S
    elif weighting == 'block':
        # Use block diagonal of S
        D_inv = get_block_diag(S, blocksize=S.blocksize[0], inv_flag=True)
        D_inv = sparse.bsr_matrix((D_inv, np.arange(D_inv.shape[0]),
                                   np.arange(D_inv.shape[0]+1)),
                                  shape=S.shape)
        D_inv_S = D_inv*S
        D_inv_S = (omega/approximate_spectral_radius(D_inv_S))*D_inv_S
    elif weighting == 'local':
        # Use the Gershgorin estimate as each row's weight, instead of a global
        # spectral radius estimate
        D = np.abs(S)*np.ones((S.shape[0], 1), dtype=S.dtype)
        D_inv = np.zeros_like(D)
        D_inv[D != 0] = 1.0 / np.abs(D[D != 0])

        D_inv_S = scale_rows(S, D_inv, copy=True)
        D_inv_S = omega*D_inv_S
    else:
        raise ValueError('Incorrect weighting option')

    if filter_entries:
        # Carry out Jacobi, but after calculating the prolongator update, U,
        # apply satisfy constraints so that U*B = 0
        P = T
        for _ in range(degree):
            U = (D_inv_S*P).tobsr(blocksize=P.blocksize)

            # Enforce U*B = 0 (1) Construct array of inv(Bi'Bi), where Bi is B
            # restricted to row i's sparsity pattern in pattern. This
            # array is used multiple times in satisfy_constraints(...).
            BtBinv = compute_BtBinv(B, U)

            # (2) Apply satisfy constraints
            satisfy_constraints(U, B, BtBinv)

            # Update P
            P = P - U

    else:
        # Carry out Jacobi as normal
        P = T
        for _ in range(degree):
            P = P - (D_inv_S*P)

    return P def richardson_prolongation_smoother(S, T, omega=4.0/3.0, degree=1):
    """Richardson prolongation smoother.

    Parameters
    ----------
    S : csr_matrix, bsr_matrix
        Sparse NxN matrix used for smoothing.  Typically, A or the
        "filtered matrix" obtained from A by lumping weak connections
        onto the diagonal of A.
    T : csr_matrix, bsr_matrix
        Tentative prolongator
    omega : scalar
        Damping parameter

    Returns
    -------
    P : csr_matrix, bsr_matrix
        Smoothed (final) prolongator defined by P = (I - omega/rho(S) S) * T
        where rho(S) is an approximation to the spectral radius of S.

    Notes
    -----
    Results using Richardson prolongation smoother are not precisely
    reproducible due to a random initial guess used for the spectral radius
    approximation.  For precise reproducibility, set numpy.random.seed(..) to
    the same value before each test.


    Examples
    --------
    >>> from pyamg.aggregation import richardson_prolongation_smoother
    >>> from pyamg.gallery import poisson
    >>> from scipy.sparse import coo_matrix
    >>> import numpy as np
    >>> data = np.ones((6,))
    >>> row = np.arange(0,6)
    >>> col = np.kron([0,1],np.ones((3,)))
    >>> T = coo_matrix((data,(row,col)),shape=(6,2)).tocsr()
    >>> T.toarray()
    array([[1., 0.],
           [1., 0.],
           [1., 0.],
           [0., 1.],
           [0., 1.],
           [0., 1.]])
    >>> A = poisson((6,),format='csr')
    >>> P = richardson_prolongation_smoother(A,T)
    >>> P.toarray()
    array([[0.64930164, 0.        ],
           [1.        , 0.        ],
           [0.64930164, 0.35069836],
           [0.35069836, 0.64930164],
           [0.        , 1.        ],
           [0.        , 0.64930164]])

    """
    weight = omega/approximate_spectral_radius(S)

    P = T
    for _ in range(degree):
        P = P - weight*(S*P)

    return P def energy_prolongation_smoother(A, T, Atilde, B, Bf, Cpt_params,
                                 krylov='cg', maxiter=4, tol=1e-8,
                                 degree=1, weighting='local',
                                 prefilter=None, postfilter=None):
    """Minimize the energy of the coarse basis functions (columns of T).

    Both root-node and non-root-node style prolongation smoothing is available,
    see Cpt_params description below.

    Parameters
    ----------
    A : csr_matrix, bsr_matrix
        Sparse NxN matrix
    T : bsr_matrix
        Tentative prolongator, a NxM sparse matrix (M < N)
    Atilde : csr_matrix
        Strength of connection matrix
    B : array
        Near-nullspace modes for coarse grid.  Has shape (M,k) where
        k is the number of coarse candidate vectors.
    Bf : array
        Near-nullspace modes for fine grid.  Has shape (N,k) where
        k is the number of coarse candidate vectors.
    Cpt_params : tuple
        Tuple of the form (bool, dict).  If the Cpt_params[0] = False, then the
        standard SA prolongation smoothing is carried out.  If True, then
        root-node style prolongation smoothing is carried out.  The dict must
        be a dictionary of parameters containing, (1) for P_I, P_I.T is the
        injection matrix for the Cpts, (2) I_F is an identity matrix for only the
        F-points (i.e. I, but with zero rows and columns for C-points) and I_C is
        the C-point analogue to I_F.  See Notes below for more information.
    krylov : string
        'cg' for SPD systems.  Solve A T = 0 in a constraint space with CG
        'cgnr' for nonsymmetric and/or indefinite systems.
        Solve A T = 0 in a constraint space with CGNR
        'gmres' for nonsymmetric and/or indefinite systems.
        Solve A T = 0 in a constraint space with GMRES
    maxiter : integer
        Number of energy minimization steps to apply to the prolongator
    tol : scalar
        Minimization tolerance
    degree : int
        Generate sparsity pattern for P based on (Atilde^degree T)
    weighting : string
        'block', 'diagonal' or 'local' construction of the diagonal preconditioning
        'local' Uses a local row-wise weight based on the Gershgorin estimate.
        Avoids any potential under-damping due to inaccurate spectral
        radius estimates.
        'block' Uses a block diagonal inverse of A if A is BSR.
        'diagonal' Uses the inverse of the diagonal of A
    prefilter : dictionary
        Filter elements by row in sparsity pattern for P to reduce operator and
        setup complexity. If None or an empty dictionary, then no dropping in P
        is done.  If postfilter has key 'k', then the largest 'k' entries  are
        kept in each row.  If postfilter has key 'theta', all entries such that
        :math:`P[i,j] < kwargs['theta']*max(abs(P[i,:]))`
        are dropped.  If postfilter['k'] and postfiler['theta'] are present,
        then they are used with the union of their patterns.
    postfilter : dictionary
        Filters elements by row in smoothed P to reduce operator complexity.
        Only supported if using the rootnode_solver. If None or an empty
        dictionary, no dropping in P is done. If postfilter has key 'k',
        then the largest 'k' entries  are kept in each row.  If postfilter
        has key 'theta', all entries such that
        :math::`P[i,j] < kwargs['theta']*max(abs(P[i,:]))`
        are dropped.  If postfilter['k'] and postfiler['theta'] are present,
        then they are used with the union of their patterns.

    Returns
    -------
    T : bsr_matrix
        Smoothed prolongator

    Notes
    -----
    Only 'diagonal' weighting is supported for the CGNR method, because
    we are working with A^* A and not A.

    When Cpt_params[0] == True, root-node style prolongation smoothing is used
    to minimize the energy of columns of T.  Essentially, an identity block is
    maintained in T, corresponding to injection from the coarse-grid to the
    fine-grid root-nodes.  See [2011OlScTu]_ for more details, and see
    util.utils.get_Cpt_params for the helper function to generate Cpt_params.

    If Cpt_params[0] == False, the energy of columns of T are still
    minimized, but without maintaining the identity block.

    See [1999cMaBrVa]_ for more details on smoothed aggregation.

    Examples
    --------
    >>> from pyamg.aggregation import energy_prolongation_smoother
    >>> from pyamg.gallery import poisson
    >>> from scipy.sparse import coo_matrix
    >>> import numpy as np
    >>> data = np.ones((6,))
    >>> row = np.arange(0,6)
    >>> col = np.kron([0,1],np.ones((3,)))
    >>> T = coo_matrix((data,(row,col)),shape=(6,2)).tocsr()
    >>> print(T.toarray())
    [[1. 0.]
     [1. 0.]
     [1. 0.]
     [0. 1.]
     [0. 1.]
     [0. 1.]]
    >>> A = poisson((6,),format='csr')
    >>> B = np.ones((2,1),dtype=float)
    >>> P = energy_prolongation_smoother(A,T,A,B, None, (False,{}))
    >>> print(P.toarray())
    [[1.         0.        ]
     [1.         0.        ]
     [0.66666667 0.33333333]
     [0.33333333 0.66666667]
     [0.         1.        ]
     [0.         1.        ]]

    References
    ----------
    .. [1999cMaBrVa] Jan Mandel, Marian Brezina, and Petr Vanek
       "Energy Optimization of Algebraic Multigrid Bases"
       Computing 62, 205-228, 1999
       http://dx.doi.org/10.1007/s006070050022

    .. [2011OlScTu] Olson, L. and Schroder, J. and Tuminaro, R.,
       "A general interpolation strategy for algebraic
       multigrid using energy minimization", SIAM Journal
       on Scientific Computing (SISC), vol. 33, pp.
       966--991, 2011.

    """
    # Test Inputs
    if maxiter < 0:
        raise ValueError('maxiter must be > 0')
    if tol > 1:
        raise ValueError('tol must be <= 1')

    if sparse.isspmatrix_csr(A):
        A = A.tobsr(blocksize=(1, 1), copy=False)
    elif sparse.isspmatrix_bsr(A):
        pass
    else:
        raise TypeError('A must be csr_matrix or bsr_matrix')

    if sparse.isspmatrix_csr(T):
        T = T.tobsr(blocksize=(1, 1), copy=False)
    elif sparse.isspmatrix_bsr(T):
        pass
    else:
        raise TypeError('T must be csr_matrix or bsr_matrix')

    if T.blocksize[0] != A.blocksize[0]:
        raise ValueError('T row-blocksize should be the same as A blocksize')

    if B.shape[0] != T.shape[1]:
        raise ValueError('B is the candidates for the coarse grid. '
                         '  num_rows(b) = num_cols(T)')

    if min(T.nnz, A.nnz) == 0:
        return T

    if not sparse.isspmatrix_csr(Atilde):
        raise TypeError('Atilde must be csr_matrix')

    if prefilter is None:
        prefilter = {}

    if postfilter is None:
        postfilter = {}

    if ('theta' in prefilter) and (prefilter['theta'] == 0):
        prefilter.pop('theta', None)

    if ('theta' in postfilter) and (postfilter['theta'] == 0):
        postfilter.pop('theta', None)

    # Prepocess Atilde, the strength matrix
    if Atilde is None:
        Atilde = sparse.csr_matrix((np.ones(len(A.indices)),
                                    A.indices.copy(), A.indptr.copy()),
                                   shape=(A.shape[0]/A.blocksize[0],
                                          A.shape[1]/A.blocksize[1]))

    # If Atilde has no nonzeros, then return T
    if min(T.nnz, A.nnz) == 0:
        return T

    # Expand allowed sparsity pattern for P through multiplication by Atilde
    if degree > 0:

        # Construct pattern by multiplying with Atilde
        T.sort_indices()
        shape = (int(T.shape[0]/T.blocksize[0]),
                 int(T.shape[1]/T.blocksize[1]))
        pattern = sparse.csr_matrix((np.ones(T.indices.shape), T.indices, T.indptr),
                                    shape=shape)

        AtildeCopy = Atilde.copy()
        for _ in range(degree):
            pattern = AtildeCopy * pattern

        # Optional filtering of sparsity pattern before smoothing
        if 'theta' in prefilter and 'k' in prefilter:
            pattern_theta = filter_matrix_rows(pattern, prefilter['theta'])
            pattern = truncate_rows(pattern, prefilter['k'])
            # Union two sparsity patterns
            pattern += pattern_theta
        elif 'k' in prefilter:
            pattern = truncate_rows(pattern, prefilter['k'])
        elif 'theta' in prefilter:
            pattern = filter_matrix_rows(pattern, prefilter['theta'])
        elif len(prefilter) > 0:
            raise ValueError('Unrecognized prefilter option')

        # unamal returns a BSR matrix with 1's in the nonzero locations
        pattern = unamal(pattern, T.blocksize[0], T.blocksize[1])
        pattern.sort_indices()

    else:
        # If degree is 0, just copy T for the sparsity pattern
        pattern = T.copy()
        if 'theta' in prefilter and 'k' in prefilter:
            pattern_theta = filter_matrix_rows(pattern, prefilter['theta'])
            pattern = truncate_rows(pattern, prefilter['k'])
            # Union two sparsity patterns
            pattern += pattern_theta
        elif 'k' in prefilter:
            pattern = truncate_rows(pattern, prefilter['k'])
        elif 'theta' in prefilter:
            pattern = filter_matrix_rows(pattern, prefilter['theta'])
        elif len(prefilter) > 0:
            raise ValueError('Unrecognized prefilter option')

        pattern.data[:] = 1.0
        pattern.sort_indices()

    # If using root nodes, enforce identity at C-points
    if Cpt_params[0]:
        pattern = Cpt_params[1]['I_F'] * pattern
        pattern = Cpt_params[1]['P_I'] + pattern

    # Construct array of inv(Bi'Bi), where Bi is B restricted to row i's
    # sparsity pattern in pattern. This array is used multiple times
    # in satisfy_constraints(...).
    BtBinv = compute_BtBinv(B, pattern)

    # If using root nodes and B has more columns that A's blocksize, then
    # T must be updated so that T*B = Bfine.  Note, if this is a 'secondpass'
    # after dropping entries in P, then we must re-enforce the constraints
    if ((Cpt_params[0] and (B.shape[1] > A.blocksize[0]))
       or ('secondpass' in postfilter)):
        T = filter_operator(T, pattern, B, Bf, BtBinv)
        # Ensure identity at C-pts
        if Cpt_params[0]:
            T = Cpt_params[1]['I_F']*T + Cpt_params[1]['P_I']

    # Iteratively minimize the energy of T subject to the constraints of
    # pattern and maintaining T's effect on B, i.e. T*B =
    # (T+Update)*B, i.e. Update*B = 0
    if krylov == 'cg':
        T = cg_prolongation_smoothing(A, T, B, BtBinv, pattern,
                                      maxiter, tol, weighting, Cpt_params)
    elif krylov == 'cgnr':
        T = cgnr_prolongation_smoothing(A, T, B, BtBinv, pattern,
                                        maxiter, tol, weighting, Cpt_params)
    elif krylov == 'gmres':
        T = gmres_prolongation_smoothing(A, T, B, BtBinv, pattern,
                                         maxiter, tol, weighting, Cpt_params)

    T.eliminate_zeros()

    # Filter entries in P, only in the rootnode case,
    # i.e., Cpt_params[0] == True
    if ((len(postfilter) == 0)
       or ('secondpass' in postfilter)
       or (Cpt_params[0] is False)):
        return T

    if 'theta' in postfilter and 'k' in postfilter:
        T_theta = filter_matrix_rows(T, postfilter['theta'])
        T_k = truncate_rows(T, postfilter['k'])

        # Union two sparsity patterns
        T_theta.data[:] = 1.0
        T_k.data[:] = 1.0
        T_filter = T_theta + T_k
        T_filter.data[:] = 1.0
        T_filter = T.multiply(T_filter)

    elif 'k' in postfilter:
        T_filter = truncate_rows(T, postfilter['k'])
    elif 'theta' in postfilter:
        T_filter = filter_matrix_rows(T, postfilter['theta'])
    else:
        raise ValueError('Unrecognized postfilter option')

    # Re-smooth T_filter and re-fit the modes B into the span.
    # Note, we set 'secondpass', because this is the second
    # filtering pass
    T = energy_prolongation_smoother(A, T_filter,
                                     Atilde, B, Bf, Cpt_params,
                                     krylov=krylov, maxiter=1,
                                     tol=1e-8, degree=0,
                                     weighting=weighting,
                                     prefilter={},
                                     postfilter={'secondpass': True})

    return T def relaxation_as_linear_operator(method, A, b):
    """Create a linear operator that applies a relaxation method to a right-hand-side.

    Parameters
    ----------
    methods : {tuple or string}
        Relaxation descriptor: Each tuple must be of the form ('method','opts')
        where 'method' is the name of a supported smoother, e.g., gauss_seidel,
        and 'opts' a dict of keyword arguments to the smoother, e.g., opts =
        {'sweep':symmetric}.  If string, must be that of a supported smoother,
        e.g., gauss_seidel.

    Returns
    -------
    linear operator that applies the relaxation method to a vector for a
    fixed right-hand-side, b.

    Notes
    -----
    This method is primarily used to improve B during the aggregation setup
    phase.  Here b = 0, and each relaxation call can improve the quality of B,
    especially near the boundaries.

    Examples
    --------
    >>> from pyamg.gallery import poisson
    >>> from pyamg.relaxation.utils import relaxation_as_linear_operator
    >>> import numpy as np
    >>> A = poisson((100,100), format='csr')           # matrix
    >>> B = np.ones((A.shape[0],1))                 # Candidate vector
    >>> b = np.zeros((A.shape[0]))                  # RHS
    >>> relax = relaxation_as_linear_operator('gauss_seidel', A, b)
    >>> B = relax*B

    """

    def unpack_arg(v):
        if isinstance(v, tuple):
            return v[0], v[1]
        return v, {}

    # setup variables
    accepted_methods = ['gauss_seidel', 'block_gauss_seidel', 'sor',
                        'gauss_seidel_ne', 'gauss_seidel_nr', 'jacobi',
                        'block_jacobi', 'richardson', 'schwarz',
                        'strength_based_schwarz', 'jacobi_ne']

    b = np.array(b, dtype=A.dtype)
    fn, kwargs = unpack_arg(method)
    lvl = MultilevelSolver.Level()
    lvl.A = A

    # Retrieve setup call from relaxation.smoothing for this relaxation method
    if not accepted_methods.__contains__(fn):
        raise NameError(f'invalid relaxation method: {fn}')
    try:
        setup_smoother = getattr(relaxation.smoothing, 'setup_' + fn)
    except NameError as e:
        raise NameError(f'invalid presmoother method: {fn}') from e

    # Get relaxation routine that takes only (A, x, b) as parameters
    relax = setup_smoother(lvl, **kwargs)

    # Define matvec
    def matvec(x):
        xcopy = x.copy()
        relax(A, xcopy, b)
        return xcopy

    return LinearOperator(A.shape, matvec, dtype=A.dtype)
from warnings import warn
from scipy.sparse import csr_matrix, isspmatrix_csr, isspmatrix_bsr,\
    SparseEfficiencyWarning
from pyamg.multilevel import MultilevelSolver
from pyamg.relaxation.smoothing import change_smoothers
from pyamg.util.utils import eliminate_diag_dom_nodes, get_blocksize,\
    levelize_strength_or_aggregation, levelize_smooth_or_improve_candidates
from pyamg.strength import classical_strength_of_connection,\
    symmetric_strength_of_connection, evolution_strength_of_connection,\
    energy_based_strength_of_connection, distance_strength_of_connection,\
    algebraic_distance, affinity_distance
from .aggregate import standard_aggregation, naive_aggregation,\
    lloyd_aggregation
from .tentative import fit_candidates
from .smooth import jacobi_prolongation_smoother,\
    richardson_prolongation_smoother, energy_prolongation_smoother
from ..relaxation.utils import relaxation_as_linear_operator
import numpy as np
    >>> x,info = cg(A, b, tol=1e-8, maxiter=30, M=M)   # solve with CG

    References
    ----------
    .. [1996VaMaBr] Vanek, P. and Mandel, J. and Brezina, M.,
       "Algebraic Multigrid by Smoothed Aggregation for
       Second and Fourth Order Elliptic Problems",
       Computing, vol. 56, no. 3, pp. 179--196, 1996.
       http://citeseer.ist.psu.edu/vanek96algebraic.html

    """
    if not (isspmatrix_csr(A) or isspmatrix_bsr(A)):
        try:
            A = csr_matrix(A)
            warn('Implicit conversion of A to CSR', SparseEfficiencyWarning)
        except BaseException as e:
            raise TypeError('Argument A must have type csr_matrix or bsr_matrix, '
                            'or be convertible to csr_matrix') from e

    A = A.asfptype()

    if symmetry not in ('symmetric', 'hermitian', 'nonsymmetric'):
        raise ValueError('Expected "symmetric", "nonsymmetric" or "hermitian" '
                         'for the symmetry parameter ')
    A.symmetry = symmetry

    if A.shape[0] != A.shape[1]:
        raise ValueError('expected square matrix')

    # Right near nullspace candidates use constant for each variable as default
    if B is None:
        B = np.kron(np.ones((int(A.shape[0]/get_blocksize(A)), 1), dtype=A.dtype),
                    np.eye(get_blocksize(A), dtype=A.dtype))
    else:
        B = np.asarray(B, dtype=A.dtype)
        if len(B.shape) == 1:
            B = B.reshape(-1, 1)
        if B.shape[0] != A.shape[0]:
            raise ValueError('The shape of near null-space modes B is incorrect')
        if B.shape[1] < get_blocksize(A):
            warn('Having less target vectors, B.shape[1], than blocksize of A '
                 'can degrade convergence factors.')

    # Left near nullspace candidates
    if A.symmetry == 'nonsymmetric':
        if BH is None:
            BH = B.copy()
        else:
            BH = np.asarray(BH, dtype=A.dtype)
            if len(BH.shape) == 1:
                BH = BH.reshape(-1, 1)
            if BH.shape[1] != B.shape[1]:
                raise ValueError('The number of left and right near null-space '
                                 'modes B and BH must be equal')
            if BH.shape[0] != A.shape[0]:
                raise ValueError('The shape of near null-space modes B is incorrect')

    # Levelize the user parameters, so that they become lists describing the
    # desired user option on each level.
    max_levels, max_coarse, strength =\
        levelize_strength_or_aggregation(strength, max_levels, max_coarse)
    max_levels, max_coarse, aggregate =\
        levelize_strength_or_aggregation(aggregate, max_levels, max_coarse)
    improve_candidates =\
        levelize_smooth_or_improve_candidates(improve_candidates, max_levels)
    smooth = levelize_smooth_or_improve_candidates(smooth, max_levels)

    # Construct multilevel structure
    levels = []
    levels.append(MultilevelSolver.Level())
    levels[-1].A = A          # matrix

    # Append near nullspace candidates
    levels[-1].B = B          # right candidates
    if A.symmetry == 'nonsymmetric':
        levels[-1].BH = BH    # left candidates

    while len(levels) < max_levels and\
            int(levels[-1].A.shape[0]/get_blocksize(levels[-1].A)) > max_coarse:
        _extend_hierarchy(levels, strength, aggregate, smooth,
                          improve_candidates, diagonal_dominance, keep)

    ml = MultilevelSolver(levels, **kwargs)
    change_smoothers(ml, presmoother, postsmoother)
    return ml


def _extend_hierarchy(levels, strength, aggregate, smooth, improve_candidates,
                      diagonal_dominance=False, keep=True):
    """Extend the multigrid hierarchy.

    Service routine to implement the strength of connection, aggregation,
    tentative prolongation construction, and prolongation smoothing.  Called by
    smoothed_aggregation_solver.

    """
    def unpack_arg(v):
        if isinstance(v, tuple):
            return v[0], v[1]
        return v, {}

    A = levels[-1].A
    B = levels[-1].B
    if A.symmetry == 'nonsymmetric':
        AH = A.H.asformat(A.format)
        BH = levels[-1].BH

    # Compute the strength-of-connection matrix C, where larger
    # C[i,j] denote stronger couplings between i and j.
    fn, kwargs = unpack_arg(strength[len(levels)-1])
    if fn == 'symmetric':
        C = symmetric_strength_of_connection(A, **kwargs)
    elif fn == 'classical':
        C = classical_strength_of_connection(A, **kwargs)
    elif fn == 'distance':
        C = distance_strength_of_connection(A, **kwargs)
    elif fn in ('ode', 'evolution'):
        if 'B' in kwargs:
            C = evolution_strength_of_connection(A, **kwargs)
        else:

            C = evolution_strength_of_connection(A, B, **kwargs)