Rendering Solutions to the Wave Equation in 3D

published 06/09/2024 • 2m reading time • 29 views

The double slit experiment

Since writing wave-sim, my GPU accelerated simulator for the 2D discretized wave equation, I have wanted to visualize the solutions in three dimensions, rather than just using a 2d image. In this article I cover the changes to wave-sim I made and the Julia code I wrote to actually render the videos using Makie.jl.

Modifying the Solver

I could either add a new renderer to my wave simulator that just shows a 3D view in real time, but that would be a lot of work. I may eventually do that, but for now I opted to use an existing plotting library. All I have to do with wave-sim is get it to output the simulation state for every tick.

Because the simulation state Buffer must be accessible to the compute and fragment shaders, it has the STORAGE buffer usage, which is incompatible with MAP_READ, which is needed to download its contents back to main system memory. To get around this, we create a new ‘staging buffer’, which is just a buffer that has the COPY_DST and MAP_READ buffer usages. We can then copy the contents of the state buffer to the staging buffer and then download the staging buffer.

Below is the code for copying the buffers, the offset is used because the current and previous two simulation states are stored in the same buffer, but we only want the current.

let offset = (self.tick % 3) * (self.size.0 * self.size.1 * 4) as u64;
encoder.copy_buffer_to_buffer(
    &self.states,
    offset,
    &self.staging_buffer,
    0,
    self.staging_buffer.size(),
);

We can then wait for the copy to complete, then download the buffer and save it to disk. Each file starts with two u32s for the state width and height, then the contents of the buffer (lots of f32s).

Code to download the buffer

pub fn download_buffer(buffer: &Buffer, gc: &GraphicsContext) -> Vec<u8> {
    let slice = buffer.slice(..);

    let (tx, rx) = crossbeam_channel::bounded(1);
    slice.map_async(MapMode::Read, move |_| tx.send(()).unwrap());

    gc.device.poll(MaintainBase::Wait);
    rx.recv().unwrap();

    let data = slice.get_mapped_range().to_vec();
    buffer.unmap();

    data
}

Code to write the file

let mut data = Vec::with_capacity(8 + snapshot.size() as usize);
let size = self.simulation.get_size();
data.extend_from_slice(&size.0.to_le_bytes());
data.extend_from_slice(&size.1.to_le_bytes());
data.extend_from_slice(&download_buffer(snapshot, gc));

fs::write(format!("states/state-{}.bin", next_id()), data).unwrap();

There are a lot of opportunities for optimization, like holding a few states in the staging buffer before downloading, not blocking the main thread while downloading, and not writing a separate file for each state, but it’s not so bad waiting for it to process just a few thousand ticks.

Rendering in 3D

I started this whole 3d rendering side quest because I remembered about Makie.jl, a visualization and plotting library for the Julia programming language. It allows you to very easily make beautiful plots and animations, it supports multiple rendering backend and for this project I will use GLMakie.

First, I’ll just define some constants.

STATE_PATH = "states"
WIDTH = 1440
HEIGHT = 900
Z_SCALE = 250

Then get all the states in the state path, and make sure they are sorted by the id number in the filename.

states = readdir(STATE_PATH)
println("Found $(length(states)) states")

sort!(states, by = x -> parse(Int, split(x, "-")[2][1:end-4]))

Then we can make the surface plot with GLMakie, using an Observable to make the plot data reactive.

GLMakie.activate!()
GLMakie.closeall()

data = Observable(load_state(states[1]))

fig = Figure(resolution = (1920, 1080))
axis = Axis3(fig[1, 1], aspect = (WIDTH, HEIGHT, Z_SCALE), azimuth = 6.275pi, elevation = 0.16pi, xlabel = "x", ylabel = "y", zlabel = "z")
surface!(axis, data, colormap = :viridis)

Finally loop through each state, load it from its file and update the data observable.

function load_state(name)
    state = read(STATE_PATH * "/" * name)
    return reshape(reinterpret(Float32, state[9:end]), (WIDTH, HEIGHT))[:, 5:end]
end

record(fig, "3d_plot.mp4", 1:length(states)) do frame
    print("\r$(round(Int, 100 * frame / length(states)))%")
    data[] = load_state(states[frame])
end

println("\nDone!")

Examples

I have rendered videos of the following demos: