How to use the pyamg.gallery.poisson function in pyamg

In this example Smoothed Aggregation (SA) is used to precondition
the LOBPCG eigensolver on a two-dimensional Poisson problem with
Dirichlet boundary conditions.
"""

import scipy
from scipy.sparse.linalg import lobpcg
import matplotlib.pyplot as plt

from pyamg import smoothed_aggregation_solver
from pyamg.gallery import poisson

N = 100
K = 9
A = poisson((N,N), format='csr')

# create the AMG hierarchy
ml = smoothed_aggregation_solver(A)

# initial approximation to the K eigenvectors
X = scipy.rand(A.shape[0], K) 

# preconditioner based on ml
M = ml.aspreconditioner()

# compute eigenvalues and eigenvectors with LOBPCG
W,V = lobpcg(A, X, M=M, tol=1e-8, largest=False)


#plot the eigenvectors
plt.figure(figsize=(9,9))

>>> ml = pyamg.aggregation.smoothed_aggregation_solver(A, max_levels=1000)
        >>> M = ml.aspreconditioner(cycle='W') # pre-conditioner
        >>> x, info = cg(A, b, tol=1e-12, M=M)

        See Also
        --------
        pyamg_index : convert grid indices into the sparse matrix indices for ``A``
        pyamg_fix : fixes stencil for indices and fixes the source for the RHS matrix (uses `pyamg_source`)
        pyamg_source : fix the RHS matrix ``b`` to a constant value
        pyamg_boundary_condition : setup the sparse matrix ``A`` to given boundary conditions (called in this routine)
        """
        from pyamg.gallery import poisson
        if dtype is None:
            dtype = self.dtype
        # Initially create the CSR matrix
        A = poisson(self.shape, dtype=dtype, format='csr')
        b = np.zeros(A.shape[0], dtype=A.dtype)

        # Now apply the boundary conditions
        self.pyamg_boundary_condition(A, b)
        return A, b