demo.py

# Illustrates the selection of aggregates in AMG based on smoothed aggregation

import numpy
from pyamg import rootnode_solver
from pyamg.gallery import load_example

data = load_example('unit_square')
#data = load_example('recirc_flow')
#data = load_example('airfoil')

XA = 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

from GeneralTriangulation import GeneralTriangulation 
import scipy.sparse as sp

#import pdb; pdb.set_trace()
V = numpy.random.rand(1000,2)
tri = GeneralTriangulation(V[:,0],V[:,1])
tri.compute_area()

# build A
A = sp.csr_matrix((V.shape[0],V.shape[0]))
import pdb; pdb.set_trace()
import numpy as np
for anum in np.arange(V.shape[0]):
    adj = np.where(np.sum(anum == tri.triangles,axis=1))
    for atri in adj:
        A[anum,tri.triangles[atri,:]] += tri.area[atri,np.newaxis]
A = -A
for anum in np.arange(V.shape[0]):
    A[anum,anum] *= -A[anum,anum]

#A = (sp.coo_matrix((numpy.ones(E.shape[0]),(E[:,0],E[:,1])),shape=(E.max()+1,E.max()+1))+sp.diags(1,0,shape=(E.max()+1,E.max()+1))).tocsr()
#A = (sp.coo_matrix((numpy.ones(E.shape[0]),(E[:,0],E[:,1])),shape=(E.max()+1,E.max()+1))).tocsr()
E = numpy.vstack((A.tocoo().row,A.tocoo().col)).T  # edges of the matrix graph

# Use Root-Node Solver
mls = rootnode_solver(A, max_levels=2, max_coarse=1, keep=True)

#import pdb; pdb.set_trace()
# AggOp[i,j] is 1 iff node i belongs to aggregate j
AggOp = mls.levels[0].AggOp

# determine which edges lie entirely inside an aggregate
# AggOp.indices[n] is the aggregate to which vertex n belongs
inner_edges = AggOp.indices[E[:,0]] == AggOp.indices[E[:,1]]  
outer_edges = -inner_edges

# Grab the root-nodes (i.e., the C/F splitting)
Cpts = mls.levels[0].Cpts
Fpts = mls.levels[0].Fpts

from draw import lineplot
from pylab import figure, axis, scatter, show, title

##
# Plot the aggregation
figure(figsize=(6,6))
title('Finest-Level Aggregation\nC-pts in Red, F-pts in Blue')
axis('equal')
lineplot(V, E[inner_edges], linewidths=3.0)
lineplot(V, E[outer_edges], linewidths=0.2)
scatter(V[:,0][Fpts], V[:,1][Fpts], c='b', s=100.0)  #plot F-nodes in blue
scatter(V[:,0][Cpts], V[:,1][Cpts], c='r', s=220.0)  #plot C-nodes in red

##
# Plot the C/F splitting
figure(figsize=(6,6))
title('Finest-Level C/F splitting\nC-pts in Red, F-pts in Blue')
axis('equal')
lineplot(V, E)
scatter(V[:,0][Cpts], V[:,1][Cpts], c='r', s=100.0)  #plot C-nodes in red
scatter(V[:,0][Fpts], V[:,1][Fpts], c='b', s=100.0)  #plot F-nodes in blue

show()
 

modified example from https://github.com/pyamg/pyamg-examples/blob/master/Rootnode/demo.py