Two-dimensional diffusion and advection - Fundamentals of Numerical Computation

We next describe how to apply the method of lines to PDEs of the form

u_t = \phi(u,u_x,u_y,u_{xx},u_{xy},u_{yy}), \quad (x,y)\in [a,b]\times [c,d].

(13.2.1)

The PDE may be of either parabolic or hyperbolic type, with the primary difference being potential restrictions on the time step size. To keep descriptions and implementations relatively simple, we will consider only periodic conditions or Dirichlet boundary conditions.

As described in Tensor-product discretizations, the rectangular domain is discretized by a grid $(x_i,y_j)$ for $i=0,\ldots,m$ and $j=0,\ldots,n$ . The solution is semidiscretized as a matrix $\mathbf{U}(t)$ such that $U_{ij}$ is the approximate solution at $(x_i,y_j,t)$ . Terms involving the spatial derivatives of $u$ are readily replaced by discrete counterparts, as shown in Table 13.2.1.

Table 13.2.1:Discrete replacements for 2D derivatives

term	discrete replacement
$u$	$\mathbf{U}$
$u_x$	$\mathbf{D}_x\mathbf{U}$
$u_y$	$\mathbf{U}\mathbf{D}_y^T$
$u_{xx}$	$\mathbf{D}_{xx}\mathbf{U}$
$u_{yy}$	$\mathbf{U}\mathbf{D}_{yy}^T$

13.2.1Matrix and vector shapes¶

Our destination is an IVP that can be solved by a Runge–Kutta or multistep solver. These solvers are intended for vector problems, but our unknowns naturally have a matrix shape, which is the most convenient for the differentiation formulas (13.1.7) and (13.1.8). Fortunately, it’s easy to translate back and forth between a matrix and an equivalent vector.

Definition 13.2.1 (vec and unvec operations)

Let $\mathbf{A}$ be an $m\times n$ matrix. Define the vec function as stacking the columns of $\mathbf{A}$ into a vector, i.e.,

\operatorname{vec}(\mathbf{A}) = \begin{bmatrix} A_{11} \\ \vdots \\ A_{m1} \\ \vdots \\ A_{1n} \\ \vdots \\ A_{m n} \end{bmatrix}.

(13.2.2)

Let $\mathbf{z}$ be a vector of length $m n$ . Define the unvec function as the inverse of vec:

\operatorname{unvec}(\mathbf{z}) = \begin{bmatrix} z_1 & z_{m+1} & \cdots & z_{m(n-1)+1} \\ z_2 & z_{m+2} & \cdots & z_{m(n-1)+2} \\ \vdots & \vdots & & \vdots \\ z_m & z_{2m} & \cdots & z_{m n} \\ \end{bmatrix}.

(13.2.3)

Suppose $\mathbf{U} = \operatorname{mtx}(u)$ is the matrix of unknowns. Table 13.2.2 shows how to convert between $\mathbf{U}$ and $\mathbf{u} = \operatorname{vec}(\mathbf{U})$ .

Table 13.2.2:vec and unvec operations

Julia

MATLAB

Python

vec	unvec
`u = vec(U)`	`U = reshape(u, m+1, n+1)`

vec	unvec
`u = U(:)`	`U = reshape(u, m+1, n+1)`

vec	unvec
`u = U.T.flatten()`	`U = np.reshape(u, (n+1, m+1)).T`

About the code

The incorporation of transposes is because NumPy uses row-major order, while MATLAB and Julia use column-major order. If performance were a concern, we would reverse our convention to avoid the transposes. We also use flatten rather than ravel to ensure making copies rather than views of the data and avoiding subtle bugs.

Example 13.2.1 (Reshaping for grid functions)

Julia

MATLAB

Python

Example 13.2.1

m = 2;
n = 3;
V = rand(1:9, m, n);
v = vec(V)

6-element Vector{Int64}:
 5
 8
 2
 2
 5
 5

The unvec operation is the inverse of vec.

unvec = z -> reshape(z, m, n)
unvec(v)

2×3 Matrix{Int64}:
 5  2  5
 8  2  5

Example 13.2.1

m = 4;  n = 3;
x = linspace(0, 2, m+1);
y = linspace(-3, 0, n+1);

f = @(x, y) cos(0.75*pi * x .* y - 0.5*pi * y);
[mtx, X, Y, vec, unvec] = tensorgrid(x, y);
F = mtx(f);
disp("function on a 4x3 grid:")
disp(F)

function on a 4x3 grid:

   -0.0000   -1.0000    0.0000    1.0000
    0.3827    0.7071    0.9239    1.0000
   -0.7071    0.0000    0.7071    1.0000
    0.9239   -0.7071   -0.3827    1.0000
   -1.0000    1.0000   -1.0000    1.0000

disp("vec(F):")
disp(vec(F))

vec(F):

