1D linear shallow water equations in JAX

Authors

Joe Wallwork

Sam Avis

Adeleke Bankole 

%matplotlib inline
from IPython.display import HTML
import matplotlib.animation
import matplotlib.pyplot as plt
import numpy as np
import jax.numpy as jnp

The shallow water equations are a depth-averaged approximation to the Navier-Stokes equations. The standard 2D formulation makes use of a horizontal velocity \(\mathbf{u}:\Omega\times(0,T)\rightarrow\mathbb{R}^2\) and a free surface elevation field \(\eta:\Omega\times(0,T)\rightarrow\mathbb{R}\), where we have a spatial domain \(\Omega\subset\mathbb{R}^2\) and time period \((0,T)\) for some \(T>0\). The free surface elevation \(\eta\) can be interpreted as the perturbation of the water surface from rest. The water depth at rest is referred to as the bathymetry, \(b\), and the total water depth is \(h=b+\eta\). The bathymetry is fixed in time, whereas the total water depth is not.

In this mini-project, we apply several simplifications:

  1. Consider a 1D spatial domain rather than a 2D spatial domain.
  2. Do not consider Coriolis forces due to planetary rotation.
  3. Do not consider nonlinear terms.

The non-rotational shallow water equations are given by \[ \begin{align} \frac{\partial u}{\partial t}+\nabla\cdot(\nabla\mathbf{u})+g\nabla\eta&=\nabla\cdot(\nu(\nabla\mathbf{u}+(\nabla\mathbf{u})^T)),\\ \frac{\partial\eta}{\partial t}+\nabla\cdot(h\mathbf{u})&=0, \end{align} \] where \(\nu\) is a viscosity coefficient, and \(g\) is the gravitational acceleration constant.

Restricting to the 1D case and ‘linearising’, the advection and viscosity terms are dropped and we are left with \[ \begin{align} \frac{\partial u}{\partial t}+g\frac{\partial\eta}{\partial x}&=0,\\ \frac{\partial\eta}{\partial t}+b\frac{\partial u}{\partial x}+u\frac{\partial b}{\partial x}&=0, \end{align} \] where now \(u:\Omega\times(0,T)\rightarrow\mathbb{R}\). Note that \(\frac{\partial b}{\partial x}\) is constant in time.

Consider the 1D spatial domain \(\Omega=(0,X)\) with \(X>0\).

Combining the state as a vector \(\mathbf{w}=(u,\eta)\) we have \[ \frac{\partial\mathbf{w}}{\partial t} +\begin{bmatrix}0&g\\b&0\end{bmatrix}\frac{\partial\mathbf{w}}{\partial x} +\begin{bmatrix}0&0\\\frac{\partial b}{\partial x}&0\end{bmatrix}\mathbf{w} =\boldsymbol0, \] which can be shown to be equivalent to a wave equation with wavespeed \(\sqrt{gb}\). We introduce the notation \[ \underline{\mathbf{A}}=\begin{bmatrix}0&g\\b&0\end{bmatrix}, \quad\underline{\mathbf{B}}=\begin{bmatrix}0&0\\b&0\end{bmatrix} \] for conciseness so that \[ \frac{\partial\mathbf{w}}{\partial t} +\underline{\mathbf{A}}\frac{\partial\mathbf{w}}{\partial x} +\frac{\partial\underline{\mathbf{B}}}{\partial x}\mathbf{w} =\boldsymbol0, \]

For initial conditions, assume zero velocity \(u(x,0)=0,\:\forall x\in(0,X)\) and spatially varying free surface \(\eta(x,0)=\eta_0(x)\).

Assume periodic boundary conditions for simplicity: \(\mathbf{w}(0,t)=\mathbf{w}(X,t),\:\forall t\in(0,T)\).

Set the spatial and temporal extents.

X = 400e3
T = 4200

Define the number of points for discretising both space and time.

nx = 400
nt = 4200

Determine the grid spacing and timestep.

assert nx > 1
assert nt > 1
x = np.linspace(0, X, nx)
t = np.linspace(0, T, nt)
dx = x[1] - x[0]
dt = t[1] - t[0]
print(f"dx = {dx:.4f}")
print(f"dt = {dt:.4f}")

Define the initial conditions as follows to correspond to the tsunami modelling problem in [1]. We assume zero horizontal velocity and apply a small perturbation to the free surface elevation to see how it propagates across the domain. Note that we only allow the elevation to be non-zero within \(25\,\mathrm{km}\) of \(x=125\,\mathrm{km}\). We will return to this later.

