feat: add svm training option

Author: MartinuzziFrancesco

ext/RCLIBSVMExt.jl

@@ -1,33 +1,71 @@
 module RCLIBSVMExt
-using ReservoirComputing
+
 using LIBSVM
+using ReservoirComputing:
+    SVMReadout, addreadout!, ReservoirChain
+import ReservoirComputing: train
 
-function ReservoirComputing.train(svr::LIBSVM.AbstractSVR,
-        states::AbstractArray, target::AbstractArray)
-    out_size = size(target, 1)
-    output_matrix = []
+function train(svr::LIBSVM.AbstractSVR,
+    states::AbstractArray, target::AbstractArray)
+    @assert size(states, 2) == size(target, 2) "states and target must share columns."
+    perm_states = permutedims(states)
+    size_target = size(target, 1)
 
-    if out_size == 1
-        output_matrix = LIBSVM.fit!(svr, states', vec(target))
+    if size_target == 1
+        vec_target = vec(target)
+        model = LIBSVM.fit!(svr, perm_states, vec_target)
+        return model
     else
-        for i in 1:out_size
-            push!(output_matrix, LIBSVM.fit!(svr, states', target[i, :]))
+        models = Vector{Any}(undef, size_target)
+        for (idx, row_target) in enumerate(eachrow(target))
+            models[idx] = LIBSVM.fit!(svr, perm_states, row_target)
         end
+        return models
     end
-
-    return OutputLayer(svr, output_matrix, out_size, target[:, end])
 end
 
-function ReservoirComputing.get_prediction(training_method::LIBSVM.AbstractSVR,
-        output_layer::AbstractArray, x::AbstractArray)
-    out = zeros(output_layer.out_size)
+_has_models(ps) = (ps isa NamedTuple) && (:models in propertynames(ps))
+
+function (svmro::SVMReadout)(inp::AbstractArray, ps, st::NamedTuple)
+    if !_has_models(ps)
+        return inp, st
+    end
+    models = getfield(ps, :models)
+
+    vec_like = false
+    if ndims(inp) == 1
+        reshaped_inp = reshape(inp, 1, :)
+        num_imp = 1
+        vec_like = true
+    elseif ndims(inp) == 2
+        if size(inp, 2) == 1
+            reshaped_inp = reshape(vec(inp), 1, :)
+            num_inp = 1
+            vec_like = true
+        else
+            reshaped_inp = permutedims(inp)
+            num_imp = size(reshaped_inp, 1)
+        end
+    else
+        throw(ArgumentError("SVMReadout expects 1D or 2D input; got size $(size(inp))"))
+    end
 
-    for i in 1:(output_layer.out_size)
-        x_new = reshape(x, 1, length(x))
-        out[i] = LIBSVM.predict(output_layer.output_matrix[i], x_new)[1]
+    if models isa AbstractVector
+        out_data = Array{float(eltype(reshaped_inp))}(undef, svmro.out_dims, num_imp)
+        @inbounds for i in 1:svmro.out_dims
+            single_out = LIBSVM.predict(models[i], reshaped_inp)
+            out_data[i, :] = single_out
+        end
+    else
+        single_out = LIBSVM.predict(models, reshaped_inp)
+        out_data = reshape(single_out, 1, :)
     end
 
-    return out
+    if vec_like
+        return vec(out_data), st
+    else
+        return out_data, st
+    end
 end
 
-end #module
+end # module

src/ReservoirComputing.jl

@@ -17,35 +17,33 @@ using WeightInitializers: DeviceAgnostic, PartialFunction, Utils
 @reexport using WeightInitializers
 @reexport using LuxCore: setup, apply
 
-
 const BoolType = Union{StaticBool,Bool,Val{true},Val{false}}
 const InputType = Tuple{<:AbstractArray,Tuple{<:AbstractArray}}
 const IntegerType = Union{Integer,StaticInteger}
 
-@compat(public, (create_states))
+#@compat(public, (create_states)) #do I need to add intialstates/parameters in compat?
 
 #layers
 include("layers/basic.jl")
 include("layers/lux_layers.jl")
 include("layers/esn_cell.jl")
+include("layers/svmreadout.jl")
 #general
 include("states.jl")
 include("predict.jl")
 include("train.jl")
-#esn
+#initializers
 include("inits/inits_components.jl")
 include("inits/esn_inits.jl")
+#full models
 include("models/esn.jl")
 include("models/deepesn.jl")
 include("models/hybridesn.jl")
 #extensions
 include("extensions/reca.jl")
 
-
-
-
 export ESNCell, StatefulLayer, Readout, ReservoirChain, Collect, collectstates, train!, predict
-
+export SVMReadout
 export Pad, Extend, NLAT1, NLAT2, NLAT3, PartialSquare, ExtendedSquare
 export StandardRidge
 export chebyshev_mapping, informed_init, logistic_mapping, minimal_init,

src/layers/basic.jl

@@ -195,6 +195,9 @@ function collectstates(rc::AbstractLuxLayer, data::AbstractArray, ps, st::NamedT
         inp_tmp = inp
         state_vec = nothing
         for (name, layer) in pairs(rc.layers)
+            if layer isa AbstractReservoirTrainableLayer
+                break
+            end
             inp_tmp, st_i = layer(inp_tmp, ps[name], newst[name])
             newst = merge(newst, (; name => st_i))
             if layer isa AbstractReservoirCollectionLayer

src/layers/svmreadout.jl

@@ -0,0 +1,81 @@
+@concrete struct SVMReadout <: AbstractReservoirTrainableLayer
+    in_dims <: IntegerType
+    out_dims <: IntegerType
+    include_collect <: StaticBool
+end
+
+function Base.show(io::IO, ro::SVMReadout)
+    print(io, "SVMReadout($(ro.in_dims) => $(ro.out_dims)")
+    ic = known(getproperty(ro, Val(:include_collect)))
+    ic === true && print(io, ", include_collect=true")
+    return print(io, ")")
+end
+
+SVMReadout(mapping::Pair{<:IntegerType,<:IntegerType}; kwargs...) =
+    SVMReadout(first(mapping), last(mapping); kwargs...)
+
+SVMReadout(in_dims::IntegerType, out_dims::IntegerType;
+    include_collect::BoolType=True()) =
+    SVMReadout(in_dims, out_dims, static(include_collect))
+
+initialparameters(::AbstractRNG, ::SVMReadout) = NamedTuple()
+parameterlength(::SVMReadout) = 0
+statelength(::SVMReadout) = 0
+outputsize(ro::SVMReadout, _, ::AbstractRNG) = (ro.out_dims,)
+
+# NOTE: forward for SVMReadout will be defined in the LIBSVM extension,
+# because it calls LIBSVM.predict.
+
+_svmreadout_include_collect(ro::SVMReadout) = begin
+    ic = known(getproperty(ro, Val(:include_collect)))
+    ic === nothing ? false : ic
+end
+
+function wrap_functions_in_chain_call(ro::SVMReadout)
+    return _svmreadout_include_collect(ro) ? (Collect(), ro) : ro
+end
+
+
+_quote_keys(t) = Expr(:tuple, (QuoteNode(s) for s in t)...)
+
+function _setmodels_rt(p::NamedTuple{K}, M) where {K}
+    keys = K
+    Kq = _quote_keys(keys)
+    idx = findfirst(==(Symbol(:models)), keys)
+
+    terms = Any[]
+    for i in 1:length(keys)
+        push!(terms, (idx === i) ? :(M) : :(getfield(p, $i)))
+    end
+
+    if idx === nothing
+        newK = _quote_keys((keys..., :models))
+        return :(NamedTuple{$newK}(($(terms...), M)))
+    else
+        return :(NamedTuple{$Kq}(($(terms...),)))
+    end
+end
+
+@generated function _addsvm(layers::NamedTuple{K}, ps::NamedTuple{K}, M) 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 SVMReadout)
+                 ? _setmodels_rt(getfield(ps, 1), M)
+                 : getfield(ps, 1))
+
+    tail_call = :(_addsvm(NamedTuple{$tailKq}(Base.tail(layers)),
+        NamedTuple{$tailKq}(Base.tail(ps)), M))
+
+    return :(NamedTuple{$Kq}(($head_val, Base.values($tail_call)...)))
+end
+
+function addreadout!(rc::ReservoirChain, models, ps::NamedTuple, st::NamedTuple)
+    @assert propertynames(rc.layers) == propertynames(ps)
+    new_ps = _addsvm(rc.layers, ps, models)
+    return new_ps, st
+end

src/predict.jl

@@ -62,7 +62,7 @@ function predict(rc::AbstractLuxLayer, data::AbstractMatrix, ps, st)
     Y = similar(y1, size(y1, 1), T)
     Y[:, 1] .= y1
 
-    @inbounds @views for t in 2:T
+    for t in 2:T
         yt, st = apply(rc, data[:, t], ps, st)
         Y[:, t] .= yt
     end

src/train.jl

@@ -94,7 +94,7 @@ function train(sr::StandardRidge, states::AbstractArray, target_data::AbstractAr
     return output_layer
 end
 
-_quote_keys(t) = Expr(:tuple, (QuoteNode(s) for s in t)...)
+#_quote_keys(t) = Expr(:tuple, (QuoteNode(s) for s in t)...)
 
 @generated function _setweight_rt(p::NamedTuple{K}, W) where {K}
     keys = K