refac: finalize HybridESN, improve ESN generics

Author: MartinuzziFrancesco

README.md

@@ -106,20 +106,10 @@ We showcase the second option:
 using ReservoirComputing
 input_size = 3
 res_size = 300
-esn = ReservoirChain(
-    StatefulLayer(
-        ESNCell(
-            input_size => res_size;
-            init_reservoir=rand_sparse(; radius=1.2, sparsity=6/300)
-        )
-    ),
-    NLAT2(),
-    Readout(res_size => input_size) # autoregressive so out_dims == in_dims
+esn = ESN(input_size, res_size, input_size;
+    init_reservoir=rand_sparse(; radius=1.2, sparsity=6/300),
+    state_modifiers=NLAT2
 )
-# alternative:
-# esn = ESN(input_size, res_size, input_size;
-#     init_reservoir=rand_sparse(; radius=1.2, sparsity=6/300)
-# )
 ```
 
 ### 3. Train the Echo State Network

ext/RCCellularAutomataExt.jl

@@ -5,42 +5,6 @@ import ReservoirComputing: RECACell, RECA
 using CellularAutomata
 using Random: randperm
 
-function (reca::RECACell)((inp, (ca_prev,)), ps, st::NamedTuple)
-    rm = reca.enc
-    T = eltype(inp)
-    ca0 = T.(encoding(rm, inp, T.(ca_prev)))
-    ca = CellularAutomaton(reca.automaton, ca0, rm.generations + 1)
-    evo = ca.evolution
-    feat2T = evo[2:end, :]
-    feats = reshape(permutedims(feat2T), rm.states_size)
-    ca_last = evo[end, :]
-    return (T.(feats), (T.(ca_last),)), st
-end
-
-function (reca::RECACell)(inp::AbstractVector, ps, st::NamedTuple)
-    ca = st.ca
-    return reca((inp, (ca,)), ps, st)
-end
-
-function RECA(in_dims::IntegerType,
-    out_dims::IntegerType,
-    automaton;
-    input_encoding::AbstractInputEncoding=RandomMapping(),
-    generations::Integer=8,
-    state_modifiers=(),
-    readout_activation=identity)
-
-    rm = create_encoding(input_encoding, in_dims, generations)
-    cell = RECACell(automaton, rm)
-
-    mods = state_modifiers isa Tuple || state_modifiers isa AbstractVector ?
-           Tuple(state_modifiers) : (state_modifiers,)
-
-    ro = Readout(rm.states_size => out_dims, readout_activation)
-
-    return ReservoirChain((StatefulLayer(cell), mods..., ro)...)
-end
-
 function RandomMapping(; permutations=8, expansion_size=40)
     RandomMapping(permutations, expansion_size)
 end
@@ -56,23 +20,6 @@ function create_encoding(rm::RandomMapping, in_dims::IntegerType, generations::I
     return RandomMaps(rm.permutations, rm.expansion_size, generations, maps, states_size, ca_size)
 end
 
-
-function reca_create_states(rm::RandomMaps, automata, input_data)
-    train_time = size(input_data, 2)
-    states = zeros(rm.states_size, train_time)
-    init_ca = zeros(rm.expansion_size * rm.permutations)
-
-    for i in 1:train_time
-        init_ca = encoding(rm, input_data[:, i], init_ca)
-        ca = CellularAutomaton(automata, init_ca, rm.generations + 1)
-        ca_states = ca.evolution[2:end, :]
-        states[:, i] = reshape(transpose(ca_states), rm.states_size)
-        init_ca = ca.evolution[end, :]
-    end
-
-    return states
-end
-
 function encoding(rm::RandomMaps, input_vector, tot_encoded_vector)
     input_size = size(input_vector, 1)
     #single_encoded_size = Int(size(tot_encoded_vector, 1)/permutations)
@@ -119,4 +66,40 @@ function mapping(input_size, mapped_vector_size)
     return randperm(mapped_vector_size)[1:input_size]
 end
 
+function (reca::RECACell)((inp, (ca_prev,)), ps, st::NamedTuple)
+    rm = reca.enc
+    T = eltype(inp)
+    ca0 = T.(encoding(rm, inp, T.(ca_prev)))
+    ca = CellularAutomaton(reca.automaton, ca0, rm.generations + 1)
+    evo = ca.evolution
+    feat2T = evo[2:end, :]
+    feats = reshape(permutedims(feat2T), rm.states_size)
+    ca_last = evo[end, :]
+    return (T.(feats), (T.(ca_last),)), st
+end
+
+function (reca::RECACell)(inp::AbstractVector, ps, st::NamedTuple)
+    ca = st.ca
+    return reca((inp, (ca,)), ps, st)
+end
+
+function RECA(in_dims::IntegerType,
+    out_dims::IntegerType,
+    automaton;
+    input_encoding::AbstractInputEncoding=RandomMapping(),
+    generations::Integer=8,
+    state_modifiers=(),
+    readout_activation=identity)
+
+    rm = create_encoding(input_encoding, in_dims, generations)
+    cell = RECACell(automaton, rm)
+
+    mods = state_modifiers isa Tuple || state_modifiers isa AbstractVector ?
+           Tuple(state_modifiers) : (state_modifiers,)
+
+    ro = Readout(rm.states_size => out_dims, readout_activation)
+
+    return ReservoirChain((StatefulLayer(cell), mods..., ro)...)
+end
+
 end #module

src/ReservoirComputing.jl

@@ -3,7 +3,6 @@ module ReservoirComputing
 using ArrayInterface: ArrayInterface
 using Compat: @compat
 using ConcreteStructs: @concrete
-#using Functors
 using LinearAlgebra: eigvals, mul!, I, qr, Diagonal, diag
 using LuxCore: AbstractLuxLayer, AbstractLuxContainerLayer, AbstractLuxWrapperLayer,
     setup, apply, replicate
@@ -20,6 +19,7 @@ using WeightInitializers: DeviceAgnostic, PartialFunction, Utils
 const BoolType = Union{StaticBool,Bool,Val{true},Val{false}}
 const InputType = Tuple{<:AbstractArray,Tuple{<:AbstractArray}}
 const IntegerType = Union{Integer,StaticInteger}
+const RCFields = (:cells, :states_modifiers, :readout)
 
 abstract type AbstractReservoirComputer{Fields} <: AbstractLuxContainerLayer{Fields} end
 
@@ -38,10 +38,10 @@ include("train.jl")
 include("inits/inits_components.jl")
 include("inits/esn_inits.jl")
 #full models
-include("models/esn_utils.jl")
+include("models/esn_generics.jl")
 include("models/esn.jl")
-include("models/deepesn.jl")
-include("models/hybridesn.jl")
+include("models/esn_deep.jl")
+include("models/esn_hybrid.jl")
 #extensions
 include("extensions/reca.jl")
 
@@ -60,7 +60,7 @@ export block_diagonal, chaotic_init, cycle_jumps, delay_line, delay_line_backwar
 export add_jumps!, backward_connection!, delay_line!, reverse_simple_cycle!,
     scale_radius!, self_loop!, simple_cycle!
 export train
-export ESN, HybridESN, KnowledgeModel, DeepESN
+export ESN, HybridESN, DeepESN
 #reca
 export RECACell, RECA
 export RandomMapping, RandomMaps

src/layers/esn_cell.jl

@@ -77,7 +77,7 @@ end
 
 function ESNCell((in_dims, out_dims)::Pair{<:IntegerType,<:IntegerType}, activation=tanh;
     use_bias::BoolType=False(), init_bias=zeros32, init_reservoir=rand_sparse,
-    init_input=weighted_init, init_state=randn32, leak_coefficient=1.0)
+    init_input=scaled_rand, init_state=randn32, leak_coefficient=1.0)
     return ESNCell(activation, in_dims, out_dims, init_bias, init_reservoir,
         init_input, init_state, leak_coefficient, use_bias)
 end

src/models/esn.jl

@@ -1,3 +1,11 @@
+"""
+    ESN(in_dims, res_dims, out_dims, activation=tanh;
+        leak_coefficient=1.0, init_reservoir=rand_sparse, init_input=scaled_rand,
+        init_bias=zeros32, init_state=randn32, use_bias=false,
+        state_modifiers=(), readout_activation=identity)
+
+Build a ESN.
+"""
 @concrete struct ESN <: AbstractEchoStateNetwork{(:cell, :states_modifiers, :readout)}
     cell
     states_modifiers
@@ -31,12 +39,17 @@ function initialstates(rng::AbstractRNG, esn::ESN)
     return (cell=st_cell, states_modifiers=st_mods, readout=st_ro)
 end
 
-function (esn::ESN)(inp, ps, st)
+function _partial_apply(esn::ESN, inp, ps, st)
     out, st_cell = apply(esn.cell, inp, ps.cell, st.cell)
     out, st_mods = _apply_seq(
         esn.states_modifiers, out, ps.states_modifiers, st.states_modifiers)
+    return out, (cell=st_cell, states_modifiers=st_mods)
+end
+
+function (esn::ESN)(inp, ps, st)
+    out, new_st = _partial_apply(esn, inp, ps, st)
     out, st_ro = apply(esn.readout, out, ps.readout, st.readout)
-    return out, (cell=st_cell, states_modifiers=st_mods, readout=st_ro)
+    return out, merge(new_st, (readout=st_ro,))
 end
 
 function resetcarry(rng::AbstractRNG, esn::ESN, st; init_carry=nothing)
@@ -59,29 +72,3 @@ function resetcarry(rng::AbstractRNG, esn::ESN, st; init_carry=nothing)
     new_cell = merge(st.cell, (; carry=new_state))
     return (cell=new_cell, states_modifiers=st.states_modifiers, readout=st.readout)
 end
-
-function collectstates(esn::ESN, data::AbstractMatrix, ps, st::NamedTuple)
-    newst = st
-    collected = Any[]
-    for inp in eachcol(data)
-        cell_y, st_cell = apply(esn.cell, inp, ps.cell, newst.cell)
-        state_t, st_mods = _apply_seq(
-            esn.states_modifiers, cell_y, ps.states_modifiers, newst.states_modifiers)
-        push!(collected, copy(state_t))
-        newst = (cell=st_cell, states_modifiers=st_mods, readout=newst.readout)
-    end
-    states = eltype(data).(reduce(hcat, collected))
-    @assert !isempty(collected)
-    states_raw = reduce(hcat, collected)
-    states = eltype(data).(states_raw)
-    return states, newst
-end
-
-function addreadout!(::ESN, output_matrix::AbstractMatrix,
-    ps::NamedTuple, st::NamedTuple)
-    @assert hasproperty(ps, :readout)
-    new_readout = _set_readout_weight(ps.readout, output_matrix)
-    return (cell=ps.cell,
-        states_modifiers=ps.states_modifiers,
-        readout=new_readout), st
-end

src/models/deepesn.jl renamed to src/models/esn_deep.jl

@@ -5,7 +5,7 @@
             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,
@@ -96,7 +96,7 @@ function initialstates(rng::AbstractRNG, desn::DeepESN)
     return (cells=st_cells, states_modifiers=st_mods, readout=st_ro)
 end
 
-function (desn::DeepESN)(inp, ps, st)
+function _partial_apply(desn::DeepESN, inp, ps, st)
     inp_t = inp
     n_layers = length(desn.cells)
     new_cell_st = Vector{Any}(undef, n_layers)
@@ -108,15 +108,19 @@ function (desn::DeepESN)(inp, ps, st)
             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 (desn::DeepESN)(inp, ps, st)
+    out, new_st = _partial_apply(desn, inp, ps, st)
+    inp_t, st_ro = apply(desn.readout, out, ps.readout, st.readout)
+    return inp_t, merge(new_st, (readout=st_ro,))
+end
+
 function resetcarry(rng::AbstractRNG, desn::DeepESN, st; init_carry=nothing)
     n_layers = length(desn.cells)
 
@@ -189,12 +193,3 @@ 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

src/models/esn_utils.jl renamed to src/models/esn_generics.jl

@@ -43,3 +43,25 @@ _set_readout_weight(ps_readout::NamedTuple, wro) = merge(ps_readout, (; weight=w
 function resetcarry(rng::AbstractRNG, esn, ps, st; init_carry=nothing)
     return ps, resetcarry(rng, esn, st; init_carry=init_carry)
 end
+
+function collectstates(esn::AbstractEchoStateNetwork, data::AbstractMatrix, ps, st::NamedTuple)
+    newst = st
+    collected = Any[]
+    for inp in eachcol(data)
+        state_t, partial_st = _partial_apply(esn, inp, ps, newst)
+        push!(collected, copy(state_t))
+        newst = merge(partial_st, (readout=newst.readout,))
+    end
+    states = eltype(data).(reduce(hcat, collected))
+    @assert !isempty(collected)
+    states_raw = reduce(hcat, collected)
+    states = eltype(data).(states_raw)
+    return states, newst
+end
+
+function addreadout!(::AbstractEchoStateNetwork, output_matrix::AbstractMatrix,
+    ps::NamedTuple, st::NamedTuple)
+    @assert hasproperty(ps, :readout)
+    new_readout = _set_readout_weight(ps.readout, output_matrix)
+    return merge(ps, (readout=new_readout,)), st
+end

src/models/esn_hybrid.jl

@@ -0,0 +1,57 @@
+@concrete struct HybridESN <: AbstractEchoStateNetwork{(:cell, :states_modifiers, :readout, :knowledge_model)}
+    cell
+    knowledge_model
+    states_modifiers
+    readout
+end
+
+function HybridESN(km,
+    km_dims::IntegerType, in_dims::IntegerType,
+    res_dims::IntegerType, out_dims::IntegerType,
+    activation=tanh;
+    state_modifiers=(),
+    readout_activation=identity,
+    include_collect::BoolType=True(),
+    kwargs...)
+
+    esn_inp_size = in_dims + km_dims
+    cell = StatefulLayer(ESNCell(esn_inp_size => res_dims, activation; kwargs...))
+    mods_tuple = state_modifiers isa Tuple || state_modifiers isa AbstractVector ?
+                 Tuple(state_modifiers) : (state_modifiers,)
+    mods = _wrap_layers(mods_tuple)
+    ro = LinearReadout(res_dims + km_dims => out_dims, readout_activation;
+        include_collect=static(include_collect))
+    km_layer = km isa WrappedFunction ? km : WrappedFunction(km)
+    return HybridESN(cell, km_layer, mods, ro)
+end
+
+function initialparameters(rng::AbstractRNG, hesn::HybridESN)
+    ps_cell = initialparameters(rng, hesn.cell)
+    ps_km = initialparameters(rng, hesn.knowledge_model)
+    ps_mods = map(l -> initialparameters(rng, l), hesn.states_modifiers) |> Tuple
+    ps_ro = initialparameters(rng, hesn.readout)
+    return (cell=ps_cell, knowledge_model=ps_km, states_modifiers=ps_mods, readout=ps_ro)
+end
+
+function initialstates(rng::AbstractRNG, hesn::HybridESN)
+    st_cell = initialstates(rng, hesn.cell)
+    st_km = initialstates(rng, hesn.knowledge_model)
+    st_mods = map(l -> initialstates(rng, l), hesn.states_modifiers) |> Tuple
+    st_ro = initialstates(rng, hesn.readout)
+    return (cell=st_cell, knowledge_model=st_km, states_modifiers=st_mods, readout=st_ro)
+end
+
+function _partial_apply(hesn::HybridESN, inp, ps, st)
+    k_t, st_km = hesn.knowledge_model(inp, ps.knowledge_model, st.knowledge_model)
+    xin = vcat(k_t, inp)
+    r, st_cell = apply(hesn.cell, xin, ps.cell, st.cell)
+    rstar, st_mods = _apply_seq(hesn.states_modifiers, r, ps.states_modifiers, st.states_modifiers)
+    feats = vcat(k_t, rstar)
+    return feats, (cell=st_cell, states_modifiers=st_mods, knowledge_model=st_km)
+end
+
+function (hesn::HybridESN)(inp, ps, st)
+    feats, new_st = _partial_apply(hesn, inp, ps, st)
+    y, st_ro = apply(hesn.readout, feats, ps.readout, st.readout)
+    return y, merge(new_st, (readout=st_ro,))
+end