The unvec operation is the inverse of vec.

disp("unvec(vec(F)):")
disp(unvec(vec(F)))

unvec(vec(F)):

   -0.0000   -1.0000    0.0000    1.0000
    0.3827    0.7071    0.9239    1.0000
   -0.7071    0.0000    0.7071    1.0000
    0.9239   -0.7071   -0.3827    1.0000
   -1.0000    1.0000   -1.0000    1.0000

Example 13.2.1

m, n = 4, 3
x = linspace(0, 2, m+1)
y = linspace(-3, 0, n+1)

f = lambda x, y: cos(0.75 * pi * x * y - 0.5 * pi * y)
mtx, X, Y, vec, unvec, _ = FNC.tensorgrid(x, y)
F = mtx(f)
print(f"function on a {m}x{n} grid:")
with printoptions(precision=4, suppress=True):
    print(F)

print("vec(F):")
with printoptions(precision=4, suppress=True):
    print(vec(F))

function on a 4x3 grid:
[[-0.     -1.      0.      1.    ]
 [ 0.3827  0.7071  0.9239  1.    ]
 [-0.7071  0.      0.7071  1.    ]
 [ 0.9239 -0.7071 -0.3827  1.    ]
 [-1.      1.     -1.      1.    ]]
vec(F):
[-0.      0.3827 -0.7071  0.9239 -1.     -1.      0.7071  0.     -0.7071
  1.      0.      0.9239  0.7071 -0.3827 -1.      1.      1.      1.
  1.      1.    ]

The unvec operation is the inverse of vec.

print("unvec(vec(F)):")
with printoptions(precision=4, suppress=True):
    print(unvec(vec(F)))

unvec(vec(F)):
[[-0.     -1.      0.      1.    ]
 [ 0.3827  0.7071  0.9239  1.    ]
 [-0.7071  0.      0.7071  1.    ]
 [ 0.9239 -0.7071 -0.3827  1.    ]
 [-1.      1.     -1.      1.    ]]

In order to modularize our codes, we use Algorithm 13.2.1 to define functions and values related to working with tensor-product grids. Its final output is to be discussed and used in Laplace and Poisson equations.

Algorithm 13.2.1 (tensorgrid)

Julia

MATLAB

Python

Create a tensor-product grid

tensorgrid.jl

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
"""
    tensorgrid(x, y)

Create a tensor grid for a rectangle from its 1d projections `x` and `y`.
Returns `unvec` to reshape a vector into a 2d array, `mtx` to evaluate a
function on the grid, `X`, `Y` to give the grid coordinates, and boolean array
`is_boundary` to identify the boundary points.
"""
function tensorgrid(x, y)
    m, n = length(x) - 1, length(y) - 1
    unvec(u) = reshape(u, m+1, n+1)
    mtx(h) = [h(x, y) for x in x, y in y]
    X = mtx((x, y) -> x)
    Y = mtx((x, y) -> y)
    is_boundary = trues(m+1, n+1)
    is_boundary[2:m, 2:n] .= false
    return mtx, X, Y, unvec, is_boundary
end

Create a tensor-product grid

tensorgrid.m

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
function [mtx, X, Y, vec, unvec, is_boundary] = tensorgrid(x, y)
% TENSORGRID Tensor-product grid.
% Input:
%   x, y           1d projections of the grid nodes (lengths m and n)
% Output:
%   mtx            evaluate a function on the grid (function)
%   X, Y           mtx applied to the coordinate functions (m×n)
%   vec            convert grid shape to vector shape (function)
%   unvec          convert vector shape to grid shape (function)
%   is_boundary    indicator of boundary nodes (logical m×n)
    m = length(x) - 1;
    n = length(y) - 1;
    vec = @(U) U(:);
    unvec = @(u) reshape(u, m+1, n+1);
    [X, Y] = ndgrid(x, y);
    function F = grideval(f)
        F = zeros(size(X));
        for k = 1:numel(X)
            F(k) = f(X(k), Y(k));
        end
    end
    mtx = @grideval;
    
    % Identify boundary points.
    is_boundary = true(m+1, n+1);
    is_boundary(2:m, 2:n) = false;
end

Create a tensor-product grid

tensorgrid.py

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
def tensorgrid(x, y):
    """
    tensorgrid(x, y)

    Create a tensor grid for a rectangle from its 1d projections x and y.
    Returns a function to reshape a 2d array to a vector, a function to reshape 
    a vector into a 2d array, a function to evaluate a function on the grid, 
    two arrays to give the grid coordinates, and a boolean array to identify 
    the boundary points.
    """
    m, n = len(x) - 1, len(y) - 1
    vec = lambda U: U.T.flatten()
    unvec = lambda u: np.reshape(u, (n+1, m+1)).T
    mtx = lambda h: np.array([[h(xi ,yj) for yj in y] for xi in x])
    X = mtx(lambda x, y: x)
    Y = mtx(lambda x, y: y)

    # Identify boundary points.
    is_boundary = np.tile(True, (m+1, n+1))
    is_boundary[1:-1, 1:-1] = False
    
    return mtx, X, Y, vec, unvec, is_boundary

