Lorenz ’96 model in JAX

Authors

Joe Wallwork

Sam Avis

Adeleke Bankole

The Lorenz ’96 model is an idealised dynamical system model formulated by Ed Lorenz in 1996 [1].

Let \(\mathbf{x}=(x_1,x_2,\dots,x_N)\) denote the prognostic variable for some \(N\geq4\). This vector variable describes the state of the system at a given time, which we can think of in a similar way as the collection of scalar variables describing atmospheric state (components of velocity, pressure, temperature, etc).

The formulation considered here is given by the ODE \[ \frac{\mathrm{d}x_i}{\mathrm{d}t}=(x_{i+1}-x_{i-2})\:x_{i-1}-x_i+F, \] where the indexing is periodic according to the following conventions:

\(x_{-1}=x_{N-1}\)
\(x_0=x_N\)
\(x_1=x_{N+1}\)

We use the NumPy implementation found on the Wikipedia page [2] but make use of jax.numpy rather than standard NumPy.

%matplotlib inline

import matplotlib.pyplot as plt
import numpy as np
import jax.numpy as jnp

Set model parameters

N = 5  # Number of variables
F = 8  # Forcing

Define the RHS function

def L96(x, t):
    """Evaluate the RHS of the Lorenz '96 model with constant forcing.

    :arg x: vector of state variables
    :arg t: current time
    :return: right-hand side of the Lorenz '96 model
    """
    return (jnp.roll(x, -1) - jnp.roll(x, 2)) * jnp.roll(x, 1) - x + F

Define a perturbation to apply to the initial conditions.

epsilon = np.zeros(N)
epsilon[0] = 0.1
epsilon = jnp.asarray(epsilon)

Define the timestep and discretise the time window we seek to integrate over.

dt = 0.01
end_time = 30.0
t = jnp.arange(0.0, end_time, dt)

Define the initial condition to be in equilibrium with constant forcing.

x0 = F * jnp.ones(N)

Define a function for time integration using explicit Euler.

def explicit_euler(epsilon):
    """Integrate over time using explicit Euler, given a perturbation of the initial condition.

    :arg epsilon: perturbation of the initial condition
    :return: solution trajectory over the time window of interest
    """
    trajectory = [x0 + epsilon]
    for ti in t:
        x_ = trajectory[-1]
        x = x_ + L96(x_, ti) * dt
        trajectory.append(x)
    return jnp.array(trajectory)

Plot the first three variables on 3D axes

trajectory = explicit_euler(epsilon)

fig = plt.figure()
ax = fig.add_subplot(projection="3d")
ax.plot(trajectory[:, 0], trajectory[:, 1], trajectory[:, 2])
ax.set_xlabel("$x_1$")
ax.set_ylabel("$x_2$")
ax.set_zlabel("$x_3$")
plt.show()

Mini-project ideas

Conduct a sensitivity analysis experiment to assess the sensitivity of the model to the perturbation epsilon. Currently, the perturbation only applies to the first component of the state variable. Consider perturbations of the other components, too.

References

[1] Lorenz, Edward (1996). “Predictability – A problem partly solved” (PDF). Seminar on Predictability, Vol. I, ECMWF. https://www.ecmwf.int/sites/default/files/elibrary/1995/10829-predictability-problem-partly-solved.pdf

[2] https://en.wikipedia.org/wiki/Lorenz_96_model