Solving the diffusion equation in a semi-infinite domain with the ultraspherical spectral method
One problem I have been interested in recently is numerical methods for computational fluid dynamics in semi-infinite domains, where you have some kind of boundary at
with a homogeneous boundary condition
to which we can compare our numerical solutions.
To work with the infinite domain, we apply a coordinate transformation (Boyd 2000)
of the interval
The forced diffusion equation in the new coordinates is
with the boundary condition
The diffusion equation becomes
where
The matrix
is the derivative operator and
I've written these operators as infinite dimensional ones, but in practice, we truncate the Chebyshev expansion of
While
We can, however, do better using the ultraspherical method of Olver and Townsend (2013). This rests on the fact that the derivatives of Chebyshev polynomials are scaled ultraspherical polynomials. Since we need two derivatives, we can convert to the ultraspherical basis of order 2 using the conversion matrices given on p. 8 and p. 12 of Olver and Townsend. This renders the second derivative matrix
Time discretization
There are many time discretizations that we could choose, especially with a simple pressure gradient forcing like
We end up solving
for
Boundary conditions
As in Olver and Townsend, we apply the boundary conditions by "boundary bordering," which is equivalent to a Chebyshev tau method (Boyd 2000). Basically we drop the bottom row of the matrix
Implementation in Julia
The ultraspherical spectral method is implemented within the excellent ApproxFun.jl package, but it is also fairly straightforward to implement using standard library routines for sparse linear algebra.
using LinearAlgebra, SparseArrays
First, we can create the derivative matrix
function chebyshev_derivative_matrix(Nz) sparse([(i < j) ? (i==0 ? 1 : 2)*j*mod(i+j,2) : 0 for i in 0:Nz-1, j in 0:Nz-1]) end function chebyshev_multiplication_matrix(Nz) G0 = sparse([((i == j+0) + (i == abs(j-0)))//2 for i in 0:Nz-1, j in 0:Nz-1]) G1 = sparse([((i == j+1) + (i == abs(j-1)))//2 for i in 0:Nz-1, j in 0:Nz-1]) G2 = sparse([((i == j+2) + (i == abs(j-2)))//2 for i in 0:Nz-1, j in 0:Nz-1]) # 0.5 * (1 - s)^2 = 3//4 T₀ - T₁ + 1//4 T₂ 3//4 * G0 - G1 + 1//4*G2 end
Using Integer
and Rational
types ensures that we can calculate the entries of these matrices without rounding errors.
The ultraspherical conversion matrices are likewise simple:
S1(Nz) = spdiagm(0=>[1;[1//(1+k) for k in 1:Nz-1]],2=>[-1//(1 + k) for k in 2:Nz-1]) S0(Nz) = spdiagm(0=>[2;ones(Int,Nz-1)] .// 2,2=>-ones(Int,Nz-2).//2)
We can now assemble the matrices for the left- and right-hand sides of the Crank-Nicolson solver
function assemble_matrices(Nz,Δt,ultraspherical) # Bottom boundary condition BC = [(-1)^k for k in 0:Nz-1]' # We make everything slightly larger to avoid truncation errors D = chebyshev_derivative_matrix(Nz + 1) G = chebyshev_multiplication_matrix(Nz + 1) if ultraspherical Δ = S1(Nz+1)*S0(Nz+1) * G*D*G*D A = S1(Nz+1)*S0(Nz+1) - Δt/2 * Δ B = S1(Nz+1)*S0(Nz+1) + Δt/2 * Δ else # Assemble the Chebyshev matrices without # the ultraspherical conversion Δ = G*D*G*D A = I - Δt/2 * Δ B = I + Δt/2 * Δ end # Apply boundary bordering and truncate properly [BC;A[1:Nz-1,1:Nz]],B[1:Nz,1:Nz] end
Our time step function simply assembles the right-hand side and then solves the equation ldiv!
to avoid some memory allocation. Since our pressure gradient forcing is spatially constant, we can add the forcing function to the first element of the right-hand side vector, which represents the ultraspherical coefficient for the constant function. To apply the boundary condition, we also use a trick that avoids having to allocate the vector [0;RHS[1:end-1]]
.
function timestep!(un,u,A,B,Δt,t) RHS = B*u RHS[1] += Δt/12 * (23 * cos(t) - 16 * cos(t - Δt) + 5 * cos(t - 2Δt)) RHS[end] = 0 circshift!(RHS,-1) ldiv!(un,A,RHS) end
To run the model for a fixed number of timesteps, we preallocate the output and loop
function run_model(U0,A,B,Δt,Nt) U = zeros(length(U0),Nt+1) U[:,1] = U0 for i in 1:Nt timestep!(view(U,:,i+1),view(U,:,i),A,B,Δt,(i-1)*Δt) end U end
We will also want the analytical solution for comparison and for our initial conditions
function stokes(z,t) sin(t) - exp(-z/sqrt(2))*sin(t - z/sqrt(2)) end
Since our result will be a vector of Chebyshev coefficients, we will also want to convert these to the values on a grid. There are multiple ways to do this, and you would normally use a fast cosine transform to implement the inverse Chebyshev transform. However, we only need to do this conversion twice, to compute the Chebyshev coefficients of the initial conditions and to compute the solution on the grid. The simplest way to do this is to compute a matrix of the Chebyshev functions using their recurrence relations
function chebyshev_matrix(x,P=length(x)) N = length(x) T = zeros(N,P) T[:,1] .= 1 T[:,2] .= x T[:,3] .= 2x.^2 .- 1 for i in 3:P-1 T[:,i+1] = 2*x.*T[:,i] .- T[:,i-1] end T end
And lastly, we need to compute the error in our solution, which we can do easily using Gauss-Chebyshev quadrature if we represent our solution on the Chebyshev roots grid.
function quadrature_error(z,u,u0) N = length(u) s = (z .- 1) ./ (z .+ 1) f = abs2.(u .- u0) .* sqrt.(1 .- s.^2) π*sum(f)/N end
Finally we wrap it all together. We'll solve the diffusion equation, but also time the solution and compute the error.
function run_test(Nz,Δt,Nt,Nq=1024;ultraspherical=true) # Chebyshev roots grid and transformed grid # Note that we use a high resolution grid here s = [cospi((2k + 1)/2Nq) for k in 0:Nq-1] z = (1 .+ s) ./ (1 .- s) T = chebyshev_matrix(s,Nq) # Initial conditions u0 = stokes.(z,0.0) # Forward Chebyshev transform to compute coefficients # Truncate from the high-order approximation U0 = (T\u0)[1:Nz] A,B = assemble_matrices(Nz,Δt,ultraspherical) Al = lu(A) # Run once to compute the results U = run_model(U0,Al,B,Δt,Nt) # Run again to time the solver t = @elapsed run_model(U0,Al,B,Δt,Nt) u = T[:,1:Nz]*U err = [quadrature_error(z,u[:,i],stokes.(z,(i-1)*Δt)) for i in 1:size(u,2)] u,t,err end
Results
If we run the model with Nz = 128
and Δt = 2π*0.01
for a single cycle, our results look something like this
Nz = 128 Δt = 2π*0.01 Nt = 100 Nq = 1024 u1,t1,err1 = run_test(Nz,Δt,Nt,Nq) s1 = [cospi((2k + 1)/2Nq) for k in 0:Nq-1] z1 = (1 .+ s1) ./ (1 .- s1) using CairoMakie tidx = Observable(1) uo = lift(tidx) do i u1[:,i] end fig = Figure() ax = Axis(fig[1,1],ylabel="Height",xlabel="Velocity") lines!(ax,uo,z1,color=:black) ylims!(ax,0,10) xlims!(ax,-1.1,1.1) record(fig, "ultraspherical_diffusion.mp4", 1:size(u1,2);framerate=30) do i tidx[] = i end
A video of a single period of the oscillating boundary layer flow
At this resolution, the numeric and analytic solutions are basically identical, so I have only plotted the numeric solution.
We can also run it at several resolutions to see how the computational time and the error scales with the grid size. We will run it for 10 cycles just to make sure that it is stable for long integration times. We will also run the Chebyshev method which results in dense matrices to compare its performance to the ultraspherical method.
Nzs = 16:16:256 res0 = [run_test(Nz,Δt,10*Nt,Nq;ultraspherical=false) for Nz in Nzs] res1 = [run_test(Nz,Δt,10*Nt,Nq;ultraspherical=true) for Nz in Nzs] us0,ts0,err0 = map(x->x[1],res0),map(x->x[2],res0),map(x->x[3][end],res0) us1,ts1,err1 = map(x->x[1],res1),map(x->x[2],res1),map(x->x[3][end],res1) fig2 = Figure() ax1 = Axis(fig2[1,1],ylabel="Time (s)",xlabel="Grid size") scatter!(ax1,Nzs,ts0,label="Chebyshev",marker=:circle) scatter!(ax1,Nzs,ts1,label="Ultraspherical",marker=:diamond) axislegend(ax1,position=:lt) ax2 = Axis(fig2[2,1],ylabel="Error",xlabel="Grid size",yscale=log10) scatter!(ax2,Nzs,err0,label="Chebyshev",marker=:circle) scatter!(ax2,Nzs,err1,label="Ultraspherical",marker=:diamond) axislegend(ax2,position=:rt) save("ultraspherical_scaling.png",fig2)

We can see that the ultraspherical method scales linearly with N while the Chebyshev method scales quadratically. The Chebyshev method has a lower error at small grid sizes, but the error converges between the two methods by N=64. At that point, most of the error is due to the time discretization, and it can be decreased further by decreasing the time step.
Spectral methods are really neat ways to solve PDEs accurately without a ton of effort, and the ultraspherical spectral method helps prevent the scaling issues that you get with a pure Chebyshev method. Using a coordinate transformation, it is also really easy to handle the semi-infinte domain that we want for theoretical bottom boundary layer studies. It is pretty straightforward to build an incompressible Navier-Stokes solver in this kind of framework, especially if we use periodic boundary conditions in the horizontal directions. The problem decouples into a set of vertical PDEs like our diffusion equation for each Fourier coefficient. One does have to address the incompressibility condition and the pressure computation, and we might see how that works later. Stay tuned!
You can find the Julia code contained in this file here.