In [ ]:
%matplotlib inline
In [ ]:
import numpy as np
from numpy import newaxis as na
import scipy
import scipy.sparse as sps
from scipy.sparse.linalg import spsolve, LinearOperator
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
In [ ]:
from pyamg.classical import ruge_stuben_solver
from pyfem.sem import SEMhat
from pyfem.topo import Interval
norm = lambda x: np.max(np.abs(x)) if len(x)>0 else 0.0
In [ ]:
N = 16
Ex = 4
Ey = 4
nx = N+1
ny = N+1
nx_dofs = N*Ex+1
ny_dofs = N*Ey+1
n_elem = Ex*Ey
periodic = True
if periodic:
nx_dofs -= 1
ny_dofs -= 1
semh = SEMhat(N)
In [ ]:
# def f(X):
# x = X[:,0]
# y = X[:,1]
# return np.sin(np.pi*x)*np.sin(np.pi*y)
# def f2(X):
# x = X[:,0]
# y = X[:,1]
# return np.sin(np.pi*x)*np.sin(np.pi*y)*2*(np.pi)**2
def f(X):
x = X[:,0]
y = X[:,1]
return np.cos(np.pi*x)*np.cos(np.pi*y)
def f2(X):
x = X[:,0]
y = X[:,1]
return np.cos(np.pi*x)*np.cos(np.pi*y)*2*(np.pi)**2
sx = 0.1
sy = 0.1
def ref_to_phys(X):
chi = X[:,0]
eta = X[:,1]
P = np.zeros_like(X)
P[:,0] = chi+sx*np.sin(np.pi*chi)*np.sin(np.pi*eta)
P[:,1] = eta+sy*np.sin(np.pi*chi)*np.sin(np.pi*eta)
return P
def calc_jacb(X):
chi = X[:,0]
eta = X[:,1]
J = np.zeros((X.shape[0],2,2))
J[:,0,0] = np.ones_like(chi) +sx*np.pi*np.cos(np.pi*chi)*np.sin(np.pi*eta)
J[:,0,1] = np.zeros_like(chi)+sx*np.pi*np.sin(np.pi*chi)*np.cos(np.pi*eta)
J[:,1,0] = np.zeros_like(eta)+sy*np.pi*np.cos(np.pi*chi)*np.sin(np.pi*eta)
J[:,1,1] = np.ones_like(eta) +sy*np.pi*np.sin(np.pi*chi)*np.cos(np.pi*eta)
return J
In [ ]:
topo = Interval()
jma = []
mqa = []
for Em in [Ey, Ex]:
vertices = np.linspace(-1, 1, Em+1)
etvm = np.zeros((Em, 2), dtype=np.int)
etvm[:,0] = np.arange(Em)
etvm[:,1] = np.arange(Em)+1
mq = topo.ref_to_phys(vertices[etvm], semh.xgll)
jacb_det0m = topo.calc_jacb(vertices[etvm])[0]
if periodic:
mq = mq.ravel()[:-1]
jma += [jacb_det0m]
mqa += [mq]
yq, xq = mqa
dy, dx = map(np.unique, mqa)
jacb_det0y, jacb_det0x = jma
X, Y = np.meshgrid(dx, dy)
dof_ref = np.zeros((len(X.ravel()), 2))
dof_ref[:,0] = X.ravel()
dof_ref[:,1] = Y.ravel()
dof_phys = ref_to_phys(dof_ref)
In [ ]:
# Build restriction operator
if periodic:
R0x = sps.eye(nx_dofs)
R0y = sps.eye(ny_dofs)
else:
R0x = sps.dia_matrix((np.ones(nx_dofs),1),
shape=(nx_dofs-2,nx_dofs))
R0y = sps.dia_matrix((np.ones(ny_dofs),1),
shape=(ny_dofs-2,ny_dofs))
R = sps.kron(R0y, R0x)
rngx = np.arange(nx_dofs)
rngy = np.arange(ny_dofs)
if not periodic:
boundary_dofs = np.hstack([rngx,
rngx+nx_dofs*(ny_dofs-1),
rngy[1:-1]*nx_dofs])
boundary_dofs = np.hstack([boundary_dofs,
rngy[1:-1]*nx_dofs+nx_dofs-1])
else:
boundary_dofs = np.array([],dtype=np.int)
boundary_dofs.sort()
In [ ]:
ax = np.repeat(semh.wgll[na,:-1], Ex, axis=0)
ax = ax.ravel()
if not periodic:
ax = np.hstack([ax,ax[0]])
ay = np.repeat(semh.wgll[na,:-1], Ey, axis=0)
ay = ay.ravel()
if not periodic:
ay = np.hstack([ay,ay[0]])
wvals = ay[:,na]*ax[na,:]
wvals *= jacb_det0x*jacb_det0y
wvals = wvals.ravel()
In [ ]:
# Make Q
Qa = []
for Em in [Ey, Ex]:
etdm = np.arange(Em*(N+1))
etdm = etdm.reshape((Em, -1))
etdm -= np.arange(Em)[:,na]
if periodic:
etdm[-1,-1] = etdm[0,0]
cols = etdm.ravel()
rows = np.arange(len(cols))
vals = np.ones(len(cols))
Q0m = sps.coo_matrix((vals,(rows,cols))).tocsr()
Qa += [Q0m]
Q = sps.kron(*Qa)
a = Q.dot(np.arange(nx_dofs*ny_dofs)).reshape((ny*Ey,nx*Ex))
etd = np.zeros((Ex*Ey, nx*ny), dtype=np.int)
ind = 0
for iy in range(Ey):
for ix in range(Ex):
indy = iy*ny
indx = ix*nx
etd[ind,:] = a[indy:indy+ny,indx:indx+nx].ravel()
ind += 1
cols = etd.ravel()
rows = np.arange(len(cols))
vals = np.ones(len(cols))
Q = sps.coo_matrix((vals,(rows,cols))).tocsr()
In [ ]:
# Compute jacobian information
Jacb = calc_jacb(dof_ref)
jacb_det = np.linalg.det(Jacb).ravel()
Jacb_inv = np.linalg.inv(Jacb)
In [ ]:
# build Gij
G0 = np.zeros_like(Jacb)
assert len(Jacb)==len(wvals)
for i in range(len(Jacb)):
G0[i,:,:] = Jacb_inv[i].dot(Jacb_inv[i].T)
G0[i,:,:] *= wvals[i]*jacb_det[i]
G11 = []
G12 = []
G21 = []
G22 = []
nn = nx*ny
s = (nn, nn)
for i in range(n_elem):
G11 += [sps.dia_matrix((G0[etd[i],0,0], 0), shape=s)]
G12 += [sps.dia_matrix((G0[etd[i],0,1], 0), shape=s)]
G21 += [sps.dia_matrix((G0[etd[i],1,0], 0), shape=s)]
G22 += [sps.dia_matrix((G0[etd[i],1,1], 0), shape=s)]
In [ ]:
# Build poisson stiffness matrix A
D1 = sps.kron(sps.eye(ny), semh.Dh)/jacb_det0x
D2 = sps.kron(semh.Dh, sps.eye(nx))/jacb_det0y
A0a = []
for i in range(n_elem):
A0a += [D1.T.dot(G11[i].dot(D1)+G12[i].dot(D2))+\
D2.T.dot(G21[i].dot(D1)+G22[i].dot(D2))]
A0 = sps.block_diag(A0a).tocsr()
A0 = Q.T.dot(A0.dot(Q))
A = R.dot(A0.dot(R.T))
# Build mass matrix B
nd = nx_dofs*ny_dofs
b = Q.T.dot((wvals*jacb_det)[etd.ravel()])
#Bl = sps.dia_matrix((wvals*jacb_det, 0), shape=(nd,nd))
# Bl is not the local mass matrix.
# I am just using bad notation here
Bl = sps.dia_matrix((b, 0), shape=(nd,nd))
Binv_data = (1.0/Bl.data).ravel()
Binv_data = R.dot(Binv_data)
if nd<=1e3:
print np.min(np.linalg.svd(A.toarray())[1])
In [ ]:
def apply_A(x):
x = Q.dot(x)
x = x.reshape((n_elem, nx*ny))
y = np.zeros_like(x)
for i in xrange(n_elem):
Dx = D1.dot(x[i])
y[i] += D1.T.dot(G11[i].dot(Dx))
y[i] += D2.T.dot(G21[i].dot(Dx))
Dx = D2.dot(x[i])
y[i] += D1.T.dot(G12[i].dot(Dx))
y[i] += D2.T.dot(G22[i].dot(Dx))
return Q.T.dot(y.ravel())
nn = nx_dofs*ny_dofs
linOp = LinearOperator((nn, nn), matvec=apply_A)
In [ ]:
fh = f2(dof_phys)
fl = fh
rhs = Bl.dot(fl)
radj = np.zeros(nx_dofs*ny_dofs)
radj[boundary_dofs] = f(dof_phys)[boundary_dofs]
rhs = R.dot(rhs-A0.dot(radj))
if periodic:
rhs -= np.mean(rhs)
In [ ]:
# Check apply_A against full matrix
norm(apply_A(rhs)-A.dot(rhs))
In [ ]:
ml = ruge_stuben_solver(A)
residuals = []
sol = R.T.dot(ml.solve(rhs, tol=1e-14,
maxiter=1000, residuals=residuals,
accel='cg'))
sol[boundary_dofs] = f(dof_phys)[boundary_dofs]
if periodic:
sol -= np.mean(sol)
print len(residuals), residuals[-1]
print
print norm(f(dof_phys)-sol)/norm(f(dof_phys))
In [ ]:
class CB(object):
def __init__(self):
self.n_iter = 0
def __call__(self, x):
self.n_iter += 1
cb = CB()
solcg, errc = sps.linalg.cg(linOp, rhs, tol=1e-14, maxiter=2000,
callback=cb)
solcg = R.T.dot(solcg)
print cb.n_iter, norm(rhs-apply_A(solcg))
print
print norm(f(dof_phys)-solcg)/norm(f(dof_phys))
print norm(sol-solcg)
In [ ]:
fig = plt.figure()
ax = fig.gca(projection='3d')
s = (ny_dofs,nx_dofs)
X, Y = dof_phys[:,0], dof_phys[:,1]
X = X.reshape(s)
Y = Y.reshape(s)
ax.plot_wireframe(X, Y, f(dof_phys).reshape(s))
ax.plot_wireframe(X, Y, sol.reshape(s),
color='g')
plt.savefig("sol.pdf")
In [ ]:
plt.scatter(dof_phys[:,0], dof_phys[:,1])
In [ ]: