feat: improve DeepESN, add better typing to models

Author: MartinuzziFrancesco

Project.toml

@@ -7,7 +7,6 @@ version = "0.11.4"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
-Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
@@ -35,7 +34,6 @@ CellularAutomata = "0.0.6"
 Compat = "4.16.0"
 ConcreteStructs = "0.2.3"
 DifferentialEquations = "7.16.1"
-Functors = "0.5.2"
 JET = "0.9.20"
 LIBSVM = "0.8"
 LinearAlgebra = "1.10"

src/ReservoirComputing.jl

@@ -6,20 +6,22 @@ using ConcreteStructs: @concrete
 #using Functors
 using LinearAlgebra: eigvals, mul!, I, qr, Diagonal, diag
 using LuxCore: AbstractLuxLayer, AbstractLuxContainerLayer, AbstractLuxWrapperLayer,
-               setup, apply, replicate
+    setup, apply, replicate
 import LuxCore: initialparameters, initialstates, statelength, outputsize
 using NNlib: fast_act, sigmoid
 using Random: Random, AbstractRNG, randperm
 using Static: StaticBool, StaticInt, StaticSymbol,
-              True, False, static, known, dynamic, StaticInteger
+    True, False, static, known, dynamic, StaticInteger
 using Reexport: Reexport, @reexport
 using WeightInitializers: DeviceAgnostic, PartialFunction, Utils
 @reexport using WeightInitializers
 @reexport using LuxCore: setup, apply, initialparameters, initialstates
 
-const BoolType = Union{StaticBool, Bool, Val{true}, Val{false}}
-const InputType = Tuple{<:AbstractArray, Tuple{<:AbstractArray}}
-const IntegerType = Union{Integer, StaticInteger}
+const BoolType = Union{StaticBool,Bool,Val{true},Val{false}}
+const InputType = Tuple{<:AbstractArray,Tuple{<:AbstractArray}}
+const IntegerType = Union{Integer,StaticInteger}
+
+abstract type AbstractReservoirComputer{Fields} <: AbstractLuxContainerLayer{Fields} end
 
 #@compat(public, (initialparameters)) #do I need to add intialstates/parameters in compat?
 
@@ -36,26 +38,27 @@ include("train.jl")
 include("inits/inits_components.jl")
 include("inits/esn_inits.jl")
 #full models
+include("models/esn_utils.jl")
 include("models/esn.jl")
 include("models/deepesn.jl")
 include("models/hybridesn.jl")
 #extensions
 include("extensions/reca.jl")
 
 export ESNCell, StatefulLayer, LinearReadout, ReservoirChain, Collect, collectstates,
-       train!,
-       predict, reset_carry
+    train!,
+    predict, resetcarry
 export SVMReadout
 export Pad, Extend, NLAT1, NLAT2, NLAT3, PartialSquare, ExtendedSquare
 export StandardRidge
 export chebyshev_mapping, informed_init, logistic_mapping, minimal_init,
-       modified_lm, scaled_rand, weighted_init, weighted_minimal
+    modified_lm, scaled_rand, weighted_init, weighted_minimal
 export block_diagonal, chaotic_init, cycle_jumps, delay_line, delay_line_backward,
-       double_cycle, forward_connection, low_connectivity, pseudo_svd, rand_sparse,
-       selfloop_cycle, selfloop_delayline_backward, selfloop_feedback_cycle,
-       selfloop_forward_connection, simple_cycle, true_double_cycle
+    double_cycle, forward_connection, low_connectivity, pseudo_svd, rand_sparse,
+    selfloop_cycle, selfloop_delayline_backward, selfloop_feedback_cycle,
+    selfloop_forward_connection, simple_cycle, true_double_cycle
 export add_jumps!, backward_connection!, delay_line!, reverse_simple_cycle!,
-       scale_radius!, self_loop!, simple_cycle!
+    scale_radius!, self_loop!, simple_cycle!
 export train
 export ESN, HybridESN, KnowledgeModel, DeepESN
 #reca

src/models/deepesn.jl

@@ -1,62 +1,200 @@
-# --- helpers ---
-function _asvec(x, num_reservoirs::Int)
-    if x === ()
-        return ntuple(_ -> nothing, num_reservoirs)
-    elseif x isa Tuple || x isa AbstractVector
-        len = length(x)
-        len == num_reservoirs && return Tuple(x)
-        len == 1 && return ntuple(_ -> x[1], num_reservoirs)
-        error("Expected length $num_reservoirs or 1 for per-layer argument, got $len")
-    else
-        return ntuple(_ -> x, num_reservoirs)
-    end
+"""
+    DeepESN(in_dims::Int,
+            res_dims::AbstractVector{<:Int},
+            out_dims,
+            activation=tanh;
+            leak_coefficient=1.0,
+            init_reservoir=rand_sparse,
+            init_input=weighted_init,
+            init_bias=zeros32,
+            init_state=randn32,
+            use_bias=false,
+            state_modifiers=(),
+            readout_activation=identity)
+
+Build a deep ESN: a stack of `StatefulLayer(ESNCell)` with optional per-layer
+state modifiers, followed by a final linear readout.
+"""
+@concrete struct DeepESN <: AbstractEchoStateNetwork{(:cells, :states_modifiers, :readout)}
+    cells
+    states_modifiers
+    readout
 end
 