u0 = jnp.zeros_like(x)
source_indices = jnp.where(((x - 125e3) / 25e3) **2 < 1.0)[0]
sr = (source_indices[0], source_indices[-1] + 1)
source = jnp.maximum(0.4 - ((x[sr[0]:sr[1]] - 125e3) / 25e3) ** 2, 0.0)  # Initial guess
eta0 = jnp.concatenate((jnp.zeros_like(x[:sr[0]]), source, jnp.zeros_like(x[sr[1]:])))

Consider a constant bathymetry (water depth at rest) of \(4\,\mathrm{km}\) and the standard gravitational acceleration constant. Note that our implementation supports a spatially varying bathymetry. We also define g using JAX syntax so that we can differentiate with respect to it.

b = jnp.ones(nx) * 4000.0
g = jnp.array([9.81])
def plot_solution(sol, axes=None, both=False):
    """Plot the solution field.

    :arg sol: the solution array corresponding to both velocity and elevation at a time level
    :kwarg axes: optional matplotlib axes object to plot on
    :kwarg both: logical flag for plotting both solution fields when True or just elevation when False
    """
    if axes is None:
        fig, axes = plt.subplots(figsize=(6, 2))
    if both:
        axes.plot(x / 10e3, sol[0::2], label=r"Velocity, $u$")
    axes.plot(x / 10e3, sol[1::2], label=r"Elevation, $\eta$")
    axes.set_xlabel(r"$x$ [km]")
    if both:
        axes.legend()
        axes.set_ylabel("Magnitude")
    else:
        axes.set_ylabel(r"Elevation, $\eta$ [m]")
        axes.set_ylim([-0.4, 0.4])
    axes.set_xlim([0, 40])
    axes.grid()
fig, axes = plt.subplots(ncols=2, figsize=(12, 2))
w0 = jnp.vstack((u0, eta0)).transpose().flatten()
plot_solution(w0, axes=axes[0], both=True)
axes[0].set_title("Initial conditions")
axes[1].set_title("Bathymetry")
axes[1].invert_yaxis()  # Flip the y-axis for the bathymetry plot
axes[1].plot(x / 10e3, b)
axes[1].set_xlabel(r"$x$ [m]")
axes[1].set_ylabel(r"Bathymetry, $b$ [km]")
axes[1].set_xlim([0, 40])
axes[1].grid()

Applying implicit Euler for timestepping gives the approximation \[ \frac{\mathbf{w}^{k+1}-\mathbf{w}^k}{\Delta t} +\underline{\mathbf{A}}\frac{\partial\mathbf{w}^{k+1}}{\partial x} +\frac{\partial\underline{\mathbf{B}}}{\partial x}\mathbf{w}^{k+1} \approx\boldsymbol0. \]

Applying a central difference for the spatial derivative gives \[ \frac{\mathbf{w}^{k+1}_i-\mathbf{w}^k_i}{\Delta t} +\underline{\mathbf{A}}\frac{\mathbf{w}^{k+1}_{i+1}-\mathbf{w}^{k+1}_{i-1}}{2\Delta x} +\frac{\underline{\mathbf{B}}_{i+1}-\underline{\mathbf{B}}_{i-1}}{2\Delta x}\mathbf{w}^{k+1}_i \approx\boldsymbol0. \]

Note We choose implicit Euler for the temporal discretisation and central differences for the spatial discretisation because this pairing is known to be stable.

Setting \(c=\frac{\Delta t}{2\Delta x}\) for conciseness and rearranging, we arrive at: \[ \mathbf{w}^{k+1}_i+c\left(\underline{\mathbf{A}}(\mathbf{w}^{k+1}_{i+1}-\mathbf{w}^{k+1}_{i-1})+(\underline{\mathbf{B}}_{i+1}-\underline{\mathbf{B}}_{i-1})\mathbf{w}^{k+1}_i\right) \approx\mathbf{w}^k_i \]

c = dt / (2 * dx)
print(f"c = {c:.4f}")

