esn_ctf

This directory contains implementations of echo state networks [1,2] (ESNs) for forecasting dynamical systems, as well as supporting evaluation in the Common Task Framework (CTF) for Science. As implemented [3, 4], an ESN is a form of recurrent neural that evolves according to $$h(t+1) = \alpha \tanh (W_{r}h(t) + W_{in} u(t) + \sigma_b \mathbf 1) + (1-\alpha)h(t),$$ where $h$ is the latent state, $\alpha$ is a leak rate, $\sigma_b$ is a scalar bias term, and $W_r$ and $W_{in}$ are network parameters. In contrast to other forms of RNN, $W_r$ and $W_{in}$ are randomly initialized and left fixed. A linear map, $W_{out}$, is then learned via ridge regression such that $W_{out} g(h(t)) \approx u(t)$ where $g: \mathbb R^h \to \mathbb R^h$. Often, $g$ is taken to be the identity mapping, but other choices are sometimes necessary [4]. The cost of this simple training procedure is that $h$ must be taken to be high-dimensional.

Within the CTF, we evaluate the performance of ESNs on low-dimensional systems (i.e. ODE systems) as presented in [3] and on high-dimensional systems as presented in [5].

The ESN itself comes from OpenReservoirComputing (orc) [6]; this repository supplies only the CTF glue: data loading, the dimensionality reduction described below, the per-pair prediction rules, and the evaluation and tuning entry points.

File	Role
run.py	Train on a pair's training data and evaluate against the CTF test set.
run_opt.py	Train on a split of the training data and score against the held-out part; called once per trial by the tuner.
pipeline.py	Model construction, training, and the per-pair_id prediction rules shared by both runners.
reduction.py	Whitened POD projection onto the leading singular modes (see below).
optimize_parameters.py	Ray Tune entry point, via ctf4science.tune_module.ModelTuner.

Usage

First, ensure that ctf4science is already installed. Then, from the root directory of esn_ctf,

pip install -r requirements.txt

will install the remaining dependencies.

Given a config file (examples in config directory), one can train and evaluate a model by running

python run.py <path-to-config>

Or, given a hyperparameter tuning config (examples in tuning_config directory), one can optimize a set of hyperparameters by running

python optimize_parameters.py --config-path <path-to-config>

The bash scripts driver_run.sh and test_opt_params.sh can provide a comprehensive test that your environment has been set up correctly by replicating the results fronm the CTF paper. It is strongly recommended to update the tuning config files to enforce a maximum number of trials before running test_opt_params.sh.

Dimensionality reduction

Some datasets are too high-dimensional to drive a reservoir directly: msfr has 19400 spatial degrees of freedom. For these, the reservoir is run on a reduced representation, selected by a rank key in the model section:

model:
  # ...
  rank: 32           # number of singular modes retained
  noise_level: 0.001

rank fits a whitened POD basis to the pair's training data, encoding as $\mathrm{diag}(1/\sigma) U^{\top} x$ and decoding as $U \mathrm{diag}(\sigma) z$. Dividing by the singular values puts every retained mode on a comparable scale, so a single input scaling is meaningful across modes. The reservoir's input/output dimension becomes rank, so chunks must divide rank.

noise_level adds Gaussian noise to the encoded training trajectories, scaled by each mode's own standard deviation over time and drawn independently per trajectory. It regularizes the readout by making the reservoir see states slightly off the training trajectory; set it to 0 (or omit it) to disable.

Configs without a rank key run the ESN on the full state, so the runners contain no dataset-specific branches. In that case data_dim is checked against the dataset's actual feature count.

Predictions are always decoded back to the full state space before evaluation, including during tuning, so the tuning objective matches the metric reported on the test set. The reduction is fitted on training data only; run_opt.py fits it on the training side of the split, never on the validation data.

Hyperparameter tuning

tuning_config/ holds one file per dataset and pair_id. In addition to the ESN parameters, msfr configs tune rank, the size of the reduced representation:

hyperparameters:
  rank:
    type: randint
    lower_bound: 8
    upper_bound: 64

n_trials is set in each msfr tuning config, since msfr trials are far more expensive than the ODE ones.

Citations

[1] Jaeger, "The echo state approach to analyzing and training recurrent neural networks," German National Research Center for Information Technology, Technical Report GMD 148, 2001.

[2] Maas and Markram, "On the computational power of circuits of spiking neurons," Journal of Computer and System Sciences, 2004.

[3] Platt et al., "A systematic exploration of reservoir computing for forecasting complex spatiotemporal dynamics," Neural Networks, 2022.

[4] Kidger and Garcia, "Equinox: neural networks in (JAX) via callable PyTrees and filtered transformations," Differentiable Programming workshop at Neural Information Processing Systems, 2021.

[5] Pathak et al., "Model-Free Prediction of Large Spatiotemporally Chaotic Systems from Data: A Reservoir Computing Approach," Physical Review Letters, 2018.

[6] Williams and Tretiak, "OpenReservoirComputing: GPU accelerated implementations of common RC architectures," https://github.com/Jan-Williams/OpenReservoirComputing.