Example: Lorenz System

import matplotlib.pyplot as plt
%matplotlib inline
import matplotlib_inline
matplotlib_inline.backend_inline.set_matplotlib_formats('svg')
import seaborn as sns

!pip install diffrax

Requirement already satisfied: diffrax in /Users/ibilion/.pyenv/versions/3.11.6/envs/me697/lib/python3.11/site-packages (0.6.2)

Requirement already satisfied: equinox>=0.11.10 in /Users/ibilion/.pyenv/versions/3.11.6/envs/me697/lib/python3.11/site-packages (from diffrax) (0.11.11)
Requirement already satisfied: jax>=0.4.38 in /Users/ibilion/.pyenv/versions/3.11.6/envs/me697/lib/python3.11/site-packages (from diffrax) (0.4.38)
Requirement already satisfied: jaxtyping>=0.2.24 in /Users/ibilion/.pyenv/versions/3.11.6/envs/me697/lib/python3.11/site-packages (from diffrax) (0.3.9)
Requirement already satisfied: lineax>=0.0.5 in /Users/ibilion/.pyenv/versions/3.11.6/envs/me697/lib/python3.11/site-packages (from diffrax) (0.0.7)
Requirement already satisfied: optimistix>=0.0.7 in /Users/ibilion/.pyenv/versions/3.11.6/envs/me697/lib/python3.11/site-packages (from diffrax) (0.0.10)
Requirement already satisfied: typeguard==2.13.3 in /Users/ibilion/.pyenv/versions/3.11.6/envs/me697/lib/python3.11/site-packages (from diffrax) (2.13.3)
Requirement already satisfied: typing-extensions>=4.5.0 in /Users/ibilion/.pyenv/versions/3.11.6/envs/me697/lib/python3.11/site-packages (from diffrax) (4.15.0)
Requirement already satisfied: jaxlib<=0.4.38,>=0.4.38 in /Users/ibilion/.pyenv/versions/3.11.6/envs/me697/lib/python3.11/site-packages (from jax>=0.4.38->diffrax) (0.4.38)
Requirement already satisfied: ml_dtypes>=0.4.0 in /Users/ibilion/.pyenv/versions/3.11.6/envs/me697/lib/python3.11/site-packages (from jax>=0.4.38->diffrax) (0.5.4)
Requirement already satisfied: numpy>=1.24 in /Users/ibilion/.pyenv/versions/3.11.6/envs/me697/lib/python3.11/site-packages (from jax>=0.4.38->diffrax) (1.26.4)
Requirement already satisfied: opt_einsum in /Users/ibilion/.pyenv/versions/3.11.6/envs/me697/lib/python3.11/site-packages (from jax>=0.4.38->diffrax) (3.4.0)
Requirement already satisfied: scipy>=1.10 in /Users/ibilion/.pyenv/versions/3.11.6/envs/me697/lib/python3.11/site-packages (from jax>=0.4.38->diffrax) (1.12.0)

Requirement already satisfied: wadler-lindig>=0.1.3 in /Users/ibilion/.pyenv/versions/3.11.6/envs/me697/lib/python3.11/site-packages (from jaxtyping>=0.2.24->diffrax) (0.1.7)

Example: Lorenz System#

Consider the Lorenz system of differential equations:

\begin{align*} \dot{x} &= \sigma(y-x),\ \dot{y} &= x(\rho-z)-y,\ \dot{z} &= xy-\beta z. \end{align*}

This system has a chaotic attractor for \(\sigma=10\), \(\beta=8/3\), and \(\rho=28\). We are going to study the sensitivity of the system to its initial conditions. Our goal is to demonstrate that local sensitivity analysis is not appropriate for this system.

Let’s code it up and see what it looks like.

import jax.numpy as jnp
from jax import vmap, jit
from functools import partial
from diffrax import diffeqsolve, Tsit5, ODETerm, SaveAt


