FDTD for Room Acoustics, Part 2: Implementing the 2D FDTD Scheme in Dr.Jit

Before going into the implementation, let us first look at what the scheme produces. The animation shows a Gaussian impulse propagating through a two-dimensional grid and reflecting from the domain boundaries. Since this first setup does not include absorption, the wave does not intentionally lose energy over time. The field may look increasingly diffuse after several reflections, but this is caused by the repeated superposition of reflected wavefronts, not by a damping mechanism.

Implementation using Dr.Jit in Python

For the implementation of the FDTD update formula for the wave equation we derived in Part 1 of this series, we will use Dr.Jit, a Python library for writing parallel programs that can run on the CPU or GPU. It also supports automatic differentiation, which makes it interesting for optimization problems.

Wave-based acoustic solvers are computationally intensive, so parallelisation is essential for acceptable performance at larger scales. The differentiability of Dr.Jit also opens up interesting possibilities: instead of only simulating a given scene, we can later ask how a scene or an initial condition would have to change in order to produce a desired acoustic result.

Compared with writing kernels directly in CUDA, Dr.Jit lets us write the program at a higher level. It then converts parts of the program into parallel kernels just in time and runs them on the selected hardware.

In the first part of the code we are dealing with library imports and check what backend and version dr.jit is using. The debugging flag can be very helpful in case we are running into trouble.

import drjit as dr
from drjit.cuda import Float, UInt
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, PillowWriter
from matplotlib.colors import TwoSlopeNorm
from pathlib import Path

print("backend:", dr.backend_v(Float),"version:", dr.__version__)

dr.set_flag(dr.JitFlag.Debug, True)

In our first function we define a spatial Gaussian initial pressure field in the center of the domain. For each grid point, we compute the squared distance r2 from the domain midpoint and evaluate a Gaussian profile. The parameter sigma controls the spatial width of the bump; here it is set to 3% of the smaller domain dimension so that the source remains compact even for rectangular domains.

The purpose of this function is not to model a physical loudspeaker yet. It simply creates a smooth, localized disturbance so that the finite-difference scheme has something to propagate.

# Gaussian initial condition centered in the domain
def gaussian(xL,xR,yB,yT,dx,dy,iy,jx):
    # coordinates of the grid nodes
    x = xL + Float(jx) * dx
    y = yB + Float(iy) * dy
    # define center of physical domain
    x0, y0 = 0.5*(xL+xR), 0.5*(yB+yT)    
    sigma = 0.03 * min(xR-xL, yT-yB)
    r2 = (x-x0)*(x-x0) + (y-y0)*(y-y0)   
    # compute initial Gaussian (negative exponential function)
    # at domain center r2 will be zero, the exponential 1. 
    u0 = dr.exp(-r2 / (2.0 * sigma * sigma))  
    return u0

Now we get to the actual main function run_sim_frames. We will break it down into several subsections for this post, but keep in mind that this is actually one big function.

Definition inputs

The function arguments define the grid resolution nx, ny, the physical domain bounds xL, xR, yB, yT, the wave speed c, the simulation end time te, the CFL safety factor, the frame recording interval, and the initial condition.

def run_sim_frames(nx=101, ny=51, 
                   xL=0.0, xR=2.0, yB=0.0, yT=1.0,
                   c=343.0, te=0.0003, cfl_scale=0.9,
                   record_every=2,source="gaussian"):  

Domain, grid, and timestep

This block turns the physical domain and desired end time into the discrete grid and time step used by the solver. The number of grid points determines the spacing Δx and Δy, while the timestep is chosen from the CFL limit derived in Part 1:

We multiply this maximum timestep by the safety factor cfl_scale, and then compute how many update steps fit into the requested simulation time te.

N = nx * ny
dx = (xR - xL) / (nx - 1)
dy = (yT - yB) / (ny - 1)

dt_max = 1.0 / (c * dr.sqrt(1.0/(dx*dx) + 1.0/(dy*dy)))
dt = cfl_scale * dt_max
nsteps = int(te / dt)

Flattened indexing

One difficulty in this Dr.Jit implementation is that we express 2D grids of field values with one-dimensional arrays. Hence, we define a basic variable k with a length N to create a 1D array of indices for the 2D grid. iy describes row indices in y direction, e.g. for nx=2, ny=3 we get an index array of [0, 0, 1, 1, 2, 2]. Similarly jx describes column indices in x direction, e.g. for nx=2, ny=3 we get an index array of [0, 1, 0, 1, 0, 1].

k  = dr.arange(UInt, N) 
iy = k // nx
jx = k - iy * nx

Then we define a boolean mask describing interior nodes which we will use later in the update loop.

is_interior = (iy > 0) & (iy < ny - 1) & (jx > 0) & (jx < nx - 1)

jxL, jxR, iyB and iyT describe arrays of left, right, bottom and top point offset indices of the 5 point stencil. For stencils with their center along the border whose neighbour points would be already off grid we clamp them to the border point. As we set out to work with a homogeneous Dirichlet boundary condition and thus the boundary gets set to zero in every iteration step, the clamped values do not contribute and we can forego adding additional grid points, commonly referred to as ghost cells.

jxL = dr.select(jx > 0,      jx - 1, jx)
jxR = dr.select(jx + 1 < nx, jx + 1, jx)
iyB = dr.select(iy > 0,      iy - 1, iy)
iyT = dr.select(iy + 1 < ny, iy + 1, iy)

Stencil indices

Then these offsets are applied to our row and column indices arrays so we get arrays of left, right, top and bottom neighbours for every stencil center index.