-function DeepESN(in_dims::Int,
-    res_dims::AbstractVector{<:Int},
-    out_dims,
+function DeepESN(in_dims::IntegerType,
+    res_dims::AbstractVector{<:IntegerType},
+    out_dims::IntegerType,
     activation=tanh;
     leak_coefficient=1.0,
     init_reservoir=rand_sparse,
-    init_input=weighted_init,
+    init_input=scaled_rand,
     init_bias=zeros32,
     init_state=randn32,
     use_bias=false,
     state_modifiers=(),
     readout_activation=identity)
 
-    num_reservoirs = length(res_dims)
+    n_layers = length(res_dims)
+    acts = _asvec(activation, n_layers)
+    leaks = _asvec(leak_coefficient, n_layers)
+    ires = _asvec(init_reservoir, n_layers)
+    iinp = _asvec(init_input, n_layers)
+    ibias = _asvec(init_bias, n_layers)
+    ist = _asvec(init_state, n_layers)
+    ub = _asvec(use_bias, n_layers)
+    mods0 = _asvec(state_modifiers, n_layers)
 
-    acts = _asvec(activation, num_reservoirs)
-    leaksv = _asvec(leak_coefficient, num_reservoirs)
-    inres = _asvec(init_reservoir, num_reservoirs)
-    ininp = _asvec(init_input, num_reservoirs)
-    inbias = _asvec(init_bias, num_reservoirs)
-    inst = _asvec(init_state, num_reservoirs)
-    ubias = _asvec(use_bias, num_reservoirs)
-    mods = _asvec(state_modifiers, num_reservoirs)
+    cells = Vector{Any}(undef, n_layers)
+    states_modifiers = Vector{Any}(undef, n_layers)
 
-    layers = Any[]
     prev = in_dims
-    for res in 1:num_reservoirs
-        cell = ESNCell(prev => res_dims[res], acts[res];
-            use_bias=static(ubias[res]),
-            init_bias=inbias[res],
-            init_reservoir=inres[res],
-            init_input=ininp[res],
-            init_state=inst[res],
-            leak_coefficient=leaksv[res])
-
-        push!(layers, StatefulLayer(cell))
-        if mods[res] !== nothing
-            push!(layers, mods[res])
-        end
-        prev = res_dims[res]
+    for idx in firstindex(res_dims):lastindex(res_dims)
+        cell = ESNCell(prev => res_dims[idx], acts[idx];
+            use_bias=static(ub[idx]),
+            init_bias=ibias[idx],
+            init_reservoir=ires[idx],
+            init_input=iinp[idx],
+            init_state=ist[idx],
+            leak_coefficient=leaks[idx])
+        cells[idx] = StatefulLayer(cell)
+        states_modifiers[idx] = mods0[idx] === nothing ? nothing : _wrap_layer(mods0[idx])
+        prev = res_dims[idx]
     end
+    mods_per_layer = map(_coerce_layer_mods, states_modifiers) |> Tuple
     ro = LinearReadout(prev => out_dims, readout_activation)
-    return ReservoirChain((layers..., ro)...)
+    return DeepESN(Tuple(cells), mods_per_layer, ro)
+end
+
+DeepESN(in_dims::Int, res_dim::Int, out_dims::Int; depth::Int=2, kwargs...) =
+    DeepESN(in_dims, fill(res_dim, depth), out_dims; kwargs...)
+
+function initialparameters(rng::AbstractRNG, desn::DeepESN)
+    ps_cells = map(l -> initialparameters(rng, l), desn.cells) |> Tuple
+    mods = desn.states_modifiers === nothing ? ntuple(_ -> (), length(desn.cells)) :
+           desn.states_modifiers
+    ps_mods = map(layer_mods ->
+            (layer_mods === nothing ? () :
+             map(l -> initialparameters(rng, l), layer_mods) |> Tuple),
+        mods) |> Tuple
+
+    ps_ro = initialparameters(rng, desn.readout)
+    return (cells=ps_cells, states_modifiers=ps_mods, readout=ps_ro)
+end
+
+function initialstates(rng::AbstractRNG, desn::DeepESN)
+    st_cells = map(l -> initialstates(rng, l), desn.cells) |> Tuple
+
+    mods = desn.states_modifiers === nothing ? ntuple(_ -> (), length(desn.cells)) :
+           desn.states_modifiers
+
+    st_mods = map(layer_mods ->
+            (layer_mods === nothing ? () :
+             map(l -> initialstates(rng, l), layer_mods) |> Tuple),
+        mods) |> Tuple
+
+    st_ro = initialstates(rng, desn.readout)
+    return (cells=st_cells, states_modifiers=st_mods, readout=st_ro)
+end
+
+function (desn::DeepESN)(inp, ps, st)
+    inp_t = inp
+    n_layers = length(desn.cells)
+    new_cell_st = Vector{Any}(undef, n_layers)
+    new_mods_st = Vector{Any}(undef, n_layers)
+    for idx in firstindex(desn.cells):lastindex(desn.cells)
+        inp_t, st_cell_i = apply(desn.cells[idx], inp_t, ps.cells[idx], st.cells[idx])
+        new_cell_st[idx] = st_cell_i
+        inp_t, st_mods_i = _apply_seq(desn.states_modifiers[idx], inp_t,
+            ps.states_modifiers[idx], st.states_modifiers[idx])
+        new_mods_st[idx] = st_mods_i
+    end
+    inp_t, st_ro = apply(desn.readout, inp_t, ps.readout, st.readout)
+
+    return inp_t, (;
+        cells=tuple(new_cell_st...),
+        states_modifiers=tuple(new_mods_st...),
+        readout=st_ro,
+    )
 end
 
-function DeepESN(in_dims::Int, res_dims::Int, out_dims::Int; depth::Int=2, kwargs...)
-    return DeepESN(in_dims, fill(res_dims, depth), out_dims; kwargs...)
+function resetcarry(rng::AbstractRNG, desn::DeepESN, st; init_carry=nothing)
+    n_layers = length(desn.cells)
+
+    @inline function _layer_outdim(idx)
+        st_i = st.cells[idx]
+        if st_i.carry === nothing
+            return desn.cells[idx].cell.out_dims
+        else
+            return size(first(st_i.carry), 1)
+        end
+    end
+
+    @inline function _init_for(idx)
+        if init_carry === nothing
+            return nothing
+        elseif init_carry isa Function
+            sz = _layer_outdim(idx)
+            return (_asvec(init_carry(rng, sz)),)
+        elseif init_carry isa Tuple || init_carry isa AbstractVector
+            f = init_carry[idx]
+            sz = _layer_outdim(idx)
+            return f === nothing ? nothing : (_asvec(f(rng, sz)),)
+        else
+            throw(ArgumentError("init_carry must be nothing, a Function, or a Tuple/Vector of Functions"))
+        end
+    end
+
+    new_cells = ntuple(idx -> begin
+            st_i = st.cells[idx]
+            new_carry = _init_for(idx)
+            merge(st_i, (; carry=new_carry))
+        end, n_layers)
+
+    return (;
+        cells=new_cells,
+        states_modifiers=st.states_modifiers,
+        readout=st.readout,
+    )
+end
+
+function collectstates(desn::DeepESN, data::AbstractMatrix, ps, st::NamedTuple)
+    newst = st
+    collected = Any[]
+    n_layers = length(desn.cells)
+    for inp in eachcol(data)
+        inp_t = inp
+        cell_st_parts = Vector{Any}(undef, n_layers)
+        mods_st_parts = Vector{Any}(undef, n_layers)
+        for idx in firstindex(desn.cells):lastindex(desn.cells)
+            inp_t, st_cell_i = apply(desn.cells[idx], inp_t, ps.cells[idx], newst.cells[idx])
+            cell_st_parts[idx] = st_cell_i
+            inp_t, st_mods_i = _apply_seq(
+                desn.states_modifiers[idx], inp_t,
+                ps.states_modifiers[idx], newst.states_modifiers[idx]
+            )
+            mods_st_parts[idx] = st_mods_i
+        end
+        push!(collected, copy(inp_t))
+        newst = (;
+            cells=tuple(cell_st_parts...),
+            states_modifiers=tuple(mods_st_parts...),
+            readout=newst.readout,
+        )
+    end
+    @assert !isempty(collected)
+    states = eltype(data).(reduce(hcat, collected))
+
+    return states, newst
+end
+
+collectstates(m::DeepESN, data::AbstractVector, ps, st::NamedTuple) =
+    collectstates(m, reshape(data, :, 1), ps, st)
+
+function addreadout!(::DeepESN, output_matrix::AbstractMatrix,
+    ps::NamedTuple, st::NamedTuple)
+    @assert hasproperty(ps, :readout)
+    new_readout = _set_readout_weight(ps.readout, output_matrix)
+    return (cells=ps.cells,
+        states_modifiers=ps.states_modifiers,
+        readout=new_readout), st
 end