refac: remove old ESN and DeepESN, implement new APIs

Author: MartinuzziFrancesco

README.md

@@ -105,19 +105,20 @@ possible parameters:
 ```julia
 input_size = 3
 res_size = 300
-esn = ESN(input_data, input_size, res_size;
-    reservoir=rand_sparse(; radius=1.2, sparsity=6 / res_size),
-    input_layer=weighted_init,
-    nla_type=NLAT2(),
-    rng=rng)
+esn = ReservoirChain(
+    StatefulLayer(ESNCell(input_size => res_size; init_reservoir=rand_sparse(; radius=1.2, sparsity=6/300))),
+    NLAT2(),
+    Readout(res_size => input_size)
+)
 ```
 
 The echo state network can now be trained and tested.
 If not specified, the training will always be ordinary least squares regression:
 
 ```julia
-output_layer = train(esn, target_data)
-output = esn(Generative(predict_len), output_layer)
+ps, st = setup(rng, esn)
+ps, st = train!(esn, input_data, target_data, ps, st)
+output, _ = predict(esn, 1250, ps, st; initialdata=test[:, 1])
 ```
 
 The data is returned as a matrix, `output` in the code above,
@@ -126,7 +127,7 @@ The results can now be easily plotted:
 
 ```julia
 using Plots
-plot(transpose(output); layout=(3, 1), label="predicted")
+plot(transpose(output); layout=(3, 1), label="predicted");
 plot!(transpose(test); layout=(3, 1), label="actual")
 ```
 

src/ReservoirComputing.jl

@@ -32,16 +32,15 @@ include("layers/esn_cell.jl")
 include("generics/states.jl")
 include("generics/predict.jl")
 include("generics/linear_regression.jl")
-#extensions
-include("extensions/reca.jl")
 #esn
 include("inits/inits_components.jl")
 include("inits/esn_inits.jl")
-include("layers/esn_reservoir_drivers.jl")
 include("models/esn.jl")
 include("models/deepesn.jl")
 include("models/hybridesn.jl")
-include("models/esn_predict.jl")
+#extensions
+include("extensions/reca.jl")
+
 
 
 
@@ -58,10 +57,8 @@ export block_diagonal, chaotic_init, cycle_jumps, delay_line, delay_line_backwar
     selfloop_forward_connection, simple_cycle, true_double_cycle
 export add_jumps!, backward_connection!, delay_line!, reverse_simple_cycle!,
     scale_radius!, self_loop!, simple_cycle!
-export RNN, MRNN, GRU, GRUParams, FullyGated, Minimal
 export train
 export ESN, HybridESN, KnowledgeModel, DeepESN
-export Generative, Predictive, OutputLayer
 #reca
 export RECA
 export RandomMapping, RandomMaps

src/generics/linear_regression.jl

@@ -29,18 +29,12 @@ function StandardRidge()
     return StandardRidge(0.0)
 end
 
