Solve the Heat Equation with non-zero Dirichlet BCs with Implicit Euler and Conjugate Gradient Linear Solvers?

Question

Many users have asked how to solve the Heat Equation, u_t = u_xx, with non-zero Dirichlet BCs and with conjugate gradients for the internal linear solver. This is a common simplified PDE problem before moving to more difficult versions of parabolic PDEs. How is this done in DifferentialEquations.jl?

Answer 1

Let's solve this problem in steps. First, let's build the linear operator for the discretized Heat Equation with Dirichlet BCs. A discussion of the discretization can be found on this Wiki page which shows that the central difference method gives a 2nd order discretization of the second derivative by (u[i-1] - 2u[i] + u[i+1])/dx^2. This is the same as multiplying by the Tridiagonal matrix of [1 -2 1]*(1/dx^2), so let's start by building this matrix:

using LinearAlgebra, OrdinaryDiffEq
x = collect(-π : 2π/511 : π)

## Dirichlet 0 BCs

u0 = @. -(x).^2 + π^2

n = length(x)
A = 1/(2π/511)^2 * Tridiagonal(ones(n-1),-2ones(n),ones(n-1))

Notice that we have implicitly simplified the end, since (u[0] - 2u[1] + u[2])/dx^2 = (- 2u[1] + u[2])/dx^2 when the left BC is zero, so the term is dropped from the matmul. We then use this discretization of the derivative to solve the Heat Equation:

function f(du,u,A,t)
    mul!(du,A,u)
end

prob = ODEProblem(f,u0,(0.0,10.0),A)
sol = solve(prob,ImplicitEuler())

using Plots
plot(sol[1])
plot!(sol[end])

Now we make the BCs non-zero. Notice that we just have to add back the u[0]/dx^2 that we previously dropped off, so we have:

## Dirichlet non-zero BCs
## Note that the operator is no longer linear
## To handle affine BCs, we add the dropped term

u0 = @. (x - 0.5).^2 + 1/12
n = length(x)
A = 1/(2π/511)^2 * Tridiagonal(ones(n-1),-2ones(n),ones(n-1))

function f(du,u,A,t)
    mul!(du,A,u)
    # Now do the affine part of the BCs
    du[1] += 1/(2π/511)^2 * u0[1]
    du[end] += 1/(2π/511)^2 * u0[end]
end

prob = ODEProblem(f,u0,(0.0,10.0),A)
sol = solve(prob,ImplicitEuler())

plot(sol[1])
plot!(sol[end])

Now let's swap out the linear solver. The documentation suggests that you should use LinSolveCG here, which looks like:

sol = solve(prob,ImplicitEuler(linsolve=LinSolveCG()))

There are some advantages to this, since it has a norm handling that helps conditioning. Howerver, the documentation also states that you can build your own linear solver routine. This is done by giving a Val{:init} dispatch that returns the type to use as the linear solver, so we do:

## Create a linear solver for CG
using IterativeSolvers

function linsolve!(::Type{Val{:init}},f,u0;kwargs...)
  function _linsolve!(x,A,b,update_matrix=false;kwargs...)
    cg!(x,A,b)
  end
end

sol = solve(prob,ImplicitEuler(linsolve=linsolve!))

plot(sol[1])
plot!(sol[end])

And there we are, non-zero Dirichlet Heat Equation with a Krylov method (conjugate gradients) for the linear solver, making it a Newton-Krylov method.