Applying SINDy equation identification to L96

Applying SINDy equation identification to L96#

Outline#

This notebook provides an example of SINDy equation identification in application to the two-scale L96 model with the use of pysindy library. The objective is to try to identify the governing ODEs for the large-scale variable ($X_i$) if we know only their time measurements. In other words, we want to see how well the SINDy model can capture the form of the governing equations for large-scale variables in the presence of the subgrid forcing by small-scale variables ($Y$) when there is no measurements of the small scales.

Import#

import numpy as np
import pysindy as ps
import matplotlib.pyplot as plt
import seaborn as sns
from L96_model import L96, L96_eq1_xdot, integrate_L96_2t, integrate_L96_1t

Routines for plotting#

colors = [
    "#1f77b4",
    "#ff7f0e",
    "#2ca02c",
    "#d62728",
    "#9467bd",
    "#8c564b",
    "#e377c2",
    "#7f7f7f",
    "#bcbd22",
    "#17becf",
]

def plot_coefficients(
    coefficients, input_features=None, feature_names=None, ax=None, **heatmap_kws
):
    if input_features is None:
        input_features = [f"$\dot x_{k}$" for k in range(coefficients.shape[0])]
    else:
        input_features = [f"$\dot {fi}$" for fi in input_features]

    if feature_names is None:
        feature_names = [f"f{k}" for k in range(coefficients.shape[1])]

    with sns.axes_style(style="white", rc={"axes.facecolor": (0, 0, 0, 0)}):
        if ax is None:
            fig, ax = plt.subplots(1, 1)

        max_mag = np.max(np.abs(coefficients))
        heatmap_args = {
            "xticklabels": input_features,
            "yticklabels": feature_names,
            "center": 0.0,
            "cmap": sns.color_palette("vlag", n_colors=20, as_cmap=True),
            "ax": ax,
            "linewidths": 0.1,
            "linecolor": "whitesmoke",
        }
        heatmap_args.update(**heatmap_kws)

        sns.heatmap(coefficients.T, **heatmap_args)

        ax.tick_params(axis="y", rotation=0)

    return ax

def signed_sqrt(x):
    return np.sign(x) * np.sqrt(np.abs(x))

def compare_coefficient_plots(coefficients, input_features=None, feature_names=None):
    n_cols = len(coefficients)
    with sns.axes_style(style="white", rc={"axes.facecolor": (0, 0, 0, 0)}):
        max_mag = np.max(np.abs(coefficients))
        fig, axs = plt.subplots(1, 1, figsize=(4, 4), sharey=True, sharex=True)
        plot_coefficients(
            signed_sqrt(coefficients),
            input_features=input_features,
            feature_names=feature_names,
            ax=axs,
            cbar=False,
            vmax=max_mag,
            vmin=-max_mag,
        )
        fig.tight_layout()

L96 two-scale model#

https://www.researchgate.net/publication/319201436/figure/fig1/AS:869115023589376@1584224577926/Visualisation-of-a-two-scale-Lorenz-96-system-with-J-8-and-K-6-Global-scale-values.png

Fig. 7 Visualisation of a two-scale Lorenz ‘96 system with J = 8 and K = 6. Global-scale variables ($X_k$) are updated based on neighbouring variables and on the local-scale variables ($Y_{j,k}$) associated with the corresponding global-scale variable. Local-scale variabless are updated based on neighbouring variables and the associated global-scale variable. The neighbourhood topology of both local and global-scale variables is circular. Image from Exploiting the chaotic behaviour of atmospheric models with reconfigurable architectures - Scientific Figure on ResearchGate..#

Two-time-scale Lorenz 96 model: $$ \frac{d}{dt} X_k = - X_{k-2} X_{k-1} + X_{k-1} X_{k+1} - X_k + F - \left( \frac{hc}{b} \right) \sum_{j=0}^{J-1} Y_{j,k} $$ \frac{d}{dt} Y_{j,k} = - c b Y_{j+1,k} (Y_{j+2,k} - Y_{j-1,k}) - c Y_{j,k} + \left( \frac{hc}{b} \right) \sum_{j=0}^{J-1} X_{k} $$

We will put K=8 large-scale variables ($X$), for each of them there is J=32 small-scale variables ($Y$). For the first example, we will consider the case of weak interaction between the scales, which is defined by the parameter $h$ in the L96 class.

K = 8  # Number of degrees of freedom for X
J = 32  # Number of degrees of freedom for Y per X
Forcing = 18.0  # Model parameter
h = 0.1  # WEAK coupling

def s(k, K):
    """A non-dimension coordinate from -1..+1 corresponding to k=0..K"""
    return 2 * (0.5 + k) / K - 1

# Initial conditions
k = np.arange(K)  # For coordinate in plots
j = np.arange(J * K)  # For coordinate in plots
# Initial conditions
X_init = s(k, K) * (s(k, K) - 1) * (s(k, K) + 1)
Y_init = 0 * s(j, J * K) * (s(j, J * K) - 1) * (s(j, J * K) + 1)

steps = 1000
dt = 0.005

# Create a "real world" with K=8 and J=32
W = L96(K, J, F=Forcing, h=h)
W.set_state(X_init, Y_init)
# Run the true state
X, Y, t = W.run(dt, dt * steps, store=True)

fig, ax = plt.subplots(1, 1, figsize=(6, 4), dpi=150)
contours = ax.contourf(k, t, X, levels=20)
fig.colorbar(contours)
ax.set(xlabel="k", ylabel="t")
plt.show();

../_images/e37cf216aaaec972b8c6f59e547a4da3866f28e031fa801e4260d16e1b7dcc40.png

The data array $X$ will be used for the system identification by SINDy method. We suppose that we do not dispose measurements for the small-scale variables $Y$, so their contribution to the large-scale dymanics will basically treated as noise.

Define SINDy model#

The goal is to identify the governing equations for the large-scale variables $X$, having the measurement for them.

differentiation_method = ps.SINDyDerivative(kind="finite_difference", k=1)

feature_library = ps.PolynomialLibrary(degree=2)

# sequential thresholded least squares optimizer
optimizer = ps.STLSQ(threshold=0.3)

model = ps.SINDy(
    differentiation_method=differentiation_method,
    feature_library=feature_library,
    optimizer=optimizer,
    feature_names=[("x" + str(i)) for i in k],
)