Feature request: Next Generation Methods

[timothyslau]

I found this new reservoir computing method after reading the paper below
https://github.com/quantinfo/ng-rc-paper-code

"Next generation reservoir computing | Nature Communications" https://www.nature.com/articles/s41467-021-25801-2

Is anyone able to convert the python code to Julia?

[MartinuzziFrancesco]

Yes, that is definitely a high priority implementation for ReservoirComputing.jl. There will be a slight refactorization of the code base in the following weeks, and after that I'll tackle NG-RC

[MartinuzziFrancesco]

the refactorization took a bit longer then expected, but once #318 is merged the NGRC can be obtained by

function NGRC(in_dims::Int, out_dims::Int; num_delays::Int = 2,
        stride::Int = 1, features = (x -> x,), readout_activation = identity,
        init_delay = zeros32, include_input=true)
    delay = DelayLayer(in_dims; num_delays = num_delays, stride = stride,
        init_delay = init_delay)
    feats = features isa Tuple ? features : (features,)
    nfl = NonlinearFeaturesLayer(features...; include_input = include_input)
    n_taps = in_dims * (num_delays + 1)
    n_feats = (1 + length(feats)) * n_taps
    ro = LinearReadout(n_feats => out_dims, readout_activation)

    return ReservoirChain(delay, nfl, ro)
end

I'll keep this issue open until I actually include the NGRC model in ReservoirComputing.jl

[MartinuzziFrancesco]

Something like this is a starting point to reproduce results from the paper

using OrdinaryDiffEq

#define lorenz system
function lorenz(du, u, p, t)
    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

#solve and take data
prob = ODEProblem(lorenz, [1.0f0, 0.0f0, 0.0f0], (0.0, 200.0), [10.0f0, 28.0f0, 8/3])
data = Array(solve(prob, ABM54(); dt=0.025))
shift = 300
train_len = 5000
predict_len = 1250

#one step ahead for generative prediction
input_data = data[:, shift:(shift + train_len - 1)]
μ = mean(input_data; dims=2)
σ = std(input_data; dims=2)

target_data = data[:, (shift + 1):(shift + train_len)]
test = data[:, (shift + train_len):(shift + train_len + predict_len - 1)]

input_data = (input_data .- μ) ./ σ
target_data = (target_data .- μ) ./ σ
test = (test .- μ) ./ σ

Δtarget_data = target_data .- input_data

rng = MersenneTwister(0)

function quadratic_monomials(x::AbstractVector)
    n = length(x)
    φ = similar(x, n * (n + 1) ÷ 2)
    idx = 1
    @inbounds for i in 1:n
        for j in i:n
            φ[idx] = x[i] * x[j]
            idx += 1
        end
    end
    return φ
end

const_feature = x -> [1.0f0]
id_feature = x -> x
quad_feature = quadratic_monomials

ngrc = NGRC(3, 3; num_delays = 1, stride = 1,
    features = (const_feature, id_feature, quad_feature), include_input = false)

ps, st = setup(rng, ngrc)
ps, st = train!(ngrc, input_data, Δtarget_data, ps, st; train_method=StandardRidge(2.5e-6))

u = copy(test[:, 1])
st_pred = st
traj = similar(test, 3, predict_len)

for t in 1:predict_len
    δu, st_pred = ngrc(u, ps, st_pred)
    u .= u .+ δu
    traj[:, t] .= u
end

traj

using Plots
plot(transpose(traj); layout=(3, 1), label="predicted");
plot!(transpose(test); layout=(3, 1), label="actual")