13.2.2Periodic end conditions¶

If the boundary conditions are periodic, then the unknowns in the method of lines are the elements of the matrix $\mathbf{U}(t)$ representing grid values of the numerical solution. For the purposes of an IVP solution, this matrix is equivalent to the vector $\mathbf{u}(t)$ defined as $\mathbf{u}=\operatorname{vec}(\mathbf{U})$ .

Example 13.2.2 (Heat equation in 2D)

We will solve a 2D heat equation, $u_t = 0.1(u_{xx} + u_{yy})$ , on the square $[-1,1]\times[-1,1]$ , with periodic behavior in both directions. The initial condition is

u(x,y,0) = \sin(4\pi x) \exp[\cos(\pi y)],

(13.2.4)

which is also periodic on the rectangle.

Julia

MATLAB

Python

Example 13.2.2

We start by defining the discretization of the rectangle.

m, n = (60, 25)
x, Dx, Dxx = FNC.diffper(m, [-1, 1])
y, Dy, Dyy = FNC.diffper(n, [-1, 1])
mtx, X, Y, unvec, _ = FNC.tensorgrid(x, y);

Here is the initial condition, evaluated on the grid.

u_init = (x, y) -> sin(4 * π * x) * exp(cos(π * y))
U₀ = mtx(u_init)
M = maximum(abs, U₀)
contour(x, y, U₀';
    color=:redsblues,  clims=(-M, M), 
    aspect_ratio=1,
    xaxis=("x", (-1, 1)),  yaxis=("y", (-1, 1)), 
    title="Initial condition" )

The following function computes the time derivative for the unknowns, which have a vector shape. The actual calculations, however, take place using the matrix shape.

function du_dt(u, α, t)
    U = unvec(u)
    Uxx = Dxx * U
    Uyy = U * Dyy'            # 2nd partials
    du_dt = α * (Uxx + Uyy)    # PDE
    return vec(du_dt)
end;

Since this problem is parabolic, a stiff time integrator is appropriate.

using OrdinaryDiffEq
IVP = ODEProblem(du_dt, vec(U₀), (0, 0.2), 0.1)
sol = solve(IVP, Rodas4P());

We can use the function sol defined above to get the solution at any time. Its output is the matrix $\mathbf{U}$ of values on the grid. An animation shows convergence of the solution toward a uniform value.

Source

anim = @animate for t in range(0, 0.2, 81)
    surface(x, y, unvec(sol(t))';
        color=:redsblues,  clims=(-M, M),
        xaxis=(L"x", (-1, 1)), 
        yaxis=(L"y", (-1, 1)), 
        zlims=(-M, M),
        title=@sprintf("Heat equation, t=%.3f", t),
        dpi=150, colorbar=:none)
end
mp4(anim, "figures/2d-heat.mp4");

Example 13.2.2

We start by defining the discretization of the rectangle.

m = 60;  n = 40;
[x, Dx, Dxx] = diffper(m, [-1, 1]);
[y, Dy, Dyy] = diffper(n, [-1, 1]);
[mtx, X, Y, vec, unvec] = tensorgrid(x, y);

Here is the initial condition, evaluated on the grid.