idxL = iy * nx + jxL
idxR = iy * nx + jxR
idxB = iyB * nx + jx
idxT = iyT * nx + jx

CFL coefficients and time buffers

The factors cx2 and cy2 are the squared Courant numbers in the x- and y-directions. They multiply the second-difference terms in the update equation. Further we set up 3 buffer arrays for the values u for previous, current and next timestep.

cx2 = (c * dt / dx) ** 2
cy2 = (c * dt / dy) ** 2

u_prev = dr.zeros(Float, N)
u_cur  = dr.zeros(Float, N)
u_next = dr.zeros(Float, N)

Initialisation of current pressure field

We fill the current field with the Gaussian initial condition and use the interior mask to keep the boundary nodes at zero.

As a simplification, we initialize the previous time level as a detached copy of the current one, u⁻¹ = u⁰. This is not the fully Taylor-consistent zero-initial-velocity start, but it keeps the first implementation simple. We will revisit this when refining the source and initialization model.

The first NumPy frame is stored immediately so that the animation includes the initial condition.

if source == "gaussian":
    init = gaussian(xL,xR,yB,yT,dx,dy,iy,jx)
    u_cur = dr.select(is_interior, init, 0.0)
    u_prev = dr.detach(u_cur)
# placeholder for other sourcetypes
# elif source == "placeholder":   
else:
    print('No valid sourcename provided')

frames = [np.array(u_cur, dtype=np.float32).reshape(ny, nx)] 

The update loop

Now we come to the actual update loop over the timesteps defined above. It is worth at this point to state again the update formula from Part 1 here:

The index arrays let us gather the left, right, bottom, and top neighbours for every grid point. Together with the current center value uc and the previous time value up, these arrays are inserted directly into the update formula to compute un, the next time level.

After computing u_next, we rotate the buffers: the current field becomes the previous field, and the newly computed field becomes the current one.

Then we write the latest values into our numpy array. Appending a NumPy frame forces Dr.Jit to evaluate the computation graph and copy the evaluated data from Dr.Jit-managed memory into a NumPy array.

Finally after all timesteps have been computed we return the frames, the timestep size and the number of timesteps from the function back to the caller.

for t in range(nsteps):
    up = u_prev
    uc = u_cur
    ul = dr.gather(Float, u_cur, idxL)
    ur = dr.gather(Float, u_cur, idxR)
    ub = dr.gather(Float, u_cur, idxB)
    ut = dr.gather(Float, u_cur, idxT)

    un = 2.0*uc - up + cx2*(ul - 2.0*uc + ur) + cy2*(ub - 2.0*uc + ut)

    u_next = dr.select(is_interior, un, 0.0)
    u_prev, u_cur = u_cur, u_next

    if (t + 1) % record_every == 0:
        frames.append(np.array(u_cur, dtype=np.float32).reshape(ny, nx))

return np.stack(frames), dt, nsteps

Calling the simulation

# domain dimensions
xL_s, xR_s = 0.0, 2.0
yB_s, yT_s = 0.0, 1.0   # rectangle if you want

frames, dt, nsteps = run_sim_frames(nx=101, ny=51,
                                    xL=xL_s, xR=xR_s, 
                                    yB=yB_s, yT=yT_s,
                                    te=0.003,
                                    source="gaussian")

print("frames:", frames.shape)
print("mean |delta| between frames:", np.mean(np.abs(np.diff(frames, axis=0))))

Visualisation

For visualization, we fix a symmetric colour scale around zero. The value v is chosen from a high quantile of the absolute field values, which avoids a few extreme values dominating the colour range. TwoSlopeNorm then maps negative and positive values symmetrically around zero, so that blue and red represent waves of opposite sign. Let's first plot the last frame of our simulation with following code for a first inspection.

v = np.quantile(np.abs(frames), 0.995)
norm = TwoSlopeNorm(vmin=-v, vcenter=0.0, vmax=v)

fig, ax = plt.subplots(figsize=(6,3.4))
im = ax.imshow(frames[-1], origin="lower", cmap="RdBu_r",norm=norm, extent=[xL_s, xR_s, yB_s, yT_s])

ax.set_aspect("equal")
ax.set_xlabel("x"); ax.set_ylabel("y")
ax.set_title("2D FDTD — Dirichlet boundary")
plt.show()

Finally we can create an animation by calling above plot function for every output frame and saving it to disk. For short animations gifs are fine, but for longer ones compressed video files are the better choice. We stick for now with gifs.

def update(i):
    im.set_data(frames[i])
    return (im,)

gifname = "fdtd_2D_Part2_GaussInit_Dirichlet.gif"
outdir = Path("../outputs")
outdir.mkdir(exist_ok=True)

gifpath = outdir / gifname

ani = FuncAnimation(fig, update, frames=frames.shape[0], interval=40, blit=False)
ani.save(gifpath, writer=PillowWriter(fps=25))
plt.close(fig)

print(rf"saved {gifpath}")

Results

The animation at the beginning of the post shows the main qualitative behaviour of this first implementation. The wavefront reflects repeatedly at the domain boundaries, and the field becomes visually more diffuse as reflected wavefronts overlap.

As discussed in Part 1, the Dirichlet boundary condition is more akin to a fixed membrane border than to a room wall. The phase change at the boundary impact is a visible sign of this: a positive maximum is reflected as a negative minimum, and vice versa. This is not yet a realistic room-acoustic boundary model, so in the next parts we will move toward rigid-wall behaviour and eventually absorption.

Code

The complete notebook for this part is available here: [GitHub link]. The code in the article is split into sections for explanation, while the notebook contains the runnable version in one place.