"""
Compute eigenvectors and eigenvalues using a preconditioned eigensolver
=======================================================================

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

import numpy as np
import scipy as sp
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 = np.random.random((A.shape[0], K))

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

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


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

for i in range(K):
    plt.subplot(3, 3, i + 1)
    plt.title(f"Eigenvector {i}")
    plt.pcolor(V[:, i].reshape(N, N))
    plt.axis("equal")
    plt.axis("off")
plt.show()