This gives rise to the block matrix system \[ \begin{bmatrix} \underline{\mathbf{I}}+c(\underline{\mathbf{B}}_2-\underline{\mathbf{B}}_{n_x}) & c\underline{\mathbf{A}} & & & -c\underline{\mathbf{A}}\\ -c\underline{\mathbf{A}} & \underline{\mathbf{I}}+c(\underline{\mathbf{B}}_3-\underline{\mathbf{B}}_1) & c\underline{\mathbf{A}}\\ & -c\underline{\mathbf{A}} & \underline{\mathbf{I}}+c(\underline{\mathbf{B}}_4-\underline{\mathbf{B}}_2) & c\underline{\mathbf{A}}\\ & & \ddots & \ddots & \ddots\\ c\underline{\mathbf{A}} & & & -c\underline{\mathbf{A}} & \underline{\mathbf{I}}+c(\underline{\mathbf{B}}_1-\underline{\mathbf{B}}_{n_x-1})\\ \end{bmatrix}\mathbf{w}^{k+1} \approx\mathbf{w}^k, \] where \(\underline{\mathbf{I}}\) is the \(2\times2\) identity matrix. Both block matrices are tridiagonal except that they have additional nonzero entries in the top-right and bottom-left entries due to the periodic boundary conditions.

The matrices are fixed in time so can be precomputed ahead-of-time. Given the state at timestep \(k\), this is a linear system we can solve to approximate the state at timestep \(k+1\).

For further convenience, define \[ \underline{\mathbf{U}}=c\begin{bmatrix} & \underline{\mathbf{A}} & & &\\ & & \underline{\mathbf{A}}\\ & & & \underline{\mathbf{A}}\\ & & & & \ddots\\ \underline{\mathbf{A}} \end{bmatrix}, \] \[ \underline{\mathbf{L}}=-c\begin{bmatrix} & & & & \underline{\mathbf{A}}\\ \underline{\mathbf{A}}\\ & \underline{\mathbf{A}}\\ & & \ddots\\ & & & \underline{\mathbf{A}} \end{bmatrix}, \] and \[ \underline{\mathbf{D}}=\begin{bmatrix} \underline{\mathbf{I}}+c(\underline{\mathbf{B}}_2-\underline{\mathbf{B}}_{n_x-1})\\ & \underline{\mathbf{I}}+c(\underline{\mathbf{B}}_3-\underline{\mathbf{B}}_1)\\ & & \ddots\\ & & & \underline{\mathbf{I}}+c(\underline{\mathbf{B}}_1-\underline{\mathbf{B}}_{n_x-1}) \end{bmatrix}, \] so that we have \[ (\underline{\mathbf{L}}+\underline{\mathbf{D}}+\underline{\mathbf{U}})\mathbf{w}^{k+1} \approx\mathbf{w}^k. \] We can define these matrices as follows:

B_diff = jnp.roll(b, -1) - jnp.roll(b, 1)
diagonal = jnp.eye(2 * nx) \
        + c * jnp.diag(jnp.concatenate((B_diff[0:1], jnp.vstack((jnp.zeros(nx-1), B_diff[1:])).transpose().flatten())), k=-1)
upper = c * (jnp.diag(jnp.concatenate((jnp.vstack((jnp.zeros(nx-1), b[:-1])).transpose().flatten(), jnp.zeros(1))), k=1) \
        + jnp.diag(jnp.concatenate((jnp.vstack((g * jnp.ones(nx-2), jnp.zeros(nx-2))).transpose().flatten(), g * jnp.ones(1))), k=3) \
        + jnp.diag(jnp.concatenate((jnp.zeros(1), g * jnp.ones(1), jnp.zeros(1))), k=-(2*nx-3)) \
        + jnp.diag(b[-1:], k=-(2*nx-1)))
lower = -c * (jnp.diag(jnp.concatenate((jnp.vstack((b[2:], jnp.zeros(nx-2))).transpose().flatten(), b[1:2])), k=-3) \
        + jnp.diag(jnp.concatenate((jnp.zeros(1), jnp.vstack((g * jnp.ones(nx-1), jnp.zeros(nx-1))).transpose().flatten())), k=-1) \
        + jnp.diag(jnp.concatenate((jnp.zeros(1), b[0:1], jnp.zeros(1))), k=2*nx-3) \
        + jnp.diag(g * jnp.ones(1), k=2*nx-1))
lhs_matrix = lower + diagonal + upper

We can plot the sparsity pattern of the matrix as follows:

plt.spy(lhs_matrix);

Solve the PDE by integrating in time and solving the linear system at each timestep. We import the tqdm utility to display a progress bar.