U0 = sin(4*pi*X) .* exp( cos(pi*Y) );
clf,  surf(X', Y', U0')
mx = max(abs(vec(U0)));
clim([-mx, mx]),  shading interp
colormap(redsblues)
xlabel('x'),  ylabel('y')  
title('Initial condition')

The following function computes the time derivative for the unknowns which have a vector shape. The actual calculations, however, take place using the matrix shape.

f13_2_heat.m

function du_dt = timederiv(t, u, p)
    [alpha, Dxx, Dyy, vec, unvec] = p{:};
    U = unvec(u);
    Uxx = Dxx * U;  Uyy = U * Dyy';     % 2nd partials
    dU_dt = alpha * (Uxx + Uyy);  % PDE
    du_dt = vec(dU_dt);
end

Since this problem is parabolic, a stiff time integrator is appropriate.

odefun = @(t, u) f13_2_heat(t, u, {0.1, Dxx, Dyy, vec, unvec});
sol = ode15s(odefun, [0, 0.2], vec(U0));
sol = solutionFcn(ivp, 0, 0.2);
U = @(t) unvec(deval(sol, t));

Unrecognized function or variable 'ivp'.

We can use the function U defined above to get the solution at any time. Its output is a matrix of values on the grid.

Source

surf(X', Y', U(0.05)')
clim([-mx, mx]),  shading interp
colormap(redsblues)
xlabel('x'),  ylabel('y')  
title('Solution at t = 0.05')

Array indices must be positive integers or logical values.

An animation shows convergence toward a uniform value.

Source

title('Heat equation on a periodic domain')
vid = VideoWriter("figures/2d-heat.mp4","MPEG-4");
vid.Quality = 85;
open(vid);
for t = linspace(0, 0.2, 61)
    cla, surf(X', Y', U(t)')
    zlim([-3, 3]),  clim([-mx, mx])
    shading interp
    str = sprintf("t = %.2f", t);
    text(-0.9, 0.75, 2, str, fontsize=14);
    writeVideo(vid, frame2im(getframe(gcf)));
end
close(vid)

Example 13.2.2

We start by defining the discretization of the rectangle.

m, n = 60, 40
x, Dx, Dxx = FNC.diffper(m, [-1, 1])
y, Dy, Dyy = FNC.diffper(n, [-1, 1])
mtx, X, Y, vec, unvec, _ = FNC.tensorgrid(x, y)

Here is the initial condition, evaluated on the grid.

u_init = lambda x, y: sin(4 * pi * x) * exp(cos(pi * y))
U0 = mtx(u_init)
mx = max(abs(U0))
pcolormesh(X, Y, U0, vmin=-mx, vmax=mx, cmap="RdBu", shading="gouraud")
axis("equal"),  colorbar()
xlabel("$x$"),  ylabel("$y$")
title("Initial condition");

The following function computes the time derivative for the unknowns, which have a vector shape. The actual calculations, however, take place using the matrix shape.

alpha = 0.1
def du_dt(t, u):
    U = unvec(u)
    Uyy = Dxx @ U
    Uxx = U @ Dyy.T
    dU_dt = alpha * (Uxx + Uyy)  # PDE
    return vec(dU_dt)

Since this problem is parabolic, a stiff time integrator is appropriate.

from scipy.integrate import solve_ivp
sol = solve_ivp(du_dt, (0, 0.2), vec(U0), method="BDF", dense_output=True)
U = lambda t: unvec(sol.sol(t))

We can use the function U defined above to get the solution at any time. Its output is a matrix of values on the grid.

pcolormesh(X.T, Y.T, U(0.02).T, 
    vmin=-mx, vmax=mx, cmap="RdBu", shading="gouraud")
axis("equal"),  colorbar()
xlabel("$x$"),  ylabel("$y$")
title("Heat equation, t=0.02");

An animation shows convergence toward a uniform value.

from matplotlib import animation
fig, ax = subplots()
obj = ax.pcolormesh(X, Y, U(0), vmin=-mx, vmax=mx, cmap="RdBu", shading="gouraud")
time_text = ax.text(0.05, 0.9, '', transform=ax.transAxes)
ax.set_xlabel("$x$"),  ax.set_ylabel("$y$")
ax.set_aspect("equal")
ax.set_title("Heat equation on a periodic domain")
def snapshot(t):
    global obj
    obj.remove()
    obj = ax.pcolormesh(X, Y, U(t), vmin=-mx, vmax=mx, cmap="RdBu", shading="gouraud")
    time_text.set_text(f"t = {t:.2f}")

anim = animation.FuncAnimation(fig, snapshot, frames=linspace(0, 0.2, 41))
anim.save("figures/heat-2d.mp4", fps=30)
close()

13.2.3Dirichlet conditions¶

In Boundaries we coped with boundary conditions by removing the boundary values from the vector of unknowns being solved in the semidiscretized ODE. Each evaluation of the time derivative required us to extend the values to include the boundaries before applying differentiation matrices in space, then remove them from the time derivative vector.

We proceed similarly here, defining chop and extend functions that convert between the full grid and the inner grid of interior values. Mathematically speaking, the chopping operation is defined by

\operatorname{chop}(\mathbf{U}) = \mathbf{E}_x \mathbf{U} \mathbf{E}_y^T,

(13.2.5)

where

\mathbf{E}_x = \begin{bmatrix} 0 & 1 & 0 & \cdots & 0 & 0 \\ 0 & 0 & 1 & \cdots & 0 & 0 \\ & & & \ddots & & \\ 0 & 0 & 0 & \cdots & 1 & 0 \end{bmatrix}

(13.2.6)

is $(m-1)\times (m+1)$ , and $\mathbf{E}_y$ is analogous but of size $(n-1)\times (n+1)$ . The left multiplication in (13.2.5) deletes the first and last row of $\mathbf{U}$ , and the right multiplication deletes its first and last column.

The extension operator is a bit more awkward to write out. It stars with appending rows and columns of zeros around the border of a matrix $\mathbf{W}$ of interior values:

\tilde{\mathbf{U}} = \mathbf{E}_x^T \mathbf{W} \mathbf{E}_y

(13.2.7)

We can then modify the new zero values to reflect the boundary conditions, via

\mathbf{U} = \operatorname{extend}(\mathbf{W}) = \tilde{\mathbf{U}} + \mathbf{G},

(13.2.8)

Finally, we combine extending and chopping with the need to reshape between vectors and matrices. The function

\operatorname{unpack}(\mathbf{w}) = \operatorname{extend}(\operatorname{unvec}(\mathbf{w}))

(13.2.9)

converts a vector of unknowns (i.e., interior values) into a matrix of grid values, including the boundary values. The function

\operatorname{pack}(\mathbf{U}) = \operatorname{vec}(\operatorname{chop}(\mathbf{U}))

(13.2.10)

reverses the transformation, which is needed after the time derivative is computed on the full grid.

Example 13.2.3 (Advection-diffusion equation in 2D)

We solve an advection-diffusion problem, $u_t + u_x = 1 + \epsilon(u_{xx} + u_{yy})$ on the square $[-1,1]^2$ , with $u=0$ on the boundary. The outline of our approach is based on Function 11.5.2 for parabolic PDEs in one space dimension.

Julia

MATLAB

Python

Example 13.2.3

The first step is to define a discretization of the domain.

m, n = 50, 36
x, Dx, Dxx = FNC.diffcheb(m, [-1, 1])
y, Dy, Dyy = FNC.diffcheb(n, [-1, 1])
mtx, X, Y, _ = FNC.tensorgrid(x, y)
U₀ = mtx( (x, y) -> (1 + y) * (1 - x)^4 * (1 + x)^2 * (1 - y^4) );

We define functions extend and chop to deal with the Dirichlet boundary conditions.

chop = U -> U[2:m, 2:n]
z = zeros(1, n-1)
extend = W -> [zeros(m+1) [z; W; z] zeros(m+1)]

#27 (generic function with 1 method)

Next, we define the pack and unpack functions, using another call to Algorithm 13.2.1 to get reshaping functions for the interior points.

_, _, _, unvec, _ = FNC.tensorgrid(x[2:m], y[2:n])
pack = U -> vec(chop(U))
unpack = w -> extend(unvec(w))

#31 (generic function with 1 method)

Now we can define and solve the IVP using a stiff solver.

function dw_dt(w, ϵ, t)
    U = unpack(w)
    Ux, Uxx = Dx * U, Dxx * U
    Uyy = U * Dyy'
    du_dt = @. 1 - Ux + ϵ * (Uxx + Uyy)
    return pack(du_dt)
end

IVP = ODEProblem(dw_dt, pack(U₀), (0.0, 2), 0.05)
w = solve(IVP, Rodas4P());

When we evaluate the solution at a particular value of $t$ , we get a vector of the interior grid values. The same unpack function above converts this to a complete matrix of grid values.

U = t -> unpack(w(t))
contour(x, y, U(0.5)';
    fill=true,  color=:blues,  levels=20, l=0,
    aspect_ratio=1,  xlabel=L"x",  ylabel=L"y",
    title="Solution at t = 0.5")

Source

anim = @animate for t in 0:0.02:2
    U = unpack(w(t))
    surface(x, y, U';
        layout=(1, 2),  size=(640, 320),
        xlabel=L"x",  ylabel=L"y",  zaxis=((0, 2.3), L"u(x,y)"),
        color=:blues,  alpha=0.66,  clims=(0, 2.3), colorbar=:none,
        title="Advection-diffusion",  dpi=150)
    contour!(x, y, U'; 
        levels=24, 
        aspect_ratio=1,  subplot=2, 
        xlabel=L"x",  ylabel=L"y",
        color=:blues,  clims=(0, 2.3),  colorbar=:none,
        title=@sprintf("t = %.2f", t))
end
mp4(anim, "figures/2d-advdiff.mp4");

Example 13.2.3

The first step is to define a discretization of the domain and the initial state.

m = 50;  n = 40;
[x, Dx, Dxx] = diffcheb(m, [-1, 1]);
[y, Dy, Dyy] = diffcheb(n, [-1, 1]);
[mtx, X, Y] = tensorgrid(x, y);
u_init = @(x, y) (1+y) .* (1-x).^4 .* (1+x).^2 .* (1-y.^4);

We define functions extend and chop to deal with the Dirichlet boundary conditions.

chop = @(U) U(2:m, 2:n);
z = zeros(1, n-1);
extend = @(W) [ zeros(m+1, 1) [z; W; z] zeros(m+1, 1)];

Now we define the pack and unpack functions, using another call to Algorithm 13.2.1 to get reshaping functions for the interior points.

[~, ~, ~, vec, unvec] = tensorgrid(x(2:m), y(2:n));
pack = @(U) vec(chop(U));
unpack = @(w) extend(unvec(w));

Now we can define and solve the IVP using a stiff solver.

f13_2_advdiff.m

function du_dt = timederiv(t, u, p)
    [ep, Dx, Dxx, Dy, Dyy, pack, unpack] = p{:};
    U = unpack(u);
    Uxx = Dxx * U;  Uyy = U * Dyy'; 
    dU_dt = 1 - Dx * U + ep * (Uxx + Uyy);  % PDE
    du_dt = pack(dU_dt);
end

odefun = @(t, u) f13_2_advdiff(t, u, {0.05, Dx, Dxx, Dy, Dyy, pack, unpack});
u_init = pack(mtx(u_init));
sol = ode15s(odefun, [0, 2], u_init);

When we evaluate the solution at a particular value of $t$ , we get a vector of the interior grid values. The unpack converts this to a complete matrix of grid values.

U = @(t) unpack(deval(sol, t));

clf,  pcolor(X', Y', U(0.5)')
clim([0, 2.3]), shading interp
axis equal,  colormap(sky), colorbar
title('Advection-diffusion at t = 0.5')  
xlabel('x'),  ylabel('y')

Source

hold on
vid = VideoWriter("figures/2d-advdiff.mp4","MPEG-4");
vid.Quality = 85;
open(vid);
title("Advection-diffusion in 2d")
for t = linspace(0, 2, 81)
    cla, pcolor(X', Y', U(t)')
    shading interp
    str = sprintf("t = %.2f", t);
    text(-1.5, 0.75, str, fontsize=14);
    writeVideo(vid, frame2im(getframe(gcf)));
end
close(vid)

Example 13.2.3

The first step is to define a discretization of the domain.

m, n = 50, 36
x, Dx, Dxx = FNC.diffcheb(m, [-1, 1])
y, Dy, Dyy = FNC.diffcheb(n, [-1, 1])
mtx, X, Y, _, _, _ = FNC.tensorgrid(x, y)
u_init = lambda x, y: (1 + y) * (1 - x)**4 * (1 + x)**2 * (1 - y**4)

We define functions extend and chop to deal with the Dirichlet boundary conditions.

def chop(U):
    return U[1:-1, 1:-1]

def extend(W):
    U = zeros((m+1, n+1))
    U[1:-1, 1:-1] = W
    return U

Now we define the pack and unpack functions, using another call to Algorithm 13.2.1 to get reshaping functions for the interior points.

_, _, _, vec, unvec, _ = FNC.tensorgrid(x[1:-1], y[1:-1])
pack = lambda U: vec(chop(U))          # restrict to interior, then vectorize
unpack = lambda w: extend(unvec(w))    # reshape, then extend to boundary

Now we can define and solve the IVP using a stiff solver.

ep = 0.05
def dw_dt(t, w):
    U = unpack(w)
    Uyy = Dxx @ U
    Uxx = U @ Dyy.T 
    dU_dt = 1 - Dx @ U + ep * (Uxx + Uyy)
    return pack(dU_dt)

U0 = mtx(u_init)
sol = solve_ivp(dw_dt, (0, 2), pack(U0), method="BDF", dense_output=True)

When we evaluate the solution at a particular value of $t$ , we get a vector of the interior grid values. The unpack converts this to a complete matrix of grid values.

U = lambda t: unpack(sol.sol(t))    # matrix-valued function of time

pcolormesh(X.T, Y.T, U(0.5).T, cmap="Blues", shading="gouraud")
colorbar()
xlabel("$x$"),  ylabel("$y$")
axis("equal"),  title("Solution at t=0.5");

mx = max([max(U(t)) for t in linspace(0, 2, 21)])
fig, ax = subplots()
obj = ax.pcolormesh(X.T, Y.T, U(0).T, vmin=0, vmax=2, cmap="Blues", shading="gouraud")
time_text = ax.text(0.05, 0.9, '', transform=ax.transAxes)
ax.set_xlabel("$x$"),  ax.set_ylabel("$y$")
ax.set_aspect("equal")
ax.set_title("Advection-diffusion in 2d")
def snapshot(t):
    global obj
    obj.remove()
    obj = ax.pcolormesh(X.T, Y.T, U(t).T, vmin=0, vmax=mx, cmap="Blues", shading="gouraud")
    time_text.set_text(f"t = {t:.2f}")

anim = animation.FuncAnimation(fig, snapshot, frames=linspace(0, 2, 81))
anim.save("figures/advdiff-2d.mp4", fps=30)
close()

13.2.4Wave equation¶

The wave equation introduces a little additional complexity. First, we write the 2D wave equation $u_{tt}=c^2(u_{xx}+u_{yy})$ in first-order form as

\begin{split} u_t &= v, \\ v_t &= c^2(u_{xx}+u_{yy}). \end{split}

(13.2.11)

Typical boundary conditions are to prescribe $u$ on the boundary and let $v$ be unspecified.

Now the grid functions are a pair of matrices $\mathbf{U}(t)$ and $\mathbf{V}(t)$ . We need to chop $\mathbf{U}$ to an interior $\mathbf{W}$ and extend back using boundary data. Note that the IVP unknowns $\mathbf{W}$ and $\mathbf{V}$ have different sizes, so there are two separate reshaping operations involved. All of these details are handled within the pack and unpack functions we create.

Example 13.2.4 (Wave equation in 2D)

We solve the wave equation with $c=1$ on the square $[-2,2]\times[-2,2]$ , with $u=0$ on the boundary.

Julia

MATLAB

Python

Example 13.2.4

We start with the discretization and initial condition.

m, n = 40, 40
x, Dx, Dxx = FNC.diffcheb(m, [-2, 2])
y, Dy, Dyy = FNC.diffcheb(n, [-2, 2])
mtx, X, Y, _ = FNC.tensorgrid(x, y)
U₀ = mtx( (x, y) -> (x + 0.2) * exp(-12 * (x^2 + y^2)) )
V₀ = zeros(size(U₀));

We need to define chopping and extension for the $u$ component. This looks the same as in Example 13.2.3.

chop = U -> U[2:m, 2:n]
extend = U -> [zeros(m+1) [zeros(1, n-1); U; zeros(1, n-1)] zeros(m+1)]

#39 (generic function with 1 method)

While vec is the same for both the interior and full grids, the unvec operation is defined differently for them.

_, _, _, unvec_v, _ = FNC.tensorgrid(x, y)
_, _, _, unvec_u, _ = FNC.tensorgrid(x[2:m], y[2:n])
N = (m-1) * (n-1)    # number of interior unknowns

pack = (U, V) -> [vec(chop(U)); vec(V)]
unpack = w -> (extend(unvec_u(w[1:N])), unvec_v(w[N+1:end]));

We can now define and solve the IVP. Since this problem is hyperbolic, a nonstiff integrator is faster than a stiff one.

function dw_dt(w, c, t)
    U, V = unpack(w)
    du_dt = V
    dv_dt = c^2 * (Dxx * U + U * Dyy')
    return pack(du_dt, dv_dt)
end

IVP = ODEProblem(dw_dt, pack(U₀, V₀), (0, 4.0), 1)
sol = solve(IVP, Tsit5())
U = t -> unpack(sol(t))[1];

Source

anim = @animate for t in 0:4/100:4
    Ut = U(t)
    surface(x, y, Ut';
        layout=(1, 2), size=(640, 320),
        xlabel=L"x",  ylabel=L"y",  zaxis=((-0.1, 0.1), L"u(x,y)"),
        color=:redsblues,  alpha=0.66,  clims=(-0.1, 0.1), colorbar=:none,
        title="Wave equation",  dpi=150)
    contour!(x, y, Ut'; 
        levels=24,  subplot=2, 
        aspect_ratio=1,
        xlabel=L"x",  ylabel=L"y",
        color=:redsblues,  clims=(-0.1, 0.1), 
        colorbar=:none,  title=@sprintf("t = %.2f", t))
end
mp4(anim, "figures/2d-wave.mp4");

Example 13.2.4

We start with the discretization and initial condition.

m = 40; n = 42;
[x, Dx, Dxx] = diffcheb(m, [-2, 2]);
[y, Dy, Dyy] = diffcheb(n, [-2, 2]);
[mtx, X, Y] = tensorgrid(x, y);

u_init = @(x, y) (x+0.2) .* exp(-12*(x.^2 + y.^2));
U0 = mtx(u_init);
V0 = zeros(size(U0));

We need to define chopping and extension for the $u$ component. This looks the same as in Example 13.2.3.

chop = @(U) U(2:m, 2:n);
z = zeros(1, n-1);
extend = @(U) [ zeros(m+1, 1) [z; U; z] zeros(m+1, 1)];

The vec and unvec operations are different for the interior-only grid and the full grid.

[~, ~, ~, vec_v, unvec_v] = tensorgrid(x, y);
[~, ~, ~, vec_u, unvec_u] = tensorgrid(x(2:m), y(2:n));
N = (m-1) * (n-1);

pack = @(U, V) [vec_u(chop(U)); vec_v(V)];
unpack = @(u) f13_2_wave_unpack(u, N, unvec_u, unvec_v, extend);

f13_2_wave_unpack.m

function [U, V] = unpack(w, N, unvec_u, unvec_v, extend)
    U = extend( unvec_u(w(1:N)) );
    V = unvec_v(w(N+1:end));
end

We can now define and solve the IVP. Since this problem is hyperbolic, a nonstiff integrator is faster than a stiff one.

f13_2_wave.m

function dw_dt = timederiv(t, w, p)
    [Dxx, Dyy, pack, unpack] = p{:};
    [U, V] = unpack(w);
    dU_dt = V;
    dV_dt = Dxx * U + U * Dyy';
    dw_dt = pack(dU_dt, dV_dt);
end

odefun = @(t, u) f13_2_wave(t, u, {Dxx, Dyy, pack, unpack});
ivp.InitialTime = 0;
z_init = pack(U0, V0);
ivp_sol = ode45(odefun, [0, 4], z_init);
sol = @(t) deval(ivp_sol, t);

clf
[U, V] = unpack(sol(0.5));
pcolor(X', Y', U')
axis equal,  clim([-0.1, 0.1])
colormap(redsblues),  shading interp
xlabel("x"),  ylabel("y")
title("Wave equation at t = 0.5")

Source

hold on
vid = VideoWriter("figures/2d-wave.mp4","MPEG-4");
vid.Quality = 85;
open(vid);
title("Wave equation in 2d")
for t = linspace(0, 4, 121)
    [U, V] = unpack(sol(t));
    cla, pcolor(X, Y, U)
    shading interp
    str = sprintf("t = %.2f", t);
    text(-3, 1.75, str, fontsize=14);
    writeVideo(vid, frame2im(getframe(gcf)));
end
close(vid)

Example 13.2.4

We start with the discretization and initial condition.

m, n = 40, 42
x, Dx, Dxx = FNC.diffcheb(m, [-2, 2])
y, Dy, Dyy = FNC.diffcheb(n, [-2, 2])
mtx, X, Y, _, _, _ = FNC.tensorgrid(x, y)

U0 = mtx(lambda x, y: (x + 0.2) * exp(-12 * (x**2 + y**2)))
V0 = zeros(U0.shape)

We need to define chopping and extension for the $u$ component. This looks the same as in Example 13.2.3.

def extend(U):
    UU = zeros((m+1, n+1))
    UU[1:-1, 1:-1] = U
    return UU

def chop(U):
    return U[1:-1, 1:-1]

_, _, _, vec_v, unvec_v, _ = FNC.tensorgrid(x, y)
_, _, _, vec_u, unvec_u, _ = FNC.tensorgrid(x[1:-1], y[1:-1])
N = (m-1) * (n-1)    # number of interior points

def pack(U, V): 
    return hstack([vec_u(chop(U)), vec_v(V)])

def unpack(w):
    U = extend(unvec_u(w[:N]))
    V = unvec_v(w[N:])
    return U, V

We can now define and solve the IVP. Since this problem is hyperbolic, a nonstiff integrator is faster than a stiff one.

def dw_dt(t, w):
    U, V = unpack(w)
    dU_dt = V
    dV_dt = Dxx @ U + U @ Dyy.T
    return pack(dU_dt, dV_dt)

from scipy.integrate import solve_ivp
sol = solve_ivp(dw_dt, (0, 4), pack(U0, V0), method="RK45", dense_output=True)
U = lambda t: unpack(sol.sol(t))[0]

fig, ax = subplots()
obj = ax.pcolormesh(X, Y, U(0), vmin=-0.1, vmax=0.1, cmap="RdBu", shading="gouraud")
time_text = ax.text(0.05, 0.9, '', transform=ax.transAxes)
ax.set_xlabel("$x$"),  ax.set_ylabel("$y$")
ax.set_aspect("equal")
ax.set_title("Wave equation in 2d")

def snapshot(t):
    global obj
    obj.remove()
    obj = ax.pcolormesh(X, Y, U(t), vmin=-0.1, vmax=0.1, cmap="RdBu", shading="gouraud")
    time_text.set_text(f"t = {t:.2f}")

anim = animation.FuncAnimation(fig, snapshot, frames=linspace(0, 4, 91))
anim.save("figures/wave-2d.mp4", fps=30);
close()

13.2.5Exercises¶

Exercise 13.2.5

From Maxwell’s equations we can find a way to convert the wave equation to a first-order form that, unlike (13.2.11), uses only first-order derivatives in space:

\begin{split} u_t &= c^2(v_y - w_x),\\ v_t &= u_y, \\ w_t &= -u_x, \end{split}

(13.2.12)

subject to $u=0$ on the boundary.

(a) ✍ Show that a solution of (13.2.12) satisfies $u_t=c^2(u_{xx}+u_{yy})$ .

(b) ⌨ Solve (13.2.12) with $c=2$ in the rectangle $x\in[-3,3]$ , $y\in[-1,1]$ , $u(x,y,0) = \exp(x-x^2)(9-x^2)(1-y^2)$ , and $v=w=0$ at $t=0$ . Use $m=50$ for $x$ and $n=25$ for $y$ , solve for $0\le t \le 6$ , and make an animation of the solution.

Preface

Tensor-product discretizations

Preface

Laplace and Poisson equations