def solve_lorenz(u0, theta):

    def vector_field(t, u, theta):
        x = u[0]
        y = u[1]
        z = u[2]
        sigma = theta[0]
        beta = theta[1]
        rho = theta[2]
        dx = sigma * (y - x)
        dy = x * (rho - z) - y
        dz = x * y - beta * z
        return jnp.array([dx, dy, dz])

    return diffeqsolve(
        ODETerm(vector_field),
        Tsit5(),
        t0=0.0,
        t1=100.0,
        dt0=0.1,
        y0=u0,
        args=theta,
        saveat=SaveAt(ts=jnp.linspace(0.0, 100.0, 10_000))
    ).ys

monte_carlo_lorenz = jit(vmap(solve_lorenz, in_axes=(0, None)))

sigma = 10.0
beta = 8.0 / 3.0
rho = 28.0
theta = jnp.array([sigma, beta, rho])
u0 = jnp.array([1.0, 1.0, 1.0]).reshape(1, 3)
ys = monte_carlo_lorenz(u0, theta)
ys.shape

(1, 10000, 3)

Plot the first time steps of the solution:

ts = jnp.linspace(0.0, 100.0, 10_000)
steps_to_plot = 1_000

fig, ax = plt.subplots(figsize=FIGURE_SIZES["half_standard"])
markevery = max(1, steps_to_plot // 12)
ax.plot(ts[:steps_to_plot], ys[0, :steps_to_plot, 0], color="0.10", linestyle="-", marker="o", markevery=markevery, markersize=3, label="x", lw=1.0)
ax.plot(ts[:steps_to_plot], ys[0, :steps_to_plot, 1], color="0.40", linestyle="--", marker="s", markevery=markevery, markersize=3, label="y", lw=1.0)
ax.plot(ts[:steps_to_plot], ys[0, :steps_to_plot, 2], color="0.65", linestyle="-.", marker="^", markevery=markevery, markersize=3, label="z", lw=1.0)
ax.legend(frameon=False)
ax.set(xlabel="t", ylabel="u(t)", title="Lorenz system")
finalize_axes(keep_box=False)

array([<Axes: title={'center': 'Lorenz system'}, xlabel='t', ylabel='u(t)'>],
      dtype=object)
The first portion of a Lorenz trajectory, with the x, y, and z state components plotted against time.

And here is the classic butterfly plot:

fig = plt.figure(figsize=FIGURE_SIZES["full_standard"])
ax = plt.axes(projection='3d')
ax.plot(ys[0, :, 0], ys[0, :, 1], ys[0, :, 2], lw=0.5, alpha=0.5)
ax.set(xlabel="x", ylabel="y", zlabel="z", title="Lorenz system")
finalize_axes(keep_box=False)

array([<Axes3D: title={'center': 'Lorenz system'}, xlabel='x', ylabel='y', zlabel='z'>],
      dtype=object)
Three-dimensional Lorenz trajectory forming the characteristic two-lobed butterfly attractor.

We now take a small cloud of initial conditions and examine how it evolves in time. We plot the points in dark gray so that the cloud remains visible in black-and-white print.

import numpy as np

mu = np.array([1.0, 1.0, 1.0])
sigma = np.array([0.001, 0.001, 0.001])
num_samples = 1_000
u0_samples = np.random.normal(mu, sigma, size=(num_samples, 3))
ys_samples = monte_carlo_lorenz(u0_samples, theta)
ys_samples.shape

(1000, 10000, 3)

Here we go, the blob at different times:

for i in [0, 100, 500, 750, 1_000, 1_500, 2_000, 5_000, 10_000]:
    fig = plt.figure(figsize=FIGURE_SIZES["full_standard"])
    ax = plt.axes(projection='3d')
    ax.plot(ys_samples[:, i, 0], ys_samples[:, i, 1], ys_samples[:, i, 2], '.', color='0.15', alpha=0.85, ms=1.0)
    ax.set(xlabel="x", ylabel="y", zlabel="z", title=f"Lorenz system at t={ts[i]:.2f}")
    finalize_axes(keep_box=False)
    ax.set_xlim(-20, 20)
    ax.set_ylim(-20, 20)
    ax.set_zlim(0, 50)
Snapshot of the Lorenz ensemble in three-dimensional phase space during its evolution toward the attractor. Snapshot of the Lorenz ensemble in three-dimensional phase space during its evolution toward the attractor. Snapshot of the Lorenz ensemble in three-dimensional phase space during its evolution toward the attractor. Snapshot of the Lorenz ensemble in three-dimensional phase space during its evolution toward the attractor. Snapshot of the Lorenz ensemble in three-dimensional phase space during its evolution toward the attractor. Snapshot of the Lorenz ensemble in three-dimensional phase space during its evolution toward the attractor. Snapshot of the Lorenz ensemble in three-dimensional phase space during its evolution toward the attractor. Snapshot of the Lorenz ensemble in three-dimensional phase space during its evolution toward the attractor. Snapshot of the Lorenz ensemble in three-dimensional phase space during its evolution toward the attractor.

You see that the tiny blob moved everywhere in the attractor.

What does this mean? We cannot predict the future of the system from the initial conditions even if we know the parameters perfectly.

We know that we will fail, but let’s check where local sensitivity analysis will get us. How far does it predict correctly?

We need to get the Jacobian with respect to the initial conditions.

from jax import jacobian

jit_solve_lorenz = jit(solve_lorenz)
jit_jac_solve_lorenz = jit(jacobian(solve_lorenz, argnums=0))

mu_lorenz = solve_lorenz(mu, theta)
jac_lorenz = jit_jac_solve_lorenz(mu, theta)

Here is the Jacobian.

jac_lorenz.shape

(10000, 3, 3)

Notice that towards the end it has quite a few NaNs:

jac_lorenz;

The NaNs are because of numerical errors. You may have to move to 64-bit floats to get rid of them. Anyway, local sensitivity analysis will break before that point. So let’s proceed. I will calculate only the pointwise variance of the time-dependent random state; the full time-time covariance matrix is too large.

var = np.einsum(
    "tij,j,tij->ti",
    jac_lorenz,
    sigma ** 2,
    jac_lorenz
)

Again, I really hope you appreciate the magic of the einsum function. Try to do the above calculation without it.

Let’s look at the mean and the variance at specific times:

mc_mean = ys_samples.mean(axis=0)

fig, ax = plt.subplots(figsize=FIGURE_SIZES["half_standard"])
ax.plot(ts[:steps_to_plot], mc_mean[:steps_to_plot, 0], label="x (Monte Carlo)", lw=1)
ax.plot(ts[:steps_to_plot], mu_lorenz[:steps_to_plot, 0], '--', label="x (Local sensitivity)", lw=1)
ax.set(xlabel="t", ylabel="u(t)", title="Lorenz system")
ax.legend(frameon=False)
finalize_axes(keep_box=False)

array([<Axes: title={'center': 'Lorenz system'}, xlabel='t', ylabel='u(t)'>],
      dtype=object)
Monte Carlo and local-sensitivity estimates of the Lorenz x-component mean versus time, agreeing initially and then diverging.

The variance breaks down even faster:

mc_var = ys_samples.var(axis=0)

steps_to_plot = 550
fig, ax = plt.subplots(figsize=FIGURE_SIZES["half_standard"])
ax.plot(ts[:steps_to_plot], mc_var[:steps_to_plot, 0], label="x (Monte Carlo)", lw=1)
ax.plot(ts[:steps_to_plot], var[:steps_to_plot, 0], '--', label="x (Local sensitivity)", lw=1)
ax.set(xlabel="t", ylabel="u(t)", title="Lorenz system")
ax.legend(frameon=False)
ax.set_xlim(0, 10)
finalize_axes(keep_box=False)

array([<Axes: title={'center': 'Lorenz system'}, xlabel='t', ylabel='u(t)'>],
      dtype=object)
Monte Carlo and local-sensitivity estimates of the Lorenz x-component variance versus time, with the local approximation breaking down early.