#Solve -(Lambda(u)*u')'+a0*u-f(x)=0
from scipy.integrate import fixed_quad
from scipy.special.orthogonal import p_roots
from scipy.sparse.linalg.dsolve import spsolve
from scipy import sparse
from numpy import arange, zeros, dot, array, exp, sqrt
from pylab import plot, figure, savefig, grid, legend, clf, pcolor, spy, axis, matshow
#finds the test function numbers on an element
def functionIndices( elem_index, num_elems, order=1 ):
#given the element index, the number of elements, and the order;
# we will find the function indices which exist on the specific element
if elem_index!=1 and elem_index!=num_elems: #give both hats on the interval
index_list = [elem_index-1,elem_index]
elif elem_index==1: #skip the known leftmost half-hat function \
index_list = [1]
else: #elem_index == num_elems skip the known rightmost half-hat function /
index_list = [num_elems-1]
if order>1: #return higher order functions too
for l_ in range(order-1):
index_list.append( num_elems + (order-1)*(elem_index-1) + l_ )
return index_list
#finds the values of u at x
def u( x , Y, mesh, num_elems, elem_index, dir_bdy, order=1):
index_list = functionIndices( elem_index, num_elems, order ) #get the functions on the current interval
list_len = len(index_list) #find how many functions are on the interval
result=0 #set to zero before summing
a = mesh[elem_index-1] #find the left endpoint of the interval
b = mesh[elem_index] #the right endpoint
x = 2.*(x-a)/(b-a)-1. #change to ref map i.e. (-1.,1.)
if list_len>order: #if the number of functions is greater than
for i in range(list_len): # the order then we are no on the left or right
result += Y[index_list[i]-1] * fn[i]( x ) # interval
elif elem_index==1: #leftmost interval
for i in range(list_len):
result += Y[index_list[i]-1] * fn[i+1]( x )
result +=dir_bdy[0] * fn[0](x) #include left Dirichlet boundary
else: #else we are on the right interval
result+=Y[index_list[0]-1] * fn[0]( x )
for i in range(list_len-1):
result += Y[index_list[i+1]-1] * fn[i+2]( x )
result += dir_bdy[1] * fn[1](x) #include right Dirichlet boundary
return result
# finds the value of du at x by first transforming to the -1 to 1 interval and then transforming back using a jacobian
def du( x , Y, mesh, num_elems, elem_index, dir_bdy, order=1):
index_list = functionIndices( elem_index, num_elems, order )
list_len = len(index_list)
result=0
a = mesh[elem_index-1]
b = mesh[elem_index]
x = 2.*(x-a)/(b-a)-1. #change to ref map
if list_len>order:
for i in range(list_len):
result += Y[index_list[i]-1] * dfn[i]( x )
elif elem_index==1:
for i in range(list_len):
result += Y[index_list[i]-1] * dfn[i+1]( x )
result +=dir_bdy[0] * dfn[0](x)
else:
result+=Y[index_list[0]-1] * dfn[0]( x )
for i in range(list_len-1):
result += Y[index_list[i+1]-1] * dfn[i+2]( x )
result +=dir_bdy[1] * dfn[1](x)
return result*2/(b-a)
# Defining shape functions and derivatives
def fn0(xi):
return (1. - xi)/2.
def dfn0(xi):
return -1./2.
def fn1(xi):
return (1. + xi)/2.
def dfn1(xi):
return 1./2.
def fn2(xi):
return xi**2 - 1.
def dfn2(xi):
return 2.*xi
def fn3(xi):
return xi*(xi**2 - 1.)
def dfn3(xi):
return 3.*xi**2 - 1.
def fn4(xi):
return xi**4 - 2.*xi**2 + 1
def dfn4(xi):
return 4*xi**3 - 4.*xi
fn = [fn0, fn1, fn2, fn3, fn4]
dfn = [dfn0, dfn1, dfn2, dfn3, dfn4]
#Finds the i'th element of F (RHS of Newton's Method)
def Fi( Y, u, dudx, f, Lambda, glob_i, order, num_elems, mesh, roots, weights, qpts_num, a0 ): #where elems is the list of elem intervals
#find number of element
elem_index = findElemNum( glob_i, num_elems, order )
#find test functions on element
test_list = functionIndices( elem_index, num_elems, order )
#INTEGRATION
a = mesh[elem_index-1]
b = mesh[elem_index]
jac = (b-a)/2.
result = 0.
vi_con = findCon( glob_i, num_elems, order, elem_index )
#integrate term lambda*u'*vi' on the element elem_index
for k in range(qpts_num):
result+= weights[k] * Lambda(u( x_ref( roots[k], a, b ) , Y, mesh, num_elems, elem_index, dir_bdy, order)) * du( x_ref( roots[k], a, b ) , Y, mesh, num_elems, elem_index, dir_bdy, order) * dfn[vi_con](roots[k])
#integrate a0*u*vi
for k in range(qpts_num):
result+= weights[k] * a0*u( x_ref( roots[k], a, b ) , Y, mesh, num_elems, elem_index, dir_bdy, order) * fn[vi_con](roots[k]) * jac
#integrate term -f(x)*vi
for k in range(qpts_num):
result+= -weights[k] * f( x_ref( roots[k], a, b ) ) * fn[vi_con](roots[k]) * jac
if glob_i<num_elems:
elem_index +=1
a = mesh[elem_index-1]
b = mesh[elem_index]
jac = (b-a)/2.
vi_con = findCon( glob_i, num_elems, order, elem_index )
#integrate term lambda*u'*vi' on the element elem_index
for k in range(qpts_num):
result+= weights[k] * Lambda(u( x_ref( roots[k], a, b ) , Y, mesh, num_elems, elem_index, dir_bdy, order)) * du( x_ref( roots[k], a, b ) , Y, mesh, num_elems, elem_index, dir_bdy, order) * dfn[vi_con](roots[k])
#integrate a0*u*vi
for k in range(qpts_num):
result+= weights[k] * a0*u( x_ref( roots[k], a, b ) , Y, mesh, num_elems, elem_index, dir_bdy, order) * fn[vi_con](roots[k]) * jac
#integrate term -f(x)*vi
for k in range(qpts_num):
result+= -weights[k] * f( x_ref( roots[k], a, b ) ) * fn[vi_con](roots[k]) * jac
return result
#Finds the i,jth element of the Jacobian Matrix of the Newton Method
def Jij( Y, u, dudx, f, Lambda, glob_i, glob_j, order, num_elems, mesh, roots, weights, qpts_num, a0 ):
#integrate dLambda*vj*du*vi + Lambda*dvj*vi + vj*vi
#find number of element for glob's
elem_index_i = findElemNum( glob_i, num_elems, order )
elem_index_j = findElemNum( glob_j, num_elems, order )
#check if the functions will overlap
overlapping_elems = 2 #first assume 2 overlaps
elem_index = elem_index_i #specifies the leftmost overlap after tests are done
if abs(elem_index_i-elem_index_j)>1:
return 0.
elif (abs(elem_index_i-elem_index_j)==1):
if glob_i>=num_elems and glob_j>=num_elems:#check if both are bubbles
return 0.
elif glob_i>=num_elems or glob_j>=num_elems: #one is bubble
if glob_i>=num_elems and (elem_index_i<elem_index_j):
return 0.
elif glob_j>=num_elems and (elem_index_j<elem_index_i):
return 0.
else: # a hat overlapping a bubble
overlapping_elems=1
if glob_j>=num_elems:
elem_index = elem_index_j
else:
elem_index = elem_index_i
else: #both are hats but with only one overlapping element
overlapping_elems = 1
if elem_index_i<elem_index_j:
elem_index = elem_index_j
else: #...j<...i
elem_index = elem_index_i
else: #same elem_index's returned
if glob_j>=num_elems or glob_i>=num_elems:
overlapping_elems=1 #a hat overlapping a bubble or bubble on bubble
#if they do overlap
#start on primary element specified by elem_index
#find test functions on element
test_list = functionIndices( elem_index, num_elems, order )
#INTEGRATION
a = mesh[elem_index-1]
b = mesh[elem_index]
jac = (b-a)/2.
result = 0
vi_con = findCon( glob_i, num_elems, order, elem_index )
vj_con = findCon( glob_j, num_elems, order, elem_index )
#integrate term dLambda*du*dvi*vj on the element elem_index
for k in range(qpts_num):
result+= weights[k] * dLambda(u( x_ref( roots[k], a, b ) , Y, mesh, num_elems, elem_index, dir_bdy, order)) * du( x_ref( roots[k], a, b ) , Y, mesh, num_elems, elem_index, dir_bdy, order) * dfn[vi_con](roots[k]) * fn[vj_con](roots[k])
#integrate Lambda*dvi*dvj
for k in range(qpts_num):
result+= weights[k] * Lambda(u( x_ref( roots[k], a, b ) , Y, mesh, num_elems, elem_index, dir_bdy, order)) * dfn[vi_con](roots[k]) * dfn[vj_con](roots[k]) / jac
#integrate term a0*vi*vj
for k in range(qpts_num):
result+= weights[k] * a0*fn[vi_con](roots[k]) * fn[vj_con](roots[k]) * jac
if overlapping_elems==2:
#INTEGRATION
elem_index = elem_index+1 #move one element to the right
a = mesh[elem_index-1]
b = mesh[elem_index]
jac = (b-a)/2.
vi_con = findCon( glob_i, num_elems, order, elem_index )
vj_con = findCon( glob_j, num_elems, order, elem_index )
#integrate term dLambda*du*dvi*vj on the element elem_index
for k in range(qpts_num):
result+= weights[k] * dLambda(u( x_ref( roots[k], a, b ) , Y, mesh, num_elems, elem_index, dir_bdy, order)) * du( x_ref( roots[k], a, b ) , Y, mesh, num_elems, elem_index, dir_bdy, order) * dfn[vi_con](roots[k]) * fn[vj_con](roots[k])
#integrate Lambda*dvi*dvj
for k in range(qpts_num):
result+= weights[k] * Lambda(u( x_ref( roots[k], a, b ) , Y, mesh, num_elems, elem_index, dir_bdy, order)) * dfn[vi_con](roots[k]) * dfn[vj_con](roots[k]) / jac
#integrate term a0*vi*vj
for k in range(qpts_num):
result+= weights[k] * a0*fn[vi_con](roots[k]) * fn[vj_con](roots[k]) * jac
return result
#assembles all i,j components of the Jacobian matrix and returns a CSR matrix
def assembleJ( Y, u, dudx, f, Lambda, order, num_elems, mesh, roots, weights, qpts_num, size, a0 ):
I=[]
J=[]
V=[]
for glob_i in range(1,size+1):
for glob_j in range(1,size+1):
I.append( glob_i-1 )
J.append( glob_j-1 )
V.append( Jij( Y, u, dudx, f, Lambda, glob_i, glob_j, order, num_elems, mesh, roots, weights, qpts_num, a0 ) )
# creating CSR matrix from arrays I, J, V
Sn = sparse.coo_matrix((V,(I,J)),shape=(size,size))
Sn = Sn.tocsr()
return Sn
#assembles all the Fi elements of the RHS vector for the Newton Method
def assembleF( Y, u, dudx, f, Lambda, order, num_elems, mesh, roots, weights, qpts_num, size, a0 ):
F = zeros( size )
for glob_i in range(1,size+1):
F[glob_i-1] = Fi( Y, u, dudx, f, Lambda, glob_i, order, num_elems, mesh, roots, weights, qpts_num, a0 )
return F
#transformation from the reference map to the (a,b) interval of the problem
def x_ref( xi, a, b ):
return (b-a)*(xi+1)/2. + a
#Higher order indices may be expressed in the form N+(p-1)*(k-1)+l where l runs from 0 to p-2
#Set this equal to the glob_i and solving for k (the element number) we find that we don't know l but the k must be an integer therefore l must satisfy a modulus formula given below.
def findL( glob_i, num_elems, order ):
if order==1 or glob_i<num_elems:
return -1
for l_ in range(order-1):
if ( glob_i - num_elems - l_ )%( order - 1)==0:
return l_ #l_ is 0 for order 2 and 1 for order 3...
print "ERROR!!! There is a bug in findL."
return -10. #this is not a good return! ERROR!
#Finds which test function to use given glob_i and the element we are on
def findCon( glob_i, num_elems, order, elem_num ):
if glob_i<num_elems:
if glob_i==elem_num:
return 1
else:
return 0
return findL( glob_i, num_elems, order )+2
#After defining the findL formula above we are able to find the element index given glob_i.
def findElemNum( glob_i, num_elems, order=1 ): #finds the element number given the data
#find number of element for glob_i
if glob_i < num_elems:
#return the left elem's index
elem_index_i = glob_i
else:
l_ = findL( glob_i, num_elems, order )
elem_index_i = 1 + ( glob_i - num_elems - l_ )/(order-1)
return elem_index_i
#makes a vector of equally space points from a to b. This is good for quick mesh's
def makeVector( a, b, n_points ):
vector_X = array( range(n_points) )/(n_points-1.)*(b-a)
for i in range(n_points):
vector_X[i] = vector_X[i]+a
return vector_X
#Find the L^2 norm of a vector Y
def norm(Y):
val = 0.
for y in Y: val += y**2
return sqrt(val)
#Solves the differential equation using Newton's Method. It terminates when F or dY are small.
def solveEquation( mesh, Y, order, a0, f, Lambda, dLambda, qpts_num, dir_bdy, tol=1e-6 ):
#roots and weights on -1 to 1
roots, weights = p_roots(qpts_num)
num_elems = len(mesh)-1
size = order*num_elems-1
#Solving
print "Solving."
iter_count = 0
while 1:
J = assembleJ( Y, u, du, f, Lambda, order, num_elems, mesh, roots, weights, qpts_num, size, a0 )
F = assembleF( Y, u, du, f, Lambda, order, num_elems, mesh, roots, weights, qpts_num, size, a0 )
norm_F = norm(F)
print "Residual norm: ", norm_F
if norm_F < tol:
print "Residual in tolerance."
break
print "Solving matrix problem."
dY = -spsolve(J, F)
Y = Y + dY
norm_dY = norm(dY)
print "Increment norm: ", norm_dY
iter_count += 1
if norm_dY < tol:
print "Increment in tolerance."
break
# Plotting
print "Plotting."
elem_sub_div = 9
x_vector = []
sol = []
for i in range(num_elems):
a = mesh[i]
b = mesh[i+1]
for j in range(elem_sub_div-1):
x = j/(elem_sub_div-1.)*(b-a) +a
x_vector.append( x )
u_val = u( x, Y, mesh, num_elems, i+1, dir_bdy, order)
sol.append( u_val )
x_vector.append(b)
sol.append( u( b, Y, mesh, num_elems, i+1, dir_bdy, order))
clf()
axis('equal')
label = "FEM solution"
plot(x_vector, sol, "g-", label=label)
legend()
grid(True)
savefig("a.png")
return iter_count
#INITIAL CONDITIONS DEFINED BY THE USER!
num_elems = 3
a = -1.
b = 1.
mesh = makeVector( a, b, num_elems+1 )
order = 4
qpts_num = order+4
tol = 1e-5
Y = zeros(order*num_elems-1 ) #initial guess
dir_bdy = [0.,0.]
#Parameters for the differential equation: -(Lambda(u)*u')'+a0*u=f(x)
a0 = 1.
def f(x):
return exp(-x)
def Lambda( u ):
return 1. + u**2 + u**6
def dLambda( u ):
return 2.*u+6.*u**5
#Finally we will solve the equation
iter_num = solveEquation( mesh, Y, order, a0, f, Lambda, dLambda, qpts_num, dir_bdy, tol )
print "Iterations: ", iter_num
##################################################################################################
##################################################################################################
##################################################################################################
Solving. Residual norm: 1.68910419612 Solving matrix problem. Increment norm: 0.504363216616 Residual norm: 0.0758100500431 Solving matrix problem. Increment norm: 0.0160598723432 Residual norm: 0.000463153485049 Solving matrix problem. Increment norm: 7.05441288309e-05 Residual norm: 1.0733192717e-08 Residual in tolerance. Plotting. Iterations: 3 
|