Simulation Class: The Wave Equation in 2D

The Wave Equation in 2D

The 1D wave equation solution from the previous post is fun to interact with, and the logical next step is to extend the solver to 2D. It turns out that this is almost trivially simple, with most of the work going into making adjustments to display and interaction with the state arrays. As we'll see, we'll make almost no changes to the physics whatsoever - this will be an ongoing theme in this class, that simulation programming is mostly about input, output, iteration over arrays, and display, and involves very little actual physics (in terms of lines of code).

The first thing we'll do is change the state array from a 1D array to a 2D array. We'll also change notation so that X & Y represent the dimensions of the array, and Z represents the vertical dimension expressed by heights in the array. Our state initialization now looks like this (showing only the lines that have changed):

float WorldSize = 10.0;
int NX = 64;
int NY = 64;
int ArraySize = NX * NY;
float DXY = WorldSize / NX;

float LX = WorldSize;
float LY = WorldSize;
float LZ = WorldSize / 2.0;

We'll keep the State array declaration exactly the same, which means that we're still declaring the individual fields themselves as one-dimensional arrays. In order to index these 1D arrays from 2D coordinates, we'll need an indexing function:

int IX( int i, int j )
{
    return i + NX*j;
}

Next we implement the initialization of the state arrays, which is almost exactly the same as in the 1D case, except using a 2D iteration and 2D noise patterns. Otherwise, it's basically identical:

void SetInitialState() {
    noiseSeed( 0 );
    for (int j = 0; j < NY; ++j) {
        for (int i = 0; i < NX; ++i) {
            float worldX = 2341.17 + DXY * ( float )i;
            float worldY = 9911.44 + DXY * ( float )j;

            float n = 0.5 * snoise( worldX * 0.0625, worldY * 0.0625 ) +
                0.4 * snoise( worldX * 0.125, worldY * 0.125 ) +
                0.3 * snoise( worldX * 0.25, worldY * 0.25 ) +
                0.2 * snoise( worldX * 0.5, worldY * 0.5 );
            n = 1.0 - abs(n);
            n = (2.0 * n) - 1.0;

            State[StateHeight][IX(i,j)] = n;

            State[StateVel][IX(i,j)] = 0.0;
        }
    }
    EnforceHeightBoundaryConditions( StateHeight );
    EnforceNeumannBoundaryConditions( StateVel );
    StateCurrentTime = 0.0;

    CopyArray( StateHeight, StateHeightPrev );
    CopyArray( StateVel, StateVelPrev );
    InputActive = false;
}

Boundary Conditions in 2D

Our boundary conditions used to only be concerned with the first and last points of the linear 1D Array. Now we have to be concerned with the entire boundary, on all four edges. The Dirichlet and Neumann boundary functions are changed thus:

void EnforceDirichletBoundaryConditions( int io_a ) {
    for (int j = 0; j < NY; ++j) {
        if (j == 0 || j == (NY-1)) {
            for (int i = 0; i < NX; ++i) {
                State[io_a][IX(i,j)] = 0.0;
            }
        } else {
            State[io_a][IX(0,j)] = 0.0;
            State[io_a][IX(NX-1,j)] = 0.0;
        }
    }
}

void EnforceNeumannBoundaryConditions( int io_v ) {
    for (int j = 0; j < NY; ++j) {
        if (j == 0) {
            for (int i = 0; i < NX; ++i) {
                State[io_v][IX(i,0)] = State[io_v][IX(i,1)];
            }
        } else if (j == (NY-1)) {
            for (int i = 0; i < NX; ++i) {
                State[io_v][IX(i,NY-1)] = State[io_v][IX(i,NY-2)];
            }
        }

        State[io_v][IX(0,j)] = State[io_v][IX(1,j)];
        State[io_v][IX(NX-1,j)] = State[io_v][IX(NX-2,j)];
    }
}

Mouse Input in 2D

Input gathering is a fairly simple process as well - we didn't explicitly go over this in the previous example, so we'll provide a quick explanation here. In the top of the code, we introduce a small amount of input state:

boolean InputActive = false;
int InputIndexX = 0;
int InputIndexY = 0;
float InputHeight = 0;

And then we've got a function that gathers input when the mouse button is pressed:

void GetInput() {
    InputActive = false;
    if ( mousePressed && mouseButton == LEFT ) {
        float mouseCellX = mouseX / ( ( float )PixelsPerCell );
        float mouseCellY = mouseY / ( ( float )PixelsPerCell );

        int iX = ( int )floor( mouseCellX + 0.5 );
        int iY = ( int )floor( mouseCellY + 0.5 );
        if ( iX > 0 && iX < NX-1 && iY > 0 && iY < NY-1 ) {
            InputIndexX = iX;
            InputIndexY = iY;
            InputHeight = 1.5;
            InputActive = true;
        }
    }
}

This input data is used in only one place - when the height field boundary conditions are enforced, we set the input height explicitly (a Dirichlet condition). This is implemented via an explicit boundary handler for height:

void EnforceHeightBoundaryConditions( int io_h ) {
    EnforceNeumannBoundaryConditions(io_h);
    if ( InputActive ) {
        State[io_h][IX(InputIndexX, InputIndexY)] = 1.0;
    }
}

Displaying the 2D Results

Before we worry about any of the physics - we'll just work on getting the drawing of the data working. This is a revisiting of our basic philosophy and workflow - define your State, figure out how to display your State, and then work on your physics!

Processing offers a 2D image type that we can use to write to the pixels directly. Near the top of the file, we'll globally create this image:

PImage StateImage = createImage( NX, NY, RGB );

And then in the display function, we'll convert our height data into pixel colors. This is a somewhat aesthetic decision - our simulation data will range from -1.0 to 1.0 in height, generally, so we'll remap those values so that smaller numbers are drawn in dark blue and bigger numbers are drawn in white. Here's the draw-image code. We'll try this out temporarily without a time step just to make sure that the display works.

// Draw height field into the image.
void DrawHeightField( int i_field ) {
    float pixr, pixg, pixb;
    float d;
    color dark_blue = color(0.01, 0.01, 0.2);
    color light_blue = color(0.9, 0.9, 1.0);
    StateImage.loadPixels();
    for (int j = 0; j < NY; ++j) {
        for (int i = 0; i < NX; ++i) {
            float d = State[i_field][IX(i,j)];

            // Renormalize
            d = constrain( (d + 1.0) / 2.0, 0.0, 1.0 );

            // Add contrast
            d = pow( d, 8 );

            // Interpolate color.

            StateImage.pixels[IX(i,j)] = lerpColor(dark_blue, light_blue, d);
        }
    }
    StateImage.updatePixels();
    image( StateImage, 0, 0, width, height );
}

void draw() {
    background( 0.5 );

    // Simulate
    //GetInput();
    //TimeStepRK4( 1.0 / 24.0 );
    DrawHeightField(StateHeight);

    // Label.
    fill( 1.0 );
    text( "2D Wave Equation Init Only", 10, 30 );
}

Let's look at that in Processing, with just the drawing of the initial state, and nothing else:

Discretizing the 2D Wave Equation

In the 1D solver, we implemented our time-step by a somewhat lengthy discretization of the 1D wave equation, which we massaged and rearranged until it was suitable for use with our simple Jacobi iterative matrix solver. We'll follow exactly the same steps again. We'll take some bigger leaps this time around, but the method is exactly the same as in the 1D case, which you can review to understand each step broken down.

The 2D Wave Equation describes the motion of waves in a height field, and can be written simply as:

\[ \begin{equation} \frac{\partial^2 h}{\partial t^2} = c^2 \nabla^2 h \end{equation} \]

Which is the same as:

\[ \begin{equation} \frac{\partial^2 h}{\partial t^2} = c^2 \Delta h \end{equation} \]

Which are both just different shorthand notations for the following differential equation:

\[ \begin{equation} \frac{\partial^2 h}{\partial t^2} = c^2 (\frac{\partial^2 h}{\partial x^2} + \frac{\partial^2 h}{\partial y^2}) \end{equation} \]

The form of this equation simply adds the second partial derivative of height with respect to y to the right hand side of the 1D wave equation. The resulting discretization is simple and looks very familiar, we follow the same series of steps as we did with the 1D equation to arrive at:

\[ \begin{equation} \begin{split} a_{{i,j}_{t+\Delta t}} = &c^2 \frac{h^\star_{{i+1,j}_{t+\Delta t}} + \ h^\star_{{i-1,j}_{t+\Delta t}} + \ h^\star_{{i,j+1}_{t+\Delta t}} + \ h^\star_{{i,j-1}_{t+\Delta t}} - \ 4 h^\star_{{i,j}_{t+\Delta t}} } {\Delta x^2} + \\ &c^2 \Delta t^2 \frac{a_{{i+1,j}_{t+\Delta t}} + \ a_{{i-1,j}_{t+\Delta t}} + \ a_{{i,j+1}_{t+\Delta t}} + \ a_{{i,j-1}_{t+\Delta t}} - \ 4 a_{{i,j}_{t+\Delta t}} } {\Delta x^2} \label{eq:AccelDisc1Subscripted} \end{split} \end{equation} \]

The above equation assumes that \(\Delta x = \Delta y\), as does the derivation below. If working with non-square cells, the discretization gets a little longer, left as an exercise for the reader! (It rarely occurs in practice). Following the same process as in the 1D case, we rearrange our terms, substituting intermediate constants \(\kappa\) and \(\gamma\):

\[ \begin{equation} \kappa = \frac{c^2 \Delta t^2}{\Delta x^2}, \ \gamma = \frac{c^2}{\Delta x^2} \label{eq:Constants} \end{equation} \]

