Implementation of a reservoir computer with functions to optimize the parameters and calculate the Lyapunov exponents
Please cite if using this code:
Platt, J. A., Penny, S. G., Smith, T. A., Chen, T.-C., & Abarbanel, H. D. I. (2022). A Systematic Exploration of Reservoir Computing for Forecasting Complex Spatiotemporal Dynamics. arXiv http://arxiv.org/abs/2201.08910
Examples are found under the examples folder in the github repo.
To install the package should be registered.
import Pkg
Pkg.add("BasicReservoirComputing")
if not then add from github
import Pkg
Pkg.add("git@github.com:japlatt/DynamicalRC.git")
The code is built around the rc object
struct rc
"""Structure holding the rc parameters and matrixes
Args:
D::Int : Input system Dimension
N::Int : Reservoir Dimension
SR::Float : Spectral Radius
ρA::Float : Density of A
α::Float : Leak Rate
σ::Float : Input Strength
random_state::Int : Random seed
σb::Float : Bias
β::Float : Tikhonov regularization
(params)::NamedTuple : Named tuple with fields SR, α, σ, σb, β
"""
function rc(D, N, ρA, Δt, SR, α, σ, σb, β;
random_state=111111)
end
function rc(D, N, ρA, Δt, params::NamedTuple; random_state=111111)
rng = MersenneTwister(random_state)
end
end
To initialize the rc we can either specify all the parameters or pass in a few along with a named tuple which contains the fields (SR, α, σ, σb, β). This functionality is mainly to help with integration with the optimization code
Once the rc has been initialized we can use the function train_RC
function train_RC(rc::rc, u; spinup = 100, readout = false)
"""Train the rc
This function sets the matrix rc.Wout
u ∈ D×T
r ∈ N×T
Args:
rc : reservoir computer
u : Array of data of shape D×Time
spinup : Amount of data to use for spinup
Returns:
if readout=true: Wout*r, the readout of the rc training data
else: return the last state of r, can be used to predict
forward from the training data
"""
which takes in the rc and the training data u. u is a matrix of equally spaced data of shape (D, Time). Training will set the field Wout in the rc.
We can do a single forecast over nsteps.
function forecast_RC(rc::rc, nsteps; uspin=nothing, r0 = nothing)
"""Forecast nsteps forward in time from the end of uspin or from rc state r0
Make sure to train the RC before forecasting. Requires either uspin or r0.
If both are provided will use uspin to set r0.
"""
The functionality is that we can either give a reservoir state r0 to start the forecast or some data uspin in order to "spin up" the reservoir. See the paper for more discussion. If we'd like to test an ensemble off different test data then we can use test_RC() to return an array of valid prediction time.
function test_RC(rc::rc, test_arr; ϵ=0.3, spinup=100)
"""Test the RC on all the forecasts in test_arr
Args:
rc : reservoir computer
test_arr : array of forecasts DxT
Returns:
valid prediction time of all the tests
"""
See the paper for the formula as to how this is calculated.
In order for the RC to be useful we must find the correct parameters for the RC for the given problem. Here we use an optimization routine, again see paper for details.
opt_lower_bounds = @with_kw (SR_lb = 0.01, α_lb = 0.0, σ_lb = 1e-12, σb_lb = 0.0, β_lb = 1e-12)
opt_upper_bounds = @with_kw (SR_ub = 2.0, α_ub = 1.0, σ_ub = 1.0, σb_ub = 4.0, β_ub = 2.0)
struct opt
"""Optimization object
Args:
train_data, valid_data_list, spinup_steps, N, ρA
lb::NamedTuple : use opt_lower_bounds to set lower bounds on params
ub::NamedTuple : use opt_upped_bounds to set upper bounds on params
"""
The opt object takes in some DxT shaped training data, a list of DxT validation data and arguments for spinup steps, reservoir size/density and lower/upper bounds. There are convenience functions for the lower and upper bounds which can be passed in with the changes specified e.g., lb = opt_lower_bounds(SR_lb = 0.5) and the rest set to default.
The optimization works using the algorithm CMAES https://github.com/jbrea/CMAEvolutionStrategy.jl but will hopefully be extended at some point to use any global optimization algorithm.
Now just call the algorithm
function (f::opt)(;maxtime=200, popsize = 15, multithread = false)
"""Call future forecasting optimization routine
This routine optimizes the rc by looking at a number of long term forecasts
of the data and comparing them.
"""
so if optimization = opt(train_data, valid_data_list, spinup_steps, N, ρA) then just call optimization() to get the best parameters. Multithreading can be a bit temperamental on some systems, make sure to initialize julia with multiple threads "julia -t2" if using this option. maxtime and popsize may need to be adjusted depending on the difficulty of the problem.
An example from the examples folder is shown below
function opt_rc(N, train_data, valid_data; Δt = 0.01)
"""Find the RC parameters through optimization"""
# define some parameters
nspin = 200
ρA = 0.02
D = size(train_data)[1]
multithread = Threads.nthreads() > 1 ? true : false
# Optimize the RC
rc_opt = opt(train_data, valid_data, nspin, N, ρA,
lb = opt_lower_bounds(), ub = opt_upper_bounds())
params, cost = rc_opt(maxtime=200, multithread=multithread)
return params
end
One way to check the fidelity of the RC is to check that its Lyapunov exponents match those of the input data. See the paper or
Jason A. Platt et al. "Robust forecasting using predictive generalized synchronization in reservoir computing". In: Chaos 31 (2021), p. 123118. URL: https://doi.org/10.1063/5.0066013. for a more detailed description.
The routine here is
function Global_LEs(rc::rc, u, nspin; num_exp = 1, renorm_steps=10)
"""Find the Lyapunov Exponents of the rc
Args:
rc : reservoir computer
u : input data, must be enough for the estimate. Shorter=less accurate
nspin : spinup steps to use
num_exp : number of exponents to compute out of N
renorm_steps : How many steps before renormalizing using QR decomposition
Returns:
Lyapunov Exponents : Array
"""
calculated with the help of DynamicalSystems.jl https://juliadynamics.github.io/DynamicalSystems.jl/dev/. u is the data over which the estimate is done. renorm_steps can help speed up the computation but will reduce accuracy if one is looking for the short term exponents. Warning: the finite time LEs are not invariant over the attractor. num_exp chooses the number of LEs to calculate.
We also expose the Tangent linear model TLM()/TLM!() in case one has need of it. For instance in data assimilation algorithms.
There may be an issue with the matrix solve and multithreading. This is sometimes solved by using MKL as a backend for the matrix operations but not always. Still trying to track this one down. If you run into issues you can try changing this.