Example: Using the Standard ESN#

This tutorial shows how to use the ESN class, covering:

building and training the network
making sample predictions and comparing the predictions to test data
storing the network weights and re-reading the network to an ESN

[1]:

import numpy as np
import matplotlib.pyplot as plt

[2]:

plt.style.use("./xesn.mplstyle")

Create training and testing data#

Here we generate a time series using the 6 Dimensional version of the Lorenz 96 model [Lorenz, 1996]. First we spinup the model dynamics before making the training data, then discard a brief transient period, and finally create the test dataset.

This uses a local module here in the documentation directory.

[3]:

from lorenz import Lorenz96

[4]:

model = Lorenz96(N=6)

[5]:

n_spinup = 2_000
n_train = 40_000
n_transient = 1_000
n_test = 1_000

n_total = n_spinup+n_train+n_transient+n_test

rs = np.random.RandomState(0)
trajectory = model.generate(n_steps=n_total, x0=rs.normal(size=(model.N,)))

[6]:

trainer = trajectory.isel(time=slice(n_spinup,n_spinup+n_train+1))
tester = trajectory.isel(time=slice(-n_test-1, None))

Normalization#

Here we normalize the data very simply, with the training mean and standard deviation. Note that we do not normalize each spatial dimension separately, following [Platt et al., 2022] and [Smith et al., 2023].

[7]:

bias = trainer.mean()
scale = trainer.std()

[8]:

trainer = (trainer - bias)/scale
tester = (tester - bias)/scale

Create the ESN#

Here we create a network that fully emulates the 6D Lorenz96 system, so it has input and output dimension equal to 6. The other parameters have been obtained from previous experiments, and are used just as an example.

[9]:

from xesn import ESN

[10]:

esn = ESN(
    n_input=model.N,
    n_output=model.N,
    n_reservoir=500,
    leak_rate=0.87,
    tikhonov_parameter=6.9e-7,
    input_kwargs={
        "factor": 0.86,
        "normalization": "svd",
        "distribution": "uniform",
        "random_seed": 0,
    },
    adjacency_kwargs={
        "factor":0.71,
        "normalization": "eig",
        "distribution": "uniform",
        "is_sparse": True,
        "connectedness": 5,
        "random_seed": 1,
    },
    bias_kwargs={
        "factor": 1.8,
        "distribution": "uniform",
        "random_seed": 2,
    },
)

And note that we can always get a printed view of the values set in the ESN

[11]:

esn

[11]:

ESN
    n_input:                6
    n_output:               6
    n_reservoir:            500
---
    leak_rate:              0.87
    tikhonov_parameter:     6.9e-07
---
    Input Matrix:
        factor              0.86
        normalization       svd
        distribution        uniform
        random_seed         0
---
    Adjacency Matrix:
        factor              0.71
        normalization       eig
        distribution        uniform
        is_sparse           True
        connectedness       5
        random_seed         1
---
    Bias Vector:
        factor              1.8
        distribution        uniform
        random_seed         2

Note: It is usually a good idea to use the random_seed to control the randomly generated matrices, so that results are consistent. Once reasonable performance is obtained, the seed can be varied to determine how dependent the results are on the exact matrices and vector chosen.

Build the network and train the readout matrix#

The .build() method creates the two random matrices and the random bias vector. Once these are generated, we can run .train() to learn the readout matrix weights.

[12]:

esn.build()

[13]:

%%time
esn.train(trainer, batch_size=5_000)

CPU times: user 15.8 s, sys: 862 ms, total: 16.7 s
Wall time: 1.8 s

Note: the batch_size parameter should be chosen so that batch_size x n_reservoir can fit comfortably in memory. Here we choose a somewhat conservative value.

Cleanup#

[23]:

from shutil import rmtree

[24]:

rmtree("esn-weights.zarr")

[ ]:

Example: Using the Standard ESN

Contents

Example: Using the Standard ESN#

Create training and testing data#

Normalization#

Create the ESN#

Build the network and train the readout matrix#

Test the network#

Store the Trained Readout Weights#

Read in the Network and Show Reproducibility#

Cleanup#