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
kron3 = lambda x,y,z: sps.kron(x,sps.kron(y,z))
In [ ]:
N = 8
Ex = 4
Ey = 4
Ez = 4
nx = N+1
ny = N+1
nz = N+1
nx_dofs = N*Ex+1
ny_dofs = N*Ey+1
nz_dofs = N*Ez+1
n_elem = Ex*Ey*Ez
periodic = True
if periodic:
nx_dofs -= 1
ny_dofs -= 1
nz_dofs -= 1
semh = SEMhat(N)
In [ ]:
# def f(X):
# x = X[:,0]
# y = X[:,1]
# z = X[:,2]
# return np.sin(np.pi*x)*np.sin(np.pi*y)*np.sin(np.pi*z)
# def f2(X):
# x = X[:,0]
# y = X[:,1]
# z = X[:,2]
# return np.sin(np.pi*x)*np.sin(np.pi*y)*np.sin(np.pi*z)*3*(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]
zeta = X[:,2]
P = np.zeros_like(X)
P[:,0] = chi +sx*np.sin(np.pi*chi)*np.sin(np.pi*eta)*np.sin(np.pi*zeta)
P[:,1] = eta +sy*np.sin(np.pi*chi)*np.sin(np.pi*eta)*np.sin(np.pi*zeta)
P[:,2] = zeta+sy*np.sin(np.pi*chi)*np.sin(np.pi*eta)*np.sin(np.pi*zeta)
return P
def calc_jacb(X):
chi = X[:,0]
eta = X[:,1]
zeta = X[:,2]
J = np.zeros((X.shape[0],3,3))
J[:,0,0] = 1+sx*np.pi*np.cos(np.pi*chi)*np.sin(np.pi*eta)*np.sin(np.pi*zeta)
J[:,0,1] = sx*np.pi*np.sin(np.pi*chi)*np.cos(np.pi*eta)*np.sin(np.pi*zeta)
J[:,0,2] = sx*np.pi*np.sin(np.pi*chi)*np.sin(np.pi*eta)*np.cos(np.pi*zeta)
J[:,1,0] = sy*np.pi*np.cos(np.pi*chi)*np.sin(np.pi*eta)*np.sin(np.pi*zeta)
J[:,1,1] = 1+sy*np.pi*np.sin(np.pi*chi)*np.cos(np.pi*eta)*np.sin(np.pi*zeta)
J[:,1,2] = sy*np.pi*np.sin(np.pi*chi)*np.sin(np.pi*eta)*np.cos(np.pi*zeta)
J[:,2,0] = sy*np.pi*np.cos(np.pi*chi)*np.sin(np.pi*eta)*np.sin(np.pi*zeta)
J[:,2,1] = sy*np.pi*np.sin(np.pi*chi)*np.cos(np.pi*eta)*np.sin(np.pi*zeta)
J[:,2,2] = 1+sy*np.pi*np.sin(np.pi*chi)*np.sin(np.pi*eta)*np.cos(np.pi*zeta)
return J
In [ ]:
topo = Interval()
jma = []
mqa = []
for Em in [Ez, 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]
zq, yq, xq = mqa
dz, dy, dx = map(np.unique, mqa)
jacb_det0z, jacb_det0y, jacb_det0x = jma
XYZ = np.zeros((nz_dofs,ny_dofs,nx_dofs,3))
XYZ[:,:,:,0] = dx[na,na,:]
XYZ[:,:,:,1] = dy[na,:,na]
XYZ[:,:,:,2] = dz[:,na,na]
dof_ref = XYZ.reshape((-1,3))
dof_phys = ref_to_phys(dof_ref)
In [ ]:
# Build restriction operator
if periodic:
R0x = sps.eye(nx_dofs)
R0y = sps.eye(ny_dofs)
R0z = sps.eye(nz_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))
R0z = sps.dia_matrix((np.ones(nz_dofs),1),
shape=(nz_dofs-2,nz_dofs))
R = kron3(R0z, R0y, R0x)
rngx = np.arange(nx_dofs)
rngy = np.arange(ny_dofs)
rngz = np.arange(nz_dofs)
if not periodic:
assert False
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]])
az = np.repeat(semh.wgll[na,:-1], Ez, axis=0)
az = az.ravel()
if not periodic:
az = np.hstack([az,az[0]])
wvals = az[:,na,na]*ay[na,:,na]*ax[na,na,:]
wvals *= jacb_det0x*jacb_det0y*jacb_det0z
wvals = wvals.ravel()
In [ ]:
# Make Q
Qa = []
for Em in [Ez, 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 = kron3(*Qa)
a = Q.dot(np.arange(nx_dofs*ny_dofs*nz_dofs)).reshape((nz*Ez,ny*Ey,nx*Ex))
etd = np.zeros((Ex*Ey*Ez, nx*ny*nz), dtype=np.int)
ind = 0
for iz in range(Ez):
for iy in range(Ey):
for ix in range(Ex):
indz = iz*nz
indy = iy*ny
indx = ix*nx
etd[ind,:] = a[indz:indz+nz,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 = []
G13 = []
G21 = []
G22 = []
G23 = []
G31 = []
G32 = []
G33 = []
nn = nx*ny*nz
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)]
G13 += [sps.dia_matrix((G0[etd[i],0,2], 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)]
G23 += [sps.dia_matrix((G0[etd[i],1,2], 0), shape=s)]
G31 += [sps.dia_matrix((G0[etd[i],2,0], 0), shape=s)]
G32 += [sps.dia_matrix((G0[etd[i],2,1], 0), shape=s)]
G33 += [sps.dia_matrix((G0[etd[i],2,2], 0), shape=s)]
In [ ]:
# Build poisson stiffness matrix A
D1 = kron3(sps.eye(nz), sps.eye(ny), semh.Dh)/jacb_det0x
D2 = kron3(sps.eye(nz), semh.Dh, sps.eye(nx))/jacb_det0y
D3 = kron3(semh.Dh, sps.eye(ny), sps.eye(nx))/jacb_det0z
A0a = []
for i in range(n_elem):
A0a += [D1.T.dot(G11[i].dot(D1)+G12[i].dot(D2)+G13[i].dot(D3))+\
D2.T.dot(G21[i].dot(D1)+G22[i].dot(D2)+G23[i].dot(D3))+\
D3.T.dot(G31[i].dot(D1)+G32[i].dot(D2)+G33[i].dot(D3))]
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*nz_dofs
b = Q.T.dot((wvals*jacb_det)[etd.ravel()])
# 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*nz))
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))
y[i] += D3.T.dot(G31[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))
y[i] += D3.T.dot(G32[i].dot(Dx))
Dx = D3.dot(x[i])
y[i] += D1.T.dot(G13[i].dot(Dx))
y[i] += D2.T.dot(G23[i].dot(Dx))
y[i] += D3.T.dot(G33[i].dot(Dx))
return Q.T.dot(y.ravel())
nn = nx_dofs*ny_dofs*nz_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*nz_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=500, 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')
ax.scatter3D(dof_phys[:,0],
dof_phys[:,1],
dof_phys[:,2], s=1)
In [ ]: