formatting, docstrings

Author: MartinuzziFrancesco

ext/RCLIBSVMExt.jl

@@ -2,7 +2,8 @@ module RCLIBSVMExt
 using ReservoirComputing
 using LIBSVM
 
-function ReservoirComputing.train(svr::LIBSVM.AbstractSVR, states, target)
+function ReservoirComputing.train(svr::LIBSVM.AbstractSVR,
+        states::AbstractArray, target::AbstractArray)
     out_size = size(target, 1)
     output_matrix = []
 
@@ -17,8 +18,8 @@ function ReservoirComputing.train(svr::LIBSVM.AbstractSVR, states, target)
     return OutputLayer(svr, output_matrix, out_size, target[:, end])
 end
 
-function ReservoirComputing.get_prediction(
-        training_method::LIBSVM.AbstractSVR, output_layer, x)
+function ReservoirComputing.get_prediction(training_method::LIBSVM.AbstractSVR,
+        output_layer::AbstractArray, x::AbstractArray)
     out = zeros(output_layer.out_size)
 
     for i in 1:(output_layer.out_size)

ext/RCMLJLinearModelsExt.jl

@@ -3,8 +3,7 @@ using ReservoirComputing
 using MLJLinearModels
 
 function ReservoirComputing.train(regressor::MLJLinearModels.GeneralizedLinearRegression,
-        states::AbstractArray{T},
-        target::AbstractArray{T};
+        states::AbstractArray{T}, target::AbstractArray{T};
         kwargs...) where {T <: Number}
     out_size = size(target, 1)
     output_layer = similar(target, size(target, 1), size(states, 1))

src/esn/deepesn.jl

@@ -11,16 +11,16 @@ struct DeepESN{I, S, N, T, O, M, B, ST, W, IS} <: AbstractEchoStateNetwork
     states::IS
 end
 
+const AbstractDriver = Union{AbstractReservoirDriver, GRU}
+
 """
     DeepESN(train_data, in_size, res_size; kwargs...)
 
 Constructs a Deep Echo State Network (ESN) model for
 processing sequential data through a layered architecture of reservoirs.
 This constructor allows for the creation of a deep learning model that
 benefits from the dynamic memory and temporal processing capabilities of ESNs,
-enhanced by the depth provided by multiple reservoir layers. It's particularly
-suited for complex sequential tasks where depth can help capture hierarchical
-temporal features.
+enhanced by the depth provided by multiple reservoir layers.
 
 # Parameters
 
@@ -62,19 +62,12 @@ train_data = rand(Float32, 3, 100)
 deepESN = DeepESN(train_data, 3, 100; depth=3, washout=100)
 ```
 """
