import pyamg
import numpy
from scipy.io import loadmat
import matplotlib.pyplot as plt
%matplotlib notebook
#data = loadmat('square.mat')
data = pyamg.gallery.load_example('airfoil')
A = data['A'].tocsr() # matrix
V = data['vertices'][:A.shape[0]] # vertices of each variable
E = numpy.vstack((A.tocoo().row,A.tocoo().col)).T # edges of the matrix graph
mls = pyamg.ruge_stuben_solver(A, max_levels=2, max_coarse=1, CF='RS',keep=True)
print(mls)
# The CF splitting, 1 == C-node and 0 == F-node
splitting = mls.levels[0].splitting
C_nodes = splitting == 1
F_nodes = splitting == 0
plt.figure(figsize=(6,6))
plt.axis('equal')
for e in E:
plt.plot(V[e,0], V[e,1], 'k-', zorder=1)
plt.scatter(V[:,0][C_nodes], V[:,1][C_nodes], c='r', s=100.0, zorder=2) #plot C-nodes in red
plt.scatter(V[:,0][F_nodes], V[:,1][F_nodes], c='b', s=100.0, zorder=2) #plot F-nodes in blue