from tqdm import tqdm
def pde_solve(source):
    """Solve the PDE.

    :arg source: the initial elevation in the source region
    :return: solution trajectory
    """
    # Set initial condition
    eta0 = jnp.concatenate((jnp.zeros_like(x[:sr[0]]), source, jnp.zeros_like(x[sr[1]:])))
    w = jnp.vstack((u0, eta0)).transpose().flatten()
    trajectory = [w]
    
    # Do the time integration
    for k, time in enumerate(tqdm(t)):
        w = jnp.linalg.solve(lhs_matrix, w)
        trajectory.append(w)
    return jnp.array(trajectory)
trajectory = pde_solve(source)

Plot the solution trajectory at a few key times.

snapshot_times = [525, 1365, 2772, 3255, 4200]
idx = 0
fig, axes = plt.subplots(nrows=5, figsize=(6, 12))
for time, sol in zip(t, trajectory):
    if time >= snapshot_times[idx]:
        plt.text(0.98, 0.95, f"t={time:.0f}", transform=axes[idx].transAxes, fontsize=12, ha="right", va="top")
        plot_solution(sol, axes=axes[idx])
        idx += 1

Produce an animation, too.

fig, axes = plt.subplots(figsize=(6, 2))
axes.axis([0, 40e3, -0.1, eta0.max()])
l, = axes.plot([],[])

frame_rate = 25

def animate(i):
    axes.clear()
    plot_solution(trajectory[frame_rate * i], axes=axes)

ani = matplotlib.animation.FuncAnimation(fig, animate, frames=len(trajectory[::frame_rate]))
plt.close()
HTML(ani.to_jshtml())

So we have a way to do a full simulation and save the entire trajectory.

When it comes to real world data, we only usually have timeseries data at a sequence of points. The devices used to measure such timeseries data in ocean modelling are often referred to as gauges. Given indices representing the gauge locations, we can create a new function that produces timeseries data at those locations.

def produce_timeseries(source, gauge_indices):
    """Solve the PDE for a given initial elevation and produce timeseries at given indices.

    :arg source: the initial elevation in the source region
    :arg gauge_locations: list of gauge indices
    :return: list of timeseries at the gauges
    """
    eta_timeseries = pde_solve(source)[:, 1::2]
    return eta_timeseries[:, gauge_indices]
gauge_indices = [0, 100, 200, 300]
timeseries = produce_timeseries(source, gauge_indices)

Let’s plot these against data from a different run with an unseen source.

gauge_data = np.load("gauge_data.npy")
fig, axes = plt.subplots(nrows=4, figsize=(6, 8))
for i in range(4):
    gauge_index = gauge_indices[i]
    plt.text(0.5, 0.95, f"Gauge at x={x[gauge_index]/1e3:.0f}km", transform=axes[i].transAxes, fontsize=12, ha="center", va="top")
    axes[i].plot(timeseries[:, gauge_index], label="Simulated")
    axes[i].plot(gauge_data[:, i], label="Data")
    axes[i].set_xlabel(r"Time, $t$ [s]")
    axes[i].set_ylabel(r"Elevation, $\eta$ [m]")
    axes[i].legend(loc="upper right")
    axes[i].grid()
def objective_function(source):
    timeseries = produce_timeseries(source, gauge_indices)
    return jnp.dot(timeseries - gauge_data, timeseries - gauge_data) ** 2

Mini-project ideas

  1. Conduct a source inversion experiment where you seek to recover the source condition source that gave rise to the data saved to file. You will do this by optimising the fit of the free surface elevation solution at a set of points against timeseries. The fit is encoded in objective_function so you should be able to invert this function for source using the optimisers covered in the course.

  2. Use differentiable programming to compute sensitivities of the model with respect to the source condition source. Consider also sensitivities of the model with respect to the gravitational constant g and with respect to the bathymetry field b. Try out spatially varying bathymetry fields.

Note on idea 1: This is an ill-posed problem! We have added some regularisation to the problem by restricting source (and hence eta0) to only be non-zero in a specific region but you may need to include a regularisation term in your objective function.

Note on idea 2: The numerical scheme is unlikely to be able to cope with shocks due to sharp gradients in bathymetry so you should avoid ‘shelf break’ step functions. You should also ensure the bathymetry field is periodic.

References

[1] Davis, B. N., & LeVeque, R. J. (2016). Adjoint methods for guiding adaptive mesh refinement in tsunami modeling. In Global tsunami science: Past and future, volume I (pp. 4055-4074). Birkhäuser, Cham.