# Solve 1D heat equation, discretized using various finite difference schemes.
#
# UC Berkeley Math 228B, Per-Olof Persson <persson@berkeley.edu>
using PyPlot, PyCall, LinearAlgebra, DifferentialEquations
@pyimport IPython.display as d
"""
mkanim(x, allu)
Utility for animating PDE solutions in a Jupyter notebook.
"""
function mkanim(x, allu)
animation = pyimport("matplotlib.animation");
fig, ax = subplots(figsize=(6,4.5))
function update(frame)
ax.clear()
ax.plot(x, allu[frame+1])
ax.set_xlim([0, 1])
ax.set_ylim([-0.1, 1.1])
ax.grid(true);
end
ani = animation.FuncAnimation(fig, update, frames=length(allu), interval=50)
close(ani._fig)
d.HTML(ani.to_jshtml())
end
mkanim
"""
function heateqn_ftcs(m=100; T=0.2, kmul=0.5)
Solves the 1D heat equation with the FTCS scheme (Forward-Time,
Centered-Space), using grid size `m` and timestep multiplier `kmul`.
Integrates until final time `T` and returns the grid x and all solutions.
"""
function heateqn_ftcs(m=100; T=0.2, kmul=0.5)
# Discretization
h = 1.0 / (m+1)
x = h * (0:m+1)
k = kmul*h^2
N = ceil(Int, T/k)
u = @. exp(-(x - 0.25)^2 / 0.1^2) + 0.1sin(10*2π*x) # Initial conditions
u[[1,end]] .= 0 # Dirichlet boundary conditions u(0) = u(1) = 0
allu = [copy(u)]
for n = 1:N
u[2:m+1] += k/h^2 * (u[1:m] .- 2u[2:m+1] + u[3:m+2])
if mod(n, 10) == 0 # Save every 10th timestep
push!(allu, copy(u))
end
end
x,allu
end
x,allu = heateqn_ftcs();
mkanim(x,allu)
"""
function heateqn_cn(m=100; T=0.2, kmul=50)
Solves the 1D heat equation with the Crank−Nicolson scheme,
using grid size `m` and timestep multiplier `kmul`.
Integrates until final time `T` and returns the grid x and all solutions.
"""
function heateqn_cn(m=100; T=0.2, kmul=50)
# Discretization
h = 1.0 / (m+1)
x = h * (0:m+1)
k = kmul*h^2
N = ceil(Int, T/k)
u = exp.(-(x .- 0.25).^2 / 0.1^2) .+ 0.1sin.(10*2π*x) # Initial conditions
u[[1,end]] .= 0 # Dirichlet boundary conditions u(0) = u(1) = 0
# Form the matrices in the Crank-Nicolson scheme (Left and right)
A = SymTridiagonal(-2ones(m), ones(m-1)) / h^2
LH = I - A*k/2
RH = I + A*k/2
allu = [copy(u)]
for n = 1:N
u[2:m+1] = LH \ (RH * u[2:m+1]) # Note parenthesis for efficient evaluation
push!(allu, copy(u)) # Save every timestep
end
x,allu
end
x,allu = heateqn_cn()
mkanim(x,allu)
"""
function heateqn_odesolver(m=100; T=0.2, nsteps=100)
Solves the 1D heat equation using Method of Lines with ODE solvers from DifferentialEquations.jl.
Grid size `m`, integrates until final time `T` and returns the grid x and a total of `nsteps` solutions.
"""
function heateqn_odesolver(m=100; T=0.2, nsteps=100)
# Discretization
h = 1.0 / (m+1)
x = h * (0:m+1)
u = exp.(-(x .- 0.25).^2 / 0.1^2) .+ 0.1sin.(10*2π*x) # Initial conditions
u[[1,end]] .= 0 # Dirichlet boundary conditions u(0) = u(1) = 0
fode(u,p,t) = ([0; u[1:m+1]] .- 2u .+ [u[2:m+2]; 0]) / h^2 # RHS in du/dt = f(u)
prob = ODEProblem(fode, u, (0,T))
sol = solve(prob, alg_hints=[:stiff], saveat=T/nsteps)
x,sol.u
end
x,allu = heateqn_odesolver()
mkanim(x,allu)