def implicitsolve(mx, mt, L, T, c, source, ix, iv, lbc, rbc, lbctype, rbctype):
"""
Solve a wave equation using an implicit scheme. Unconditionally stable.
"""
xs, ts, deltax, deltat, lmbda = initialise(mx, mt, L, T, c)
# make sure these functions are vectorized
ix, iv, lbc, rbc, source = numpify_many((ix, 'x'), (iv, 'x'), (lbc, 't'),
(rbc, 't'), (source, 'x t'))
# construct matrices for implicit scheme
A_IW = tridiag(mx + 1, -0.5 * lmbda**2, 1 + lmbda**2, -0.5 * lmbda**2)
B_IW = tridiag(mx + 1, 0.5 * lmbda**2, -1 - lmbda**2, 0.5 * lmbda**2)
# corrections for Neumann conditions
A_IW[0, 1] *= 2
B_IW[0, 1] *= 2
A_IW[mx, mx - 1] *= 2
B_IW[mx, mx - 1] *= 2
# range of rows of A_EW to use
a = 1 if lbctype == 'Dirichlet' else 0
b = mx if rbctype == 'Dirichlet' else mx + 1
# initial condition vectors
U = ix(xs).copy()
V = iv(xs).copy()
# set first two time steps
u_jm1 = U
u_j = spsolve(A_IW, U - deltat * B_IW.dot(V))
# initialise u at next time step
u_jp1 = np.zeros(xs.size)
v = np.zeros(xs.size)
for j in ts[1:-1]:
v[a:b] = B_IW[a:b, a:b].dot(u_jm1[a:b]) + 2 * u_j[a:b]
addboundaries(v, lbctype, rbctype,
0.5 * lmbda**2 * (lbc(j) + lbc(j + deltat)),
0.5 * lmbda**2 * (rbc(j) + rbc(j + deltat)),
-lmbda**2 * deltax * (lbc(j) + lbc(j + deltat)),
lmbda**2 * deltax * (rbc(j) + rbc(j + deltat)))
u_jp1[a:b] = spsolve(A_IW[a:b, a:b], v[a:b])
if lbctype == 'Dirichlet':
u_jp1[0] = lbc(j + deltat)
if rbctype == 'Dirichlet':
u_jp1[mx] = rbc(j + deltat)
# add source to inner terms
u_jp1[1:-1] += deltat * source(xs[1:-1], j)
# update u_jm1 and u_j
u_jm1[:], u_j[:] = u_j[:], u_jp1[:]
return xs, u_j, lmbda
def cranknicholson(mx, mt, L, T, kappa, source, ic, lbc, rbc, lbctype,
rbctype):
"""Crank-Nicholson finite-difference scheme (implicit) for solving
parabolic PDE problems. Unconditionally stable"""
# initialise parameters
xs, ts, deltax, deltat, lmbda = initialise(mx, mt, L, T, kappa)
# make sure these functions are vectorized
ic, lbc, rbc, source = numpify_many((ic, 'x'), (lbc, 't'), (rbc, 't'),
(source, 'x t'))
# Construct Crank-Nicholson matrices
A_CN = tridiag(mx + 1, -0.5 * lmbda, 1 + lmbda, -0.5 * lmbda)
B_CN = tridiag(mx + 1, 0.5 * lmbda, 1 - lmbda, 0.5 * lmbda)
# modify first and last row for Neumann conditions
A_CN[0, 1] *= 2
A_CN[mx, mx - 1] *= 2
B_CN[0, 1] *= 2
B_CN[mx, mx - 1] *= 2
# restrict range of A_CN and B_CN in case of Dirichlet conditions
a = 1 if lbctype == 'Dirichlet' else 0
b = mx if rbctype == 'Dirichlet' else mx + 1
# initialise first time steps
u_j = ic(xs).copy()
u_jp1 = np.zeros(xs.size)
v = np.zeros(xs.size)
# Solve the PDE at each time step
for j in ts[:-1]:
v[a:b] = B_CN[a:b, a:b].dot(u_j[a:b])
addboundaries(v, lbctype, rbctype,
0.5 * lmbda * (lbc(j) + lbc(j + deltat)),
0.5 * lmbda * (rbc(j) + rbc(j + deltat)),
-lmbda * deltax * (lbc(j) + lbc(j + deltat)),
lmbda * deltax * (rbc(j) + rbc(j + deltat)))
u_jp1[a:b] = spsolve(A_CN[a:b, a:b], v[a:b])
# fix Dirichlet boundary conditions
if lbctype == 'Dirichlet':
u_jp1[0] = lbc(j)
if rbctype == 'Dirichlet':
u_jp1[mx] = rbc(j)
# add source to inner terms
u_jp1[1:-1] += 0.5 * deltat * (source(xs[1:-1], j) +
source(xs[1:-1], j + deltat))
u_j[:] = u_jp1[:]
return xs, u_j, lmbda
def explicitsolve(mx, mt, L, T, c, source, ix, iv, lbc, rbc, lbctype, rbctype):
"""
Solve a wave equation using an explicit scheme. Conditionally
stable for lambda <= 1. lambda = 1 is the 'magic step'.
"""
xs, ts, deltax, deltat, lmbda = initialise(mx, mt, L, T, c)
# make sure these functions are vectorized
ix, iv, lbc, rbc, source = numpify_many((ix, 'x'), (iv, 'x'), (lbc, 't'),
(rbc, 't'), (source, 'x t'))
# Construct explicit wave matrix
A_EW = tridiag(mx + 1, lmbda**2, 2 - 2 * lmbda**2, lmbda**2)
# Changes for Neumann boundary conditions
A_EW[0, 1] *= 2
A_EW[mx, mx - 1] *= 2
# range of rows of A_EW to use
a = 1 if lbctype == 'Dirichlet' else 0
b = mx if rbctype == 'Dirichlet' else mx + 1
# initial condition vectors
U = ix(xs).copy()
V = iv(xs).copy()
# set first two time steps
u_jm1 = U
u_j = 0.5 * A_EW.dot(U) + deltat * V
# initialise u at next time step
u_jp1 = np.zeros(xs.size)
for j in ts[1:-1]:
u_jp1[a:b] = A_EW[a:b, a:b].dot(u_j[a:b]) - u_jm1[a:b]
addboundaries(u_jp1, lbctype, rbctype, lmbda**2 * lbc(j - deltat),
lmbda**2 * rbc(j - deltat),
-2 * lmbda**2 * deltax * lbc(j - deltat),
2 * lmbda**2 * deltax * rbc(j - deltat))
# fix Dirichlet boundary conditions
if lbctype == 'Dirichlet':
u_jp1[0] = lbc(j + deltat)
if rbctype == 'Dirichlet':
u_jp1[mx] = rbc(j + deltat)
# add source to inner terms
u_jp1[1:-1] += deltat * source(xs[1:-1], j)
# update u_jm1 and u_j
u_jm1[:], u_j[:] = u_j[:], u_jp1[:]
return xs, u_j, lmbda
def backwardeuler(mx, mt, L, T, kappa, source, ic, lbc, rbc, lbctype, rbctype):
"""Backward Euler finite-difference scheme (implicit) for solving
parabolic PDE problems. Unconditionally stable"""
# initialise parameters
xs, ts, deltax, deltat, lmbda = initialise(mx, mt, L, T, kappa)
# make sure these functions are vectorized
ic, lbc, rbc, source = numpify_many((ic, 'x'), (lbc, 't'), (rbc, 't'),
(source, 'x t'))
# Construct backward Euler matrix
B_FE = tridiag(mx + 1, -lmbda, 1 + 2 * lmbda, -lmbda)
# modify first and last row for Neumann conditions
B_FE[0, 1] *= 2
B_FE[mx, mx - 1] *= 2
# restrict range of B_FE to use in case of Dirichlet conditions
a = 1 if lbctype == 'Dirichlet' else 0
b = mx if rbctype == 'Dirichlet' else mx + 1
# initialise first time steps
u_j = ic(xs).copy()
u_jp1 = np.zeros(xs.size)
# Solve the PDE at each time step
for j in ts[:-1]:
# modifications to the boundaries
addboundaries(
u_j,
lbctype,
rbctype,
lmbda * lbc(j + deltat), # Dirichlet u[1]
lmbda * rbc(j + deltat), # Dirichlet u[-2]
-2 * lmbda * deltax * lbc(j + deltat), # Neumann u[0]
2 * lmbda * deltax * rbc(j + deltat)) # Neumann u[-1]
u_jp1[a:b] = spsolve(B_FE[a:b, a:b], u_j[a:b])
# fix Dirichlet boundary conditions
if lbctype == 'Dirichlet':
u_jp1[0] = lbc(j + deltat)
if rbctype == 'Dirichlet':
u_jp1[mx] = rbc(j + deltat)
# add source to inner terms
u_jp1[1:-1] += deltat * source(xs[1:-1], j + deltat)
u_j[:] = u_jp1[:]
return xs, u_j, lmbda
def backwardeuler(mx, mt, L, T, kappa, source, ic, lbc, rbc, lbctype, rbctype):
"""Backward Euler finite-difference scheme (implicit) for solving
parabolic PDE problems. Unconditionally stable"""
# initialise
xs, ts, deltax, deltat, lmbda = initialise(mx, mt, L, T, kappa)
ic, lbc, rbc, source = numpify_many((ic, 'x'), (lbc, 'x'), (rbc, 'x'),
(source, 'x t'))
u_j = ic(xs)
print(type(u_j))
plot_solution(xs, u_j)
# if boundary conditions don't match initial conditions
if lbctype == 'Dirichlet':
u_j[0] = lbc(0)
if rbctype == 'Dirichlet':
u_j[mx] = rbc(0)
u_jp1 = np.zeros(xs.size)
# Construct forward Euler matrix
B_FE = tridiag(mx + 1, -lmbda, 1 + 2 * lmbda, -lmbda)
# modify first and last row for Neumann conditions
B_FE[0, 1] *= 2
B_FE[mx, mx - 1] *= 2
# range of rows of B_FE to use
a, b = matrixrowrange(mx, lbctype, rbctype)
print(a, b)
# Solve the PDE at each time step
for t in ts[:-1]:
addboundaries(u_j, lbctype, rbctype, lmbda * lbc(t + deltat),
lmbda * rbc(t + deltat),
-2 * lmbda * deltax * lbc(t + deltat),
2 * lmbda * deltax * rbc(t + deltat))
u_jp1[a:b] = spsolve(B_FE[a:b, a:b], u_j[a:b])
# fix Dirichlet boundary conditions
if lbctype == 'Dirichlet':
u_jp1[0] = lbc(t + deltat)
if rbctype == 'Dirichlet':
u_jp1[mx] = rbc(t + deltat)
# add source to inner terms
u_jp1[1:-1] += deltat * source(xs[1:-1], t + deltat)
u_j[:] = u_jp1[:]
return xs, u_j
def tsunami_solve(mx, mt, L, T, h0, h, wave):
"""Variable wavespeed problem assumes periodic boundary on the left and
an open boundary on the right"""
xs, ts, deltax, deltat, lmbda = initialise(mx, mt, L, T, np.sqrt(h0))
h, wave = numpify_many((h, 'x'), (wave, 'x'))
delta = deltat / deltax
# construct explicit wave matrix for variable wave speed problem
# with open boundary conditions initially
lower = [delta**2 * h(i - 0.5 * deltax) for i in xs[1:-1]] + [2 * lmbda**2]
main = [2*(1 + lmbda - lmbda**2)] + \
[2 - delta**2*(h(i + 0.5*deltax) + h(i - 0.5*deltax)) for i in xs[1:-1]] + \
[2*(1 + lmbda - lmbda**2)]
upper = [2 * lmbda**2] + [delta**2 * h(i + 0.5 * deltax) for i in xs[1:-1]]
A_EW = tridiag(mx + 1, lower, main, upper)
# initial condition vectors
u = [np.zeros(xs.size) for i in range(mt + 1)]
U = wave(xs).copy()
# set first two time steps
u[0] = U
u[1] = 0.5 * A_EW.dot(U)
# keep track of the first time the right side becomes non-zero
zero_right = True
for j in range(1, mt):
u[j + 1] = A_EW.dot(u[j]) - u[j - 1]
# correction for open boundary condition
u[j + 1][mx] /= (1 + 2 * lmbda)
# switch from open to periodic boundary
if zero_right and u[j + 1][mx] > 1e-5:
zero_right = False
# correction for open boundary conditon on the left before
# it becomes a periodic condition
if zero_right:
u[j + 1][0] /= (1 + 2 * lmbda)
else:
# periodic condition
u[j + 1][0] = u[j + 1][mx]
return xs, u
def forwardeuler(mx, mt, L, T, kappa, source, ic, lbc, rbc, lbctype, rbctype):
"""Forward euler finite-difference scheme (explicit) for solving
parabolic PDE problems. Values of lambda > 1/2 will cause the scheme
to become unstable"""
# initialise
xs, ts, deltax, deltat, lmbda = initialise(mx, mt, L, T, kappa)
# make sure these functions are vectorized
ic, lbc, rbc, source = numpify_many((ic, 'x'), (lbc, 't'), (rbc, 't'),
(source, 'x t'))
# Construct forward Euler matrix
A_FE = tridiag(mx + 1, lmbda, 1 - 2 * lmbda, lmbda)
# modify first and last row for Neumann conditions
A_FE[0, 1] *= 2
A_FE[mx, mx - 1] *= 2
# initialise first time steps
u_j = ic(xs).copy()
u_jp1 = np.zeros(xs.size)
# range of rows of A_FE to use
a = 1 if lbctype == 'Dirichlet' else 0
b = mx if rbctype == 'Dirichlet' else mx + 1
# Solve the PDE at each time step
for j in ts[:-1]:
u_jp1[a:b] = A_FE[a:b, a:b].dot(u_j[a:b])
addboundaries(u_jp1, lbctype, rbctype, lmbda * lbc(j), lmbda * rbc(j),
-2 * lmbda * deltax * lbc(j),
2 * lmbda * deltax * rbc(j))
# fix Dirichlet boundary conditions
if lbctype == 'Dirichlet':
u_jp1[0] = lbc(j + deltat)
if rbctype == 'Dirichlet':
u_jp1[mx] = rbc(j + deltat)
# add source to inner terms
u_jp1[1:-1] += deltat * source(xs[1:-1], j)
u_j[:] = u_jp1[:]
return xs, u_j, lmbda
def backwardeuler2(mx, mt, L, T, kappa, source, ic, lbc, rbc, lbc_ab, rbc_ab):
"""
Second implementation of backward Euler that takes mixed boundary
conditions
"""
# initialise parameters
xs, ts, deltax, deltat, lmbda = initialise(mx, mt, L, T, kappa)
# make sure these functions are vectorized
ic, lbc, rbc, source = numpify_many((ic, 'x'), (lbc, 't'), (rbc, 't'),
(source, 'x t'))
# Parameters needed to construct the matrix
alpha1, beta1 = lbc_ab
alpha2, beta2 = rbc_ab
# Construct the backward Euler matrix, first and last row are
# boundary condition equations
lower = (mx - 1) * [-lmbda] + [-beta2]
main = [alpha1*deltax - beta1] + (mx-1)*[1+2*lmbda] + \
[beta2 + alpha2*deltax]
upper = [beta1] + (mx - 1) * [-lmbda]
A_BE = tridiag(mx + 1, lower, main, upper)
# initialise the first time steps
u_j = ic(xs).copy()
u_jp1 = np.zeros(xs.size)
# Solve the PDE: loop over all time points
for j in ts[:-1]:
# Add boundary conditions to vector u_j
u_j[0] = deltax * lbc(j * deltat)
u_j[mx] = deltax * rbc(j * deltat)
# Backward euler timestep
u_jp1 = spsolve(A_BE, u_j)
# add source function
u_jp1[1:mx] += deltat * source(xs[1:-1], j * deltat)
# Update u_j
u_j[:] = u_jp1[:]
return xs, u_j, lmbda
def SOR(mx, my, Lx, Ly,
lbc, rbc, tbc, bbc,
maxcount, maxerr, omega, u_exact, plot=True):
"""Solve Laplace's equation on a rectangle with Dirichlet boundary
conditions, using the SOR method"""
# set up the numerical environment variables
xs = np.linspace(0, Lx, mx+1) # mesh points in x
ys = np.linspace(0, Ly, my+1) # mesh points in y
deltax = xs[1] - xs[0] # gridspacing in x
deltay = ys[1] - ys[0] # gridspacing in y
lambdasqr = (deltax/deltay)**2 # mesh number
u_old = np.zeros((xs.size,ys.size)) # u at current time step
u_new = np.zeros((xs.size,ys.size)) # u at next time step
u_true = np.zeros((xs.size,ys.size)) # exact solution
R = np.zeros((xs.size, ys.size)) # residual (goes to 0 with convergence)
bbc, tbc, lbc, rbc, u_exact = numpify_many((bbc, 'x'), (tbc, 'x'),
(lbc, 'x'), (rbc, 'x'),
(u_exact, 'x y'))
# intialise the boundary conditions, for both timesteps
u_old[1:-1,0] = bbc(xs[1:-1])
u_old[1:-1,-1] = tbc(xs[1:-1])
u_old[0,1:-1] = lbc(ys[1:-1])
u_old[-1,1:-1] = rbc(ys[1:-1])
u_new[:]=u_old[:]
if u_exact:
# true solution values on the grid
for i in range(0,mx+1):
for j in range(0,my+1):
u_true[i,j] = u_exact(xs[i],ys[j])
count = 1
err = maxerr+1
while err>maxerr and count<maxcount:
for j in range(1,my):
for i in range(1,mx):
# calculate residual
R[i,j] = (u_new[i-1,j] + u_old[i+1,j] - \
2*(1+lambdasqr)*u_old[i,j] + \
lambdasqr*(u_new[i,j-1]+u_old[i,j+1]) )/(2*(1+lambdasqr))
# SOR step
u_new[i,j] = u_old[i,j] + omega*R[i,j]
err = np.max(np.abs(u_new-u_old))
u_old[:] = u_new[:]
count=count+1
if u_exact:
# calculate the error, compared to the true solution
err_true = np.max(np.abs(u_new[1:-1,1:-1]-u_true[1:-1,1:-1]))
else:
err_true = 0
if plot:
# plot the resulting solution
xx = np.append(xs,xs[-1]+deltax)-0.5*deltax # cell corners needed for pcolormesh
yy = np.append(ys,ys[-1]+deltay)-0.5*deltay
plt.pcolormesh(xx,yy,u_new.T)
cb = plt.colorbar()
cb.set_label('u(x,y)')
plt.xlabel('x'); plt.ylabel('y')
plt.show()
return u_new, count, err, err_true