-function train!(rc::ReservoirChain, train_data::AbstractArray,
-    target_data::AbstractArray, ps, st::NamedTuple, sr::StandardRidge=StandardRidge(0.0);
+function train!(rc::ReservoirChain, train_data, target_data, ps, st, sr=StandardRidge(0.0);
     return_states::Bool=false)
-    states = collectstates(rc, train_data, ps, st)
-    readout = train(sr, states, target_data)
-    ps, st = addreadout!(rc, readout, ps, st)
-
-    if return_states
-        return (ps, st), states
-    else
-        return ps, st
-    end
+    states, new_st = collectstates(rc, train_data, ps, st)
+    W = train(sr, states, target_data)
+    ps2, _ = addreadout!(rc, W, ps, new_st)
+    return return_states ? ((ps2, new_st), states) : (ps2, new_st)
 end
 
 function train(sr::StandardRidge, states::AbstractArray, target_data::AbstractArray)
@@ -53,16 +47,50 @@ function train(sr::StandardRidge, states::AbstractArray, target_data::AbstractAr
     return output_layer
 end
 
-function addreadout!(rc::ReservoirChain, readout_matrix::AbstractArray, ps, st::NamedTuple) #make sure the compile infers
-    ro_param = (; weight=readout_matrix)
-    new_ps = (;)
-    for ((name, layer), param) in zip(pairs(rc.layers), ps)
-        if layer isa Readout
-            param = merge(param, ro_param)
-        end
-        new_ps = merge(new_ps, (; name => param))
+_quote_keys(t) = Expr(:tuple, (QuoteNode(s) for s in t)...)
+
+@generated function _setweight_rt(p::NamedTuple{K}, W) where {K}
+    keys = K
+    Kq = _quote_keys(keys)
+    idx = findfirst(==(Symbol(:weight)), keys)
+
+    terms = Any[]
+    for i in 1:length(keys)
+        push!(terms, (idx === i) ? :(W) : :(getfield(p, $i)))
+    end
+
+    if idx === nothing
+        newK = _quote_keys((keys..., :weight))
+        return :(NamedTuple{$newK}(($(terms...), W)))
+    else
+        return :(NamedTuple{$Kq}(($(terms...),)))
     end
-    return new_ps, st
 end
 
-#use a recursion to make it more compiler safe
+@generated function _addreadout(layers::NamedTuple{K}, ps::NamedTuple{K}, W) where {K}
+    if length(K) == 0
+        return :(NamedTuple())
+    end
+    tailK = Base.tail(K)
+    Kq = _quote_keys(K)
+    tailKq = _quote_keys(tailK)
+
+    head_val = :((getfield(layers, 1) isa Readout)
+                 ? _setweight_rt(getfield(ps, 1), W)
+                 : getfield(ps, 1))
+
+    tail_call = :(_addreadout(NamedTuple{$tailKq}(Base.tail(layers)),
+        NamedTuple{$tailKq}(Base.tail(ps)),
+        W))
+
+    return :(NamedTuple{$Kq}(($head_val, Base.values($tail_call)...)))
+end
+
+function addreadout!(rc::ReservoirChain,
+    W::AbstractMatrix,
+    ps::NamedTuple,
+    st::NamedTuple)
+    @assert propertynames(rc.layers) == propertynames(ps)
+    new_ps = _addreadout(rc.layers, ps, W)
+    return new_ps, st
+end

src/generics/predict.jl

@@ -1,123 +1,3 @@
-abstract type AbstractOutputLayer end
-abstract type AbstractPrediction end
-
-#general output layer struct
-struct OutputLayer{T,I,S,L} <: AbstractOutputLayer
-    training_method::T
-    output_matrix::I
-    out_size::S
-    last_value::L
-end
-
-function Base.show(io::IO, ol::OutputLayer)
-    print(io, "OutputLayer successfully trained with output size: ", ol.out_size)
-end
-
-#prediction types
-"""
-    Generative(prediction_len)
-
-A prediction strategy that enables models to generate autonomous multi-step
-forecasts by recursively feeding their own outputs back as inputs for
-subsequent prediction steps.
-
-# Parameters
-
-  - `prediction_len`: The number of future steps to predict.
-
-# Description
-
-The `Generative` prediction method allows a model to perform multi-step
-forecasting by using its own previous predictions as inputs for future predictions.
-
-At each step, the model takes the current input, generates a prediction,
-and then incorporates that prediction into the input for the next step.
-This recursive process continues until the specified
-number of prediction steps (`prediction_len`) is reached.
-"""
-struct Generative{T} <: AbstractPrediction
-    prediction_len::T
-end
-
-struct Predictive{I,T} <: AbstractPrediction
-    prediction_data::I
-    prediction_len::T
-end
-
-"""
-    Predictive(prediction_data)
-
-A prediction strategy for supervised learning tasks,
-where a model predicts labels based on a provided set
-of input features (`prediction_data`).
-
-# Parameters
-
-  - `prediction_data`: The input data used for prediction, `feature` x `sample`
-
-# Description
-
-The `Predictive` prediction method uses the provided input data
-(`prediction_data`) to produce corresponding labels or outputs based
-on the learned relationships in the model.
-"""
-function Predictive(prediction_data::AbstractArray)
-    prediction_len = size(prediction_data, 2)
-    return Predictive(prediction_data, prediction_len)
-end
-
-function obtain_prediction(rc::AbstractReservoirComputer, prediction::Generative,
-    x, output_layer::AbstractOutputLayer, args...;
-    initial_conditions=output_layer.last_value)
-    #x = last_state
-    prediction_len = prediction.prediction_len
-    train_method = output_layer.training_method
-    out_size = output_layer.out_size
-    output = output_storing(train_method, out_size, prediction_len, typeof(rc.states))
-    out = initial_conditions
-
-    for i in 1:prediction_len
-        x, x_new = next_state_prediction!(rc, x, out, i, args...)
-        out_tmp = get_prediction(train_method, output_layer, x_new)
-        out = store_results!(train_method, out_tmp, output, i)
-    end
-
-    return output
-end
-
-function obtain_prediction(rc::AbstractReservoirComputer, prediction::Predictive,
-    x, output_layer::AbstractOutputLayer, args...; kwargs...)
-    prediction_len = prediction.prediction_len
-    train_method = output_layer.training_method
-    out_size = output_layer.out_size
-    output = output_storing(train_method, out_size, prediction_len, typeof(rc.states))
-
-    for i in 1:prediction_len
-        y = @view prediction.prediction_data[:, i]
-        x, x_new = next_state_prediction!(rc, x, y, i, args...)
-        out_tmp = get_prediction(train_method, output_layer, x_new)
-        out = store_results!(output_layer.training_method, out_tmp, output, i)
-    end
-
-    return output
-end
-
-#linear models
-function get_prediction(training_method, output_layer::AbstractOutputLayer, x)
-    return output_layer.output_matrix * x
-end
-
-#single matrix for other training methods
-function output_storing(training_method, out_size, prediction_len, storing_type)
-    return adapt(storing_type, zeros(out_size, prediction_len))
-end
-
-#general storing -> single matrix
-function store_results!(training_method, out, output, i)
-    output[:, i] = out
-    return out
-end
-
 function predict(rc, steps::Int, ps, st; initialdata=nothing)
     if initialdata == nothing
         initialdata = rand(Float32, 3)

src/layers/esn_cell.jl

@@ -1,3 +1,67 @@
+@doc raw"""
+    ESNCell(in_dims => out_dims, [activation];
+        use_bias=false, init_bias=rand32,
+        init_reservoir=rand_sparse, init_input=weighted_init,
+        init_state=randn32, leak_coefficient=1.0)
+
+Echo State Network (ESN) recurrent cell with optional leaky integration.
+
+## Equations
+
+```math
+\begin{aligned}
+    \tilde{\mathbf{h}}(t) &= \phi\!\left(\mathbf{W}_{in}\,\mathbf{x}(t) +
+        \mathbf{W}_{res}\,\mathbf{h}(t-1) + \mathbf{b}\right) \\
+    \mathbf{h}(t) &= (1-\alpha)\,\mathbf{h}(t-1) + \alpha\,\tilde{\mathbf{h}}(t)
+\end{aligned}
+```
+## Arguments
+
+  - `in_dims`: Input dimension.
+  - `out_dims`: Reservoir (hidden state) dimension.
+  - `activation`: Activation function. Default: `tanh`.
+
+## Keyword arguments
+
+  - `use_bias`: Whether to include a bias term. Default: `false`.
+  - `init_bias`: Initializer for the bias. Used only if `use_bias=true`.
+      Default is `rand32`.
+  - `init_reservoir`: Initializer for the reservoir matrix `W_res`.
+    Default is [`rand_sparse`](@ref).
+  - `init_input`: Initializer for the input matrix `W_in`.
+  - `init_state`: Initializer for the hidden state when an external
+    state is not provided. Default is `randn32`.
+  - `leak_coefficient`: Leak rate `α ∈ (0,1]`. Default: `1.0`.
+
+## Inputs
+
+  - **Case 1:** `x :: AbstractArray (in_dims, batch)`
+    A fresh state is created via `init_state`; the call is forwarded to Case 2.
+  - **Case 2:** `(x, (h,))` where `h :: AbstractArray (out_dims, batch)`
+    Computes the update and returns the new state.
+
+In both cases, the forward returns `((h_new, (h_new,)), st_out)` where `st_out`
+contains any updated internal state.
+
+## Returns
+
+  - Output/hidden state `h_new :: out_dims` and state tuple `(h_new,)`.
+  - Updated layer state (NamedTuple).
+
+## Parameters
+
+Created by `initialparameters(rng, esn)`:
+
+  - `input_matrix :: (out_dims × in_dims)` — `W_in`
+  - `reservoir_matrix :: (out_dims × out_dims)` — `W_res`
+  - `bias :: (out_dims,)` — present only if `use_bias=true`
+
+## States
+
+Created by `initialstates(rng, esn)`:
+
+  - `rng`: a replicated RNG used to sample initial hidden states when needed.
+"""
 @concrete struct ESNCell <: AbstractReservoirRecurrentCell
     activation
     in_dims <: IntegerType