-function DeepESN(train_data,
-        in_size::Int,
-        res_size::Int;
-        depth::Int=2,
-        input_layer=fill(scaled_rand, depth),
-        bias=fill(zeros32, depth),
-        reservoir=fill(rand_sparse, depth),
-        reservoir_driver=RNN(),
-        nla_type=NLADefault(),
-        states_type=StandardStates(),
-        washout::Int=0,
-        rng=Utils.default_rng(),
-        matrix_type=typeof(train_data))
+function DeepESN(train_data::AbstractArray, in_size::Int, res_size::Int; depth::Int=2,
+        input_layer=fill(scaled_rand, depth), bias=fill(zeros32, depth),
+        reservoir=fill(rand_sparse, depth), reservoir_driver::AbstractDriver=RNN(),
+        nla_type::NonLinearAlgorithm=NLADefault(),
+        states_type::AbstractStates=StandardStates(), washout::Int=0,
+        rng::AbstractRNG=Utils.default_rng(), matrix_type=typeof(train_data))
     if states_type isa AbstractPaddedStates
         in_size = size(train_data, 1) + 1
         train_data = vcat(adapt(matrix_type, ones(1, size(train_data, 2))),

src/esn/esn.jl

@@ -12,6 +12,8 @@ struct ESN{I, S, N, T, O, M, B, ST, W, IS} <: AbstractEchoStateNetwork
     states::IS
 end
 
+const AbstractDriver = Union{AbstractReservoirDriver, GRU}
+
 """
     ESN(train_data; kwargs...) -> ESN
 
@@ -52,17 +54,12 @@ julia> esn = ESN(train_data, 10, 300; washout=10)
 ESN(10 => 300)
 ```
 """
-function ESN(train_data,
-        in_size::Int,
-        res_size::Int;
-        input_layer=scaled_rand,
-        reservoir=rand_sparse,
-        bias=zeros32,
-        reservoir_driver=RNN(),
-        nla_type=NLADefault(),
-        states_type=StandardStates(),
-        washout=0,
-        rng=Utils.default_rng(),
+function ESN(train_data::AbstractArray, in_size::Int, res_size::Int;
+        input_layer=scaled_rand, reservoir=rand_sparse, bias=zeros32,
+        reservoir_driver::AbstractDriver=RNN(),
+        nla_type::NonLinearAlgorithm=NLADefault(),
+        states_type::AbstractStates=StandardStates(),
+        washout::Int=0, rng::AbstractRNG=Utils.default_rng(),
         matrix_type=typeof(train_data))
     if states_type isa AbstractPaddedStates
         in_size = size(train_data, 1) + 1
@@ -85,11 +82,9 @@ function ESN(train_data,
 end
 
 function (esn::AbstractEchoStateNetwork)(prediction::AbstractPrediction,
-        output_layer::AbstractOutputLayer;
-        last_state=esn.states[:, [end]],
+        output_layer::AbstractOutputLayer; last_state=esn.states[:, [end]],
         kwargs...)
     pred_len = prediction.prediction_len
-
     return obtain_esn_prediction(esn, prediction, last_state, output_layer;
         kwargs...)
 end
@@ -133,12 +128,9 @@ julia> output_layer = train(esn, rand(Float32, 3, 90))
 OutputLayer successfully trained with output size: 3
 ```
 """
-function train(esn::AbstractEchoStateNetwork,
-        target_data,
-        training_method=StandardRidge();
-        kwargs...)
+function train(esn::AbstractEchoStateNetwork, target_data::AbstractArray,
+        training_method=StandardRidge(); kwargs...)
     states_new = esn.states_type(esn.nla_type, esn.states, esn.train_data[:, 1:end])
-
     return train(training_method, states_new, target_data; kwargs...)
 end
 

src/esn/esn_reservoir_drivers.jl

@@ -22,14 +22,10 @@ specified reservoir driver.
     update.
 """
 function create_states(reservoir_driver::AbstractReservoirDriver,
-        train_data,
-        washout,
-        reservoir_matrix,
-        input_matrix,
-        bias_vector)
+        train_data::AbstractArray, washout::Int, reservoir_matrix::AbstractMatrix,
+        input_matrix::AbstractMatrix, bias_vector::AbstractArray)
     train_len = size(train_data, 2) - washout
     res_size = size(reservoir_matrix, 1)
-
     states = adapt(typeof(train_data), zeros(res_size, train_len))
     tmp_array = allocate_tmp(reservoir_driver, typeof(train_data), res_size)
     _state = adapt(typeof(train_data), zeros(res_size, 1))
@@ -51,14 +47,10 @@ function create_states(reservoir_driver::AbstractReservoirDriver,
 end
 
 function create_states(reservoir_driver::AbstractReservoirDriver,
-        train_data,
-        washout,
-        reservoir_matrix::Vector,
-        input_matrix,
-        bias_vector)
+        train_data::AbstractArray, washout::Int, reservoir_matrix::Vector,
+        input_matrix::AbstractArray, bias_vector::AbstractArray)
     train_len = size(train_data, 2) - washout
     res_size = sum([size(reservoir_matrix[i], 1) for i in 1:length(reservoir_matrix)])
-
     states = adapt(typeof(train_data), zeros(res_size, train_len))
     tmp_array = allocate_tmp(reservoir_driver, typeof(train_data), res_size)
     _state = adapt(typeof(train_data), zeros(res_size))

src/esn/hybridesn.jl

@@ -12,6 +12,8 @@ struct HybridESN{I, S, V, N, T, O, M, B, ST, W, IS} <: AbstractEchoStateNetwork
     states::IS
 end
 
+const AbstractDriver = Union{AbstractReservoirDriver, GRU}
+
 struct KnowledgeModel{T, K, O, I, S, D}
     prior_model::T
     u0::K
@@ -91,19 +93,12 @@ traditional Echo State Networks with a predefined knowledge model [^Pathak2018].
     "Hybrid Forecasting of Chaotic Processes:
     Using Machine Learning in Conjunction with a Knowledge-Based Model" (2018).
 """
-function HybridESN(model,
-        train_data,
-        in_size::Int,
-        res_size::Int;
-        input_layer=scaled_rand,
-        reservoir=rand_sparse,
-        bias=zeros32,
-        reservoir_driver=RNN(),
-        nla_type=NLADefault(),
-        states_type=StandardStates(),
-        washout=0,
-        rng=Utils.default_rng(),
-        T=Float32,
+function HybridESN(model::KnowledgeModel, train_data::AbstractArray,
+        in_size::Int, res_size::Int; input_layer=scaled_rand, reservoir=rand_sparse,
+        bias=zeros32, reservoir_driver::AbstractDriver=RNN(),
+        nla_type::NonLinearAlgorithm=NLADefault(),
+        states_type::AbstractStates=StandardStates(), washout::Int=0,
+        rng::AbstractRNG=Utils.default_rng(), T=Float32,
         matrix_type=typeof(train_data))
     train_data = vcat(train_data, model.model_data[:, 1:(end - 1)])
 
@@ -130,8 +125,7 @@ function HybridESN(model,
 end
 
 function (hesn::HybridESN)(prediction::AbstractPrediction,
-        output_layer::AbstractOutputLayer;
-        last_state=hesn.states[:, [end]],
+        output_layer::AbstractOutputLayer; last_state::AbstractArray=hesn.states[:, [end]],
         kwargs...)
     km = hesn.model
     pred_len = prediction.prediction_len
@@ -148,10 +142,8 @@ function (hesn::HybridESN)(prediction::AbstractPrediction,
         kwargs...)
 end
 
-function train(hesn::HybridESN,
-        target_data,
-        training_method=StandardRidge();
-        kwargs...)
+function train(hesn::HybridESN, target_data::AbstractArray,
+        training_method=StandardRidge(); kwargs...)
     states = vcat(hesn.states, hesn.model.model_data[:, 2:end])
     states_new = hesn.states_type(hesn.nla_type, states, hesn.train_data[:, 1:end])
 

src/predict.jl

@@ -61,16 +61,13 @@ 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)
+function Predictive(prediction_data::AbstractArray)
     prediction_len = size(prediction_data, 2)
     Predictive(prediction_data, prediction_len)
 end
 
-function obtain_prediction(rc::AbstractReservoirComputer,
-        prediction::Generative,
-        x,
-        output_layer,
-        args...;
+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
@@ -88,12 +85,8 @@ function obtain_prediction(rc::AbstractReservoirComputer,
     return output
 end
 
-function obtain_prediction(rc::AbstractReservoirComputer,
-        prediction::Predictive,
-        x,
-        output_layer,
-        args...;
-        kwargs...)
+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
@@ -110,7 +103,7 @@ function obtain_prediction(rc::AbstractReservoirComputer,
 end
 
 #linear models
-function get_prediction(training_method, output_layer, x)
+function get_prediction(training_method, output_layer::AbstractOutputLayer, x)
     return output_layer.output_matrix * x
 end
 

src/train/linear_regression.jl

@@ -1,3 +1,62 @@
+@doc raw"""
+
+    StandardRidge([Type], [reg])
+
+Returns a training method for `train` based on ridge regression.
+The equations for ridge regression are as follows:
+
+```math
+\mathbf{w} = (\mathbf{X}^\top \mathbf{X} + 
+\lambda \mathbf{I})^{-1} \mathbf{X}^\top \mathbf{y}
+```
+
+# Arguments
+ - `Type`: type of the regularization argument. Default is inferred internally,
+   there's usually no need to tweak this
+ - `reg`: regularization coefficient. Default is set to 0.0 (linear regression).
+
+# Examples
+```jldoctest
+julia> ridge_reg = StandardRidge()
+StandardRidge(0.0)
+
+julia> ol = train(ridge_reg, rand(Float32, 10, 10), rand(Float32, 10, 10))
+OutputLayer successfully trained with output size: 10
+
+julia> ol.output_matrix #visualize output matrix
+10×10 Matrix{Float32}:
+  0.456574   -0.0407612   0.121963   …   0.859327   -0.127494    0.0572494
+  0.133216   -0.0337922   0.0185378      0.24077     0.0297829   0.31512
+  0.379672   -1.24541    -0.444314       1.02269    -0.0446086   0.482282
+  1.18455    -0.517971   -0.133498       0.84473     0.31575     0.205857
+ -0.119345    0.563294    0.747992       0.0102919   1.509      -0.328005
+ -0.0716812   0.0976365   0.628654   …  -0.516041    2.4309     -0.113402
+  0.0153872  -0.52334     0.0526867      0.729326    2.98958     1.32703
+  0.154027    0.6013      1.05548       -0.0840203   0.991182   -0.328555
+  1.11007    -0.0371736  -0.0529418      0.186796   -1.21815     0.204838
+  0.282996   -0.263799    0.132079       0.875417    0.497951    0.273423
+
+julia> ridge_reg = StandardRidge(0.001) #passing a value 
+StandardRidge(0.001)
+
+julia> ol = train(ridge_reg, rand(Float16, 10, 10), rand(Float16, 10, 10))
+OutputLayer successfully trained with output size: 10
+
+julia> ol.output_matrix
+10×10 Matrix{Float16}:
+ -1.251     3.074   -1.566     -0.10297  …   0.3823   1.341    -1.77    -0.445
+  0.11017  -2.027    0.8975     0.872       -0.643    0.02615   1.083    0.615
+  0.2634    3.514   -1.168     -1.532        1.486    0.1255   -1.795   -0.06555
+  0.964     0.9463  -0.006855  -0.519        0.0743  -0.181    -0.433    0.06793
+ -0.389     1.887   -0.702     -0.8906       0.221    1.303    -1.318    0.2634
+ -0.1337   -0.4453  -0.06866    0.557    …  -0.322    0.247     0.2554   0.5933
+ -0.6724    0.906   -0.547      0.697       -0.2664   0.809    -0.6836   0.2358
+  0.8843   -3.664    1.615      1.417       -0.6094  -0.59      1.975    0.4785
+  1.266    -0.933    0.0664    -0.4497      -0.0759  -0.03897   1.117    0.3152
+  0.6353    1.327   -0.6978    -1.053        0.8037   0.6577   -0.7246   0.07336
+
+```
+"""
 struct StandardRidge
     reg::Number
 end
@@ -10,13 +69,23 @@ function StandardRidge()
     return StandardRidge(0.0)
 end
 
-function train(sr::StandardRidge,
-        states,
-        target_data)
+function train(sr::StandardRidge, states::AbstractArray, target_data::AbstractArray)
     #A = states * states' + sr.reg * I
     #b = states * target_data
     #output_layer = (A \ b)'
-    output_layer = Matrix(((states * states' + sr.reg * I) \
+
+    if size(states, 2) != size(target_data, 2)
+        throw(DimensionMismatch("\n" *
+            "\n" *
+            "  - Number of columns in `states`: $(size(states, 2))\n" *
+            "  - Number of columns in `target_data`: $(size(target_data, 2))\n" *
+            "The dimensions of `states` and `target_data` must align for training." *
+            "\n"
+        ))
+    end
+
+    T = eltype(states)
+    output_layer = Matrix(((states * states' + T(sr.reg) * I) \
                            (states * target_data'))')
     return OutputLayer(sr, output_layer, size(target_data, 1), target_data[:, end])
 end