import pyamg
import numpy as np
from scipy.io import loadmat
import matplotlib.pyplot as plt
%matplotlib inline
n=100
sten = pyamg.gallery.diffusion.diffusion_stencil_2d(epsilon=0.0001, theta=np.pi/4, type='FE')
A = pyamg.gallery.stencil_grid(sten, (n,n), format='csr')
x1d = np.linspace(0, 1, n)
X, Y = np.meshgrid(x1d, x1d)
ml = pyamg.ruge_stuben_solver(A, max_coarse=200, strength='evolution', keep=True)
ml
splitting0 = ml.levels[0].splitting
plt.figure(figsize=(12,12))
X = X.ravel()
Y = Y.ravel()
I = np.where(splitting0==0)[0]
plt.scatter(X[I], Y[I], color='b', s=8, marker='s', lw=0)
I = np.where(splitting0==1)[0]
plt.scatter(X[I], Y[I], color='b', s=16, marker='s', lw=0)
splitting1 = ml.levels[1].splitting
I1 = np.where(splitting1==1)[0]
plt.scatter(X[I[I1]], Y[I[I1]], color='r', s=22, marker='s', lw=0)
ml.cycle_complexity(cycle='AMLI')