import numpy as np
from scipy.sparse import coo_matrix, csr_matrix
from scipy.sparse.linalg import spsolve
from numba import jit
def _flattened_grid_indices(resolution):
m,n = resolution
j = np.arange(n-1)
nodes = np.empty( (4, (m-1)*(n-1)), dtype=int )
for i in range(m-1):
nodes[0,i*(n-1)+j] = i*n + j
nodes[1,:] = nodes[0,:] + n
nodes[2,:] = nodes[0,:] + (n+1)
nodes[3,:] = nodes[0,:] + 1
return nodes
[docs]def load_vector(grid, f):
m,n = grid.resolution
hx,hy = grid.increment
f_at_mid_pts = f[0:m-1,0:n-1] + f[0:m-1,1:n] + f[1:m,0:n-1] + f[1:m,1:n]
f_at_mid_pts *= 0.25*0.25*hx*hy
F = np.zeros_like(f)
F[0:m-1,0:n-1] = f_at_mid_pts
F[0:m-1, 1:n ] += f_at_mid_pts
F[ 1:m, 0:n-1] += f_at_mid_pts
F[ 1:m, 1:n ] += f_at_mid_pts
return F.flatten()
_local_stiff0 = np.array([[ 2.0,-2.0,-1.0, 1.0 ],
[-2.0, 2.0, 1.0,-1.0 ],
[-1.0, 1.0, 2.0,-2.0 ],
[ 1.0,-1.0,-2.0, 2.0 ]]) / 6.0
_local_stiff1 = np.array([[ 2.0, 1.0,-1.0,-2.0 ],
[ 1.0, 2.0,-2.0,-1.0 ],
[-1.0,-2.0, 2.0, 1.0 ],
[-2.0,-1.0, 1.0, 2.0 ]]) / 6.0
_local_mass = np.array([[ 4.0, 2.0, 1.0, 2.0 ],
[ 2.0, 4.0, 2.0, 1.0 ],
[ 1.0, 2.0, 4.0, 2.0 ],
[ 2.0, 1.0, 2.0, 4.0 ]]) / 36.0
[docs]def assemble_system_matrix(grid, alpha, beta):
m,n = grid.resolution
hx,hy = grid.increment
N = m*n
stiff = (hy/hx) * _local_stiff0 + (hx/hy) * _local_stiff1
mass = hx*hy * _local_mass
if np.isscalar(alpha):
a = np.empty((m-1)*(n-1))
a[:] = alpha
else:
a = 0.25 * (alpha[0:m-1,0:n-1] + alpha[0:m-1,1:n] + \
alpha[ 1:m, 0:n-1] + alpha[ 1:m, 1:n])
a = a.flatten()
if np.isscalar(beta):
b = np.empty((m-1)*(n-1))
b[:] = beta
else:
b = 0.25 * (beta[0:m-1,0:n-1] + beta[0:m-1,1:n] + \
beta[ 1:m, 0:n-1] + beta[ 1:m, 1:n])
b = b.flatten()
nodes = _flattened_grid_indices( (m,n) )
A = csr_matrix((N,N))
for i,I in enumerate(nodes):
for j,J in enumerate(nodes):
A = A + coo_matrix( (a*mass[i,j] + b*stiff[i,j], (I,J)),
shape=(N,N) ).tocsr()
return A
@jit( nopython = True )
def _ellipt_matvec_loops(a,M,b,A,u,m,n):
v = np.zeros((m,n))
for i in range(m-1):
for j in range(n-1):
v1 = M[0,0] * u[ i, j] + M[0,3] * u[ i, j+1] + \
M[0,1] * u[i+1,j] + M[0,2] * u[i+1,j+1]
v2 = A[0,0] * u[ i, j] + A[0,3] * u[ i, j+1] + \
A[0,1] * u[i+1,j] + A[0,2] * u[i+1,j+1]
v[i,j] += a[i,j] * v1 + b[i,j] * v2
v1 = M[3,0] * u[ i, j] + M[3,3] * u[ i, j+1] + \
M[3,1] * u[i+1,j] + M[3,2] * u[i+1,j+1]
v2 = A[3,0] * u[ i, j] + A[3,3] * u[ i, j+1] + \
A[3,1] * u[i+1,j] + A[3,2] * u[i+1,j+1]
v[i,j+1] += a[i,j] * v1 + b[i,j] * v2
v1 = M[1,0] * u[ i, j] + M[1,3] * u[ i, j+1] + \
M[1,1] * u[i+1,j] + M[1,2] * u[i+1,j+1]
v2 = A[1,0] * u[ i, j] + A[1,3] * u[ i, j+1] + \
A[1,1] * u[i+1,j] + A[1,2] * u[i+1,j+1]
v[i+1,j] += a[i,j] * v1 + b[i,j] * v2
v1 = M[2,0] * u[ i, j] + M[2,3] * u[ i, j+1] + \
M[2,1] * u[i+1,j] + M[2,2] * u[i+1,j+1]
v2 = A[2,0] * u[ i, j] + A[2,3] * u[ i, j+1] + \
A[2,1] * u[i+1,j] + A[2,2] * u[i+1,j+1]
v[i+1,j+1] += a[i,j] * v1 + b[i,j] * v2
return v.reshape(m*n)
[docs]def ellipt_matvec(grid, alpha, beta, u):
m,n = grid.resolution
hx,hy = grid.increment
u = u.reshape((m,n))
if np.isscalar(alpha):
a = np.empty((m-1,n-1))
a[:] = alpha
else:
a = 0.25 * (alpha[0:m-1,0:n-1] + alpha[0:m-1,1:n] + \
alpha[ 1:m, 0:n-1] + alpha[ 1:m, 1:n])
if np.isscalar(beta):
b = np.empty((m-1,n-1))
b[:] = beta
else:
b = 0.25 * (beta[0:m-1,0:n-1] + beta[0:m-1,1:n] + \
beta[ 1:m, 0:n-1] + beta[ 1:m, 1:n])
A = (hy/hx) * _local_stiff0 + (hx/hy) * _local_stiff1
M = (hx*hy) * _local_mass
v = _ellipt_matvec_loops(a,M,b,A,u,m,n)
return v
@jit( nopython = True )
def _inv_diag_ellipt_matvec_loops(a,mass,b,stiff,u,m,n):
d = np.zeros((m,n))
v = np.empty((m,n))
M0 = mass[0,0]; M1 = mass[1,1]; M2 = mass[2,2]; M3 = mass[3,3]
A0 = stiff[0,0]; A1 = stiff[1,1]; A2 = stiff[2,2]; A3 = stiff[3,3]
for i in range(m-1):
for j in range(n-1):
d[ i, j ] += a[i,j] * M0 + b[i,j] * A0
d[ i, j+1] += a[i,j] * M3 + b[i,j] * A3
d[i+1, j ] += a[i,j] * M1 + b[i,j] * A1
d[i+1,j+1] += a[i,j] * M2 + b[i,j] * A2
for i in range(0,m):
for j in range(0,n):
v[i,j] = u[i,j] / d[i,j]
return v.reshape(m*n)
[docs]def inv_diag_ellipt_matvec(grid, alpha, beta, u):
m,n = grid.resolution
hx,hy = grid.increment
u = u.reshape((m,n))
if np.isscalar(alpha):
a = np.empty((m-1,n-1))
a[:] = alpha
else:
a = 0.25 * (alpha[0:m-1,0:n-1] + alpha[0:m-1,1:n] + \
alpha[ 1:m, 0:n-1] + alpha[ 1:m, 1:n])
if np.isscalar(beta):
b = np.empty((m-1,n-1))
b[:] = beta
else:
b = 0.25 * (beta[0:m-1,0:n-1] + beta[0:m-1,1:n] + \
beta[ 1:m, 0:n-1] + beta[ 1:m, 1:n])
stiff = (hy/hx) * _local_stiff0 + (hx/hy) * _local_stiff1
mass = (hx*hy) * _local_mass
v = _inv_diag_ellipt_matvec_loops(a,mass,b,stiff,u,m,n)
return v
[docs]def solve_elliptic_pde(grid, alpha, beta, f, g=0.0):
if grid.dim() != 2:
raise ValueError("Grid should be two-dimensional!")
if g != 0.0:
raise ValueError("Nonhomogeneous Neumann boundary condition with g != 0.0 has not been implemented yet!")
# rhs[:] += assemble_neumann_bc( grid, g )
rhs = load_vector( grid, f )
A = assemble_system_matrix( grid, alpha, beta )
u = spsolve( A, rhs )
u = u.reshape( grid.resolution )
return u
###########################################################################
[docs]class ReferenceElement:
def __init__(self):
self._quadrature = None
[docs] def quadrature(self, degree=None, order=None):
raise NotImplementedError("This function has not been implemented.")
[docs] def local_to_global(self, s, mesh, mask=None, coarsened=False):
raise NotImplementedError("This function has not been implemented.")
[docs] def interpolate(self,u,s):
raise NotImplementedError("This function has not been implemented.")