Leapfrog Implementation

We just discussed the leapfrog method, and why it superior to a forward difference method. The implementation is not too difficult, but it is worth sketching here.

To refresh our memory, lets rewrite the expressions for the time evolved wave function. For the real part of the wave function we have

Re (Ψ (x, t + Δ t)) \approx Re (Ψ (x, t - Δ t)) - \frac{Im (Ψ (x + Δ x, t)) - 2 Im (Ψ (x, t)) + Im (Ψ (x - Δ x, t))}{{(Δ x)}^{2}} Δ t + 2 V (x) Im (Ψ (x, t)) Δ t

and to evolve the imaginary component of the wave function in time

Im (Ψ (x, t + Δ t)) \approx Im (Ψ (x, t - Δ t)) + \frac{Re (Ψ (x + Δ x, t)) - 2 Re (Ψ (x, t)) + Re (Ψ (x - Δ x, t))}{{(Δ x)}^{2}} Δ t - 2 V (x) Re (Ψ (x, t)) Δ t

Remember that the leapfrog involved another array of wave function values from $Ψ (x, t - Δ t)$ , so we need to add another array to our shader.

The Leapfrog method adds

Ψ (x, t - Δ t)

and hence another array to our calculations.

Start by adding the new wave function array.


    // Group 1, changes on each iteration
    // Older wave function at t-Δt.
    @group(1) @binding(0) var<storage, read> oldWaveFunction : array<vec2f>;
    // Current wave function at t.
    @group(1) @binding(1) var<storage, read> waveFunction : array<vec2f>;
    // The updated wave function at t+Δt.
    @group(1) @binding(2) var<storage, read_write> updatedWaveFunction : array<vec2f>;

The thread identity directly determines the target wave function element.


    let index = global_id.x;

Within the computations we use ${Δ x}^{2}$ without the factor of 2 from earlier implementations. This has been canceled against the common factor of 2 with $Δ t$ .


    let dx2 = dx*dx;

Precompute the elements of the wave function expressions, this time including the oldWaveFunction value.


    let twoV = 2.0*parameters.potential[index];
    let oldWaveFunctionAtX = oldWaveFunction[index];
    let waveFunctionAtX = waveFunction[index];
    let waveFunctionAtXPlusDx = waveFunction[min(index+1, parameters.xResolution-1)];
    let waveFunctionAtXMinusDx = waveFunction[max(index-1, 0)];

Then we have a pretty straightforward representation of our update equations in code.


    updatedWaveFunction[index].x = oldWaveFunctionAtX.x
                                   - ((waveFunctionAtXPlusDx.y - 2.0*waveFunctionAtX.y + waveFunctionAtXMinusDx.y)
                                      / dx2 - twoV*waveFunctionAtX.y) * parameters.dt;

    updatedWaveFunction[index].y = oldWaveFunctionAtX.y
                                   + ((waveFunctionAtXPlusDx.x - 2.0*waveFunctionAtX.x + waveFunctionAtXMinusDx.x)
                                      / dx2 - twoV*waveFunctionAtX.x) * parameters.dt;

These initial conditions exploded using the forward difference. The leapfrog method is far more stable.

We need to do a bit of setup before running our new shader. Start by allocating the third wave function storage buffer, which duplicates the setup for the prior two buffers.


    waveFunctionBuffer2 = device.createBuffer({
      label: "Wave function 2",
      size: 2*xResolution*Float32Array.BYTES_PER_ELEMENT,
      usage: debug ? GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC : GPUBufferUsage.STORAGE
    });

Now that we have our three buffers, we need to describe their use. As expected, the first two buffers are read only, while we write to the third.


    const waveFunctionBindGroupLayout = device.createBindGroupLayout({
      label: "Wave function data.",
      entries: [
        {
          binding: 0,
          visibility: GPUShaderStage.COMPUTE,
          buffer: {
            type: "read-only-storage"
          }
        },
        {
          binding: 1,
          visibility: GPUShaderStage.COMPUTE,
          buffer: {
            type: "read-only-storage"
          }
        },
        {
          binding: 2,
          visibility: GPUShaderStage.COMPUTE,
          buffer: {
            type: "storage"
          }
        }
      ]
    });

The natural next step is to associate specific wave function buffers with these bindings. In the earlier example, we ping-ponged back and forth between two buffers for $Ψ (x, t)$ and $Ψ (x, t + Δ t)$ . This time, we cycle among three buffers for $Ψ (x, t - Δ t)$ , $Ψ (x, t)$ and $Ψ (x, t + Δ t)$ .


    waveFunctionBindGroup = new Array(3);

Our first set of buffers makes the obvious mapping of bindings to wave function buffers. 0 => waveFunctionBuffer0, 1 => waveFunctionBuffer1, and 2 => waveFunctionBuffer2.


    waveFunctionBindGroup[0] = device.createBindGroup({
      layout: waveFunctionBindGroupLayout,
      entries: [
        {
          binding: 0,
          resource: {
            buffer: waveFunctionBuffer0
          }
        },
        {
          binding: 1,
          resource: {
            buffer: waveFunctionBuffer1
          }
        },
        {
          binding: 2,
          resource: {
            buffer: waveFunctionBuffer2
          }
        }
      ]
    });

We allocate two more bind groups so that invocations of the shader cycle through them so that on each successive invocation, $i$ , we cycle buffers so that $Ψ (x, t + Δ t) \mapsto Ψ (x, t)$ , $Ψ (x, t) \mapsto Ψ (x, t - Δ t)$ , and $Ψ (x, t - Δ t) \mapsto Ψ (x, t + Δ t)$ . Thus, every third iteration, we return to our original configuration.


i%3	binding 0: $Ψ (x, t - Δ t)$	binding 1: $Ψ (x, t)$	binding 2: $Ψ (x, t + Δ t)$
0	waveFunctionBuffer0	waveFunctionBuffer1	waveFunctionBuffer2
1	waveFunctionBuffer1	waveFunctionBuffer2	waveFunctionBuffer0
2	waveFunctionBuffer2	waveFunctionBuffer0	waveFunctionBuffer1

Initial values

In our initial case, we need wave function values for both $Ψ (x, t - Δ t)$ and $Ψ (x, t)$ so clearly we also need to adapt out initialization process. As a first approximation, we simply set both of these to the same values.

The initializer needs to write to two wave function arrays.


    // Initial wave function at t=0.
    @group(2) @binding(0) var<storage, read_write> waveFunction0 : array<vec2f>;
    @group(2) @binding(1) var<storage, read_write> waveFunction1 : array<vec2f>;

At the end of the shader, simply copy the value to the second wave function buffer.


    waveFunction0[index] = computePsi(index);
    waveFunction1[index] = waveFunction0[index];

This time there are two storage arrays as inputs, both of which are writable.


    waveFunctionBindGroupLayout = device.createBindGroupLayout({
      label: "Wave function layout",
      entries: [
        {
          binding: 0,
          visibility: GPUShaderStage.COMPUTE,
          buffer: {
            type: "storage"
          }
        },
        {
          binding: 1,
          visibility: GPUShaderStage.COMPUTE,
          buffer: {
            type: "storage"
          }
        }
      ]
    });

It is natural to fill in values for waveFunctionBuffer0 and waveFunctionBuffer1 from the FDTD shader.


    waveFunctionBindGroup = device.createBindGroup({
      layout: waveFunctionBindGroupLayout,
      entries: [
        {
          binding: 0,
          resource: {
            buffer: waveFunctionBuffer0
          }
        },
        {
          binding: 1,
          resource: {
            buffer: waveFunctionBuffer1
          }
        }
      ]});