Hybrid Echo State Networks

Following the idea of giving physical information to machine learning models, the hybrid echo state networks (Pathak et al., 2018) try to achieve this results by feeding model data into the ESN. In this example, it is explained how to create and leverage such models in ReservoirComputing.jl.

Generating the data

For this example, we are going to forecast the Lorenz system. As usual, the data is generated leveraging DifferentialEquations.jl:

using DifferentialEquations

u0 = [1.0, 0.0, 0.0]
tspan = (0.0, 1000.0)
datasize = 100000
tsteps = range(tspan[1], tspan[2]; length=datasize)

function lorenz(du, u, p, t)
    p = [10.0, 28.0, 8 / 3]
    du[1] = p[1] * (u[2] - u[1])
    du[2] = u[1] * (p[2] - u[3]) - u[2]
    du[3] = u[1] * u[2] - p[3] * u[3]
end

ode_prob = ODEProblem(lorenz, u0, tspan)
ode_sol = solve(ode_prob; saveat=tsteps)
ode_data = Array(ode_sol)

train_len = 10000

input_data = ode_data[:, 1:train_len]
target_data = ode_data[:, 2:(train_len + 1)]
test_data = ode_data[:, (train_len + 1):end][:, 1:1000]

predict_len = size(test_data, 2)
tspan_train = (tspan[1], ode_sol.t[train_len])

(0.0, 99.9909999099991)

Building the Hybrid Echo State Network

To feed the data to the ESN, it is necessary to create a suitable function.

function prior_model_data_generator(u0, tspan, tsteps, model=lorenz)
    prob = ODEProblem(lorenz, u0, tspan)
    sol = Array(solve(prob; saveat=tsteps))
    return sol
end

prior_model_data_generator (generic function with 2 methods)

Given the initial condition, time span, and time steps, this function returns the data for the chosen model. Now, using the KnowledgeModel method, it is possible to input all this information to HybridESN.

using ReservoirComputing, Random
Random.seed!(42)

km = KnowledgeModel(prior_model_data_generator, u0, tspan_train, train_len)

in_size = 3
res_size = 300
hesn = HybridESN(km,
    input_data,
    in_size,
    res_size;
    reservoir=rand_sparse)

HybridESN{Int64, Matrix{Float64}, KnowledgeModel{typeof(Main.prior_model_data_generator), Vector{Float64}, Tuple{Float64, Float64}, Float64, Int64, Matrix{Float64}}, NLADefault, Matrix{Float32}, RNN{typeof(NNlib.tanh_fast), Float64}, Matrix{Float32}, Vector{Float32}, StandardStates, Int64, Matrix{Float64}}(300, [1.0 0.9179237970373996 … -14.533407004892574 -14.160669947811542; 0.0 0.2663425318599414 … -11.389720568651494 -9.80136971808906; … ; 0.0 0.2663425318599414 … -11.351122193592321 -9.754515413100073; 0.0 0.001263920272114893 … 38.09226520138708 38.59080979025016], KnowledgeModel{typeof(Main.prior_model_data_generator), Vector{Float64}, Tuple{Float64, Float64}, Float64, Int64, Matrix{Float64}}(Main.prior_model_data_generator, [1.0, 0.0, 0.0], (0.0, 99.9909999099991), 0.01000010000100001, 10000, [1.0 0.9179237970373996 … -14.15636475973194 -13.660821936113981; 0.0 0.2663425318599414 … -9.754515413100073 -8.16854022779646; 0.0 0.001263920272114893 … 38.59080979025016 38.80935591958125]), NLADefault(), Float32[-0.027369846 0.06270313 … -0.057819426 -0.059205364; -0.09963039 0.08069054 … 0.09205323 -0.09963995; … ; 0.058445442 -0.04866575 … -0.032288957 -0.004228306; -0.039661705 0.025667286 … -0.056607306 0.024241282], RNN{typeof(NNlib.tanh_fast), Float64}(NNlib.tanh_fast, 1.0), Float32[0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0; … ; 0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0], Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0  …  0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], StandardStates(), 0, [-0.060532171476285215 -0.03894680842580137 … -0.9999056993887195 -0.9999051895219654; -0.04534558826291831 -0.0020470180860934876 … -0.9999698335476644 -0.9999463777914631; … ; 0.13726723587626294 0.12708944051237314 … 0.831484636957554 0.8088595977117276; -0.11582089793675862 -0.1393103388746963 … 0.9999561276090689 0.9999551878333757])

Training and Prediction

The training and prediction of the Hybrid ESN can proceed as usual:

output_layer = train(hesn, target_data, StandardRidge(0.3))
output = hesn(Generative(predict_len), output_layer)

3×1000 Matrix{Float64}:
 -13.0627   -12.3793   -11.635    -10.8469   …  10.6882   10.0021    9.30554
  -6.63501   -5.22513   -3.94791   -2.82303      3.93303   3.07455   2.34147
  38.7623    38.4863    38.0168    37.3907      36.2556   35.6568   34.9744

It is now possible to plot the results, leveraging Plots.jl:

using Plots
lorenz_maxlyap = 0.9056
predict_ts = tsteps[(train_len + 1):(train_len + predict_len)]
lyap_time = (predict_ts .- predict_ts[1]) * (1 / lorenz_maxlyap)

p1 = plot(lyap_time, [test_data[1, :] output[1, :]]; label=["actual" "predicted"],
    ylabel="x(t)", linewidth=2.5, xticks=false, yticks=-15:15:15);
p2 = plot(lyap_time, [test_data[2, :] output[2, :]]; label=["actual" "predicted"],
    ylabel="y(t)", linewidth=2.5, xticks=false, yticks=-20:20:20);
p3 = plot(lyap_time, [test_data[3, :] output[3, :]]; label=["actual" "predicted"],
    ylabel="z(t)", linewidth=2.5, xlabel="max(λ)*t", yticks=10:15:40);

plot(p1, p2, p3; plot_title="Lorenz System Coordinates",
    layout=(3, 1), xtickfontsize=12, ytickfontsize=12, xguidefontsize=15,
    yguidefontsize=15,
    legendfontsize=12, titlefontsize=20)

References