Giving the following equation which expresses the relationship between the acceleration at the index \((i,j)\) and its spatially adjacent neighbors, dropping the time subscript because it is unnecessary:

\[ \begin{equation} \begin{split} (1 + 4 \kappa) a_{i,j} + &( -\kappa ) ( a_{i-1,j} + a_{i+1,j} + a_{i,j-1} + a_{i,j+1} ) = \\ &\gamma ( h^\star_{i-1,j} + \ h^\star_{i+1,j} + \ h^\star_{i,j-1} + \ h^\star_{i,j+1} - \ 4 h^\star_{i,j} ) \\ \label{eq:AccelOneMatrixRow} \end{split} \end{equation} \]

The above equation \(\eqref{eq:AccelOneMatrixRow}\) represents a single row in the linear matrix we are solving. Once again, we'll use a jacobi solver, computing a new guess for acceleration \(a^{\star k+1}\) in each iteration \(k+1\) from the acceleration \(a^{\star k}\) of the previous iteration, which looks like this:

\[ \begin{eqnarray} b_{i,j} &=& \gamma ( h^\star_{i-1,j} + \ h^\star_{i+1,j} + \ h^\star_{i,j-1} + \ h^\star_{i,j+1} - \ 4 h^\star_{i,j} ) \\ c^{\star k}_{i,j} &=& \kappa ( a^{\star k}_{i-1,j} + \ a^{\star k}_{i+1,j} + \ a^{\star k}_{i,j-1} + \ a^{\star k}_{i,j+1} ) \\ a^{\star k+1}_{i,j} &=& \frac{b_{i,j} + c^{\star k}_{i,j}}{1 + 4 \kappa}\\ \label{eq:JacobiIteration} \end{eqnarray} \]

This translates very directly into code:

// Jacobi iteration to get temp acceleration
void JacobiIterationAccel( int i_aOld, int o_aNew, int i_hStar, float i_dt ) {
    float kappa = sq( WaveSpeed ) * sq( i_dt ) / sq( DXY );
    float gamma = sq( WaveSpeed ) / sq( DXY );

    for (int j = 1; j < NY-1; ++j) {
        for (int i = 1; i < NX-1; ++i) {
            float a_left = State[i_aOld][IX(i-1,j)];
            float a_right = State[i_aOld][IX(i+1,j)];
            float a_down = State[i_aOld][IX(i,j-1)];
            float a_up = State[i_aOld][IX(i,j+1)];

            float h_star_left = State[i_hStar][IX(i-1,j)];
            float h_star_right = State[i_hStar][IX(i+1,j)];
            float h_star_down = State[i_hStar][IX(i,j-1)];
            float h_star_up = State[i_hStar][IX(i,j+1)];
            float h_star_cen = State[i_hStar][IX(i,j)];

            float b = gamma *
                (h_star_left + h_star_right + h_star_down + h_star_up -
                 (4.0 * h_star_cen));

            float c = kappa * (a_left + a_right + a_down + a_up);

            State[o_aNew][IX(i,j)] = (b + c) / (1.0 + kappa);
        }
    }

    EnforceDirichletBoundaryConditions( o_aNew );
}

Arduously Refactoring What's Left

We've made these changes so far in our port from 1D to 2D:

• Updated the initial conditions to work in 2D
• Updated boundary condition enforcement to work in 2D
• Added 2D drawing code
• Updated the jacobi iteration to work in 2D

Now we need to undertake the work of porting the remainder of the solver, which is still the bulk of the code, by length. Maybe go get a coffee and have a stretch before this undertaking. Ready?

And... we're done. (We didn't have to do anything at all)

By organizing our state vector into generic arrays, and sticking with an underlying one-dimensional layout of the arrays themselves, most of the actual solver machinery - the estimation of new fields, the time integration - works without modification. Additionally, it's modular. We can swap out RK4 for RK2, or try a first order time step to see its effect on stability. This modularity and flexibility is the payoff for the good design practices we established at the outset, and it translates into the design and implementation of bigger systems, even (especially) production-quality solvers. It also tends to abstract out parallelism very effectively.

I've added a little extra feature to this final solver - hit the 't' key to cycle between first-order, RK2, and RK4 time steps. It's interesting to notice the increase in stability with each one, though the gap between first-order and RK2 is much bigger than the gap between RK2 and RK4.

There's a bit of aliasing around the mouse input point, and the waves created have a small amount of aliasing to them as well. As an exercise for the reader, consider these two approaches towards fixing these problems. Firstly, instead of just drawing a single one-pixel dot as the mouse input boundary condition, draw a smooth shape instead. This is challenging because the application of boundary conditions needs to be idempotent - meaning once they're applied, applying them a second time should have no effect. Secondly, we're using a first-order discretization of space in the discretization of our laplacian. If you used a higher-order spatial discretization, what would that look like? How would it change the Jacobi iteration?

Download

All of the files associated with this class are available via a public github repository, found here: https://github.com/blackencino/SimClass