test: add tests for ngrc

Author: MartinuzziFrancesco

test/esn/test_esn.jl test/models/test_esn.jltest/esn/test_esn.jl renamed to test/models/test_esn.jl

@@ -1,3 +1,9 @@
+using Test
+using Random
+using ReservoirComputing
+using Static
+using LinearAlgebra
+
 @testset "DelayESN" begin
     rng = MersenneTwister(123)
 
@@ -10,12 +16,12 @@
         desn = DelayESN(in_dims, res_dims, out_dims;
             num_delays = num_delays, stride = 1)
 
-        @test desn isa ReservoirComputer
+        @test desn isa DelayESN
 
         reservoir = desn.reservoir
         @test reservoir isa StatefulLayer
 
-        mods = desn.state_modifiers
+        mods = desn.states_modifiers
         @test mods isa Tuple
         @test !isempty(mods)
         @test first(mods) isa DelayLayer
@@ -31,28 +37,6 @@
         @test Int(ro.out_dims) == out_dims
     end
 
-    @testset "setup and forward pass shapes" begin
-        in_dims = 4
-        res_dims = 10
-        out_dims = 3
-        num_delays = 1
-
-        desn = DelayESN(in_dims, res_dims, out_dims;
-            num_delays = num_delays, stride = 1)
-
-        ps, st = setup(rng, desn)
-
-        x = rand(rng, Float32, in_dims)
-        y, st2 = desn(x, ps, st)
-        @test size(y) == (out_dims,)
-
-        X = rand(rng, Float32, in_dims, 7)
-        Y, st3 = desn(X, ps, st2)
-        @test size(Y) == (out_dims, 7)
-
-        @test propertynames(st3) == propertynames(st2)
-    end
-
     @testset "num_delays changes readout input dim" begin
         in_dims = 2
         res_dims = 6

test/esn/test_esn_deep.jl test/models/test_esn_deep.jltest/esn/test_esn_deep.jl renamed to test/models/test_esn_deep.jl

@@ -2,6 +2,7 @@ using Test
 using Random
 using ReservoirComputing
 using Static
+using LinearAlgebra
 
 const _I32 = (m, n) -> Matrix{Float32}(I, m, n)
 const _Z32 = m -> zeros(Float32, m)
@@ -35,11 +36,11 @@ end
             state_modifiers = (),
             readout_activation = identity)
         ps, st = setup(rng, hesn)
-        @test haskey(ps, :cell) && haskey(ps, :knowledge_model) &&
+        @test haskey(ps, :reservoir) && haskey(ps, :knowledge_model) &&
               haskey(ps, :states_modifiers) && haskey(ps, :readout)
-        @test size(ps.cell.input_matrix) == (res_dims, in_dims + km_dims)
+        @test size(ps.reservoir.input_matrix) == (res_dims, in_dims + km_dims)
         @test size(ps.readout.weight) == (out_dims, res_dims + km_dims)
-        @test haskey(st, :cell) && haskey(st, :knowledge_model) &&
+        @test haskey(st, :reservoir) && haskey(st, :knowledge_model) &&
               haskey(st, :states_modifiers) && haskey(st, :readout)
     end
 

test/esn/test_esn_delay.jl test/models/test_esn_delay.jltest/esn/test_esn_delay.jl renamed to test/models/test_esn_delay.jl

@@ -0,0 +1,103 @@
+using Random
+using ReservoirComputing
+using LuxCore
+using Static
+
+@testset "NGRC" begin
+    rng = MersenneTwister(1234)
+
+    const_feature = x -> Float32[1.0]
+    square_feature = x -> x .^ 2
+
+    @testset "constructor & composition" begin
+        ngrc = NGRC(3, 2; num_delays = 1, stride = 2,
+            features = (const_feature, square_feature), include_input = True(),
+            init_delay = zeros32, readout_activation = tanh)
+
+        @test ngrc isa NGRC
+        @test ngrc.reservoir isa DelayLayer
+        @test ngrc.readout isa LinearReadout
+
+        dl = ngrc.reservoir
+        @test dl.in_dims == 3
+        @test dl.num_delays == 1
+        @test dl.stride == 2
+
+        @test !isempty(ngrc.states_modifiers)
+        first_mod = getfield(ngrc.states_modifiers, 1)
+        @test first_mod isa NonlinearFeaturesLayer
+    end
+
+    @testset "initialparameters & initialstates" begin
+        ngrc = NGRC(3, 2; num_delays = 1, features = (square_feature,),
+            include_input = True())
+
+        ps = initialparameters(rng, ngrc)
+        st = initialstates(rng, ngrc)
+
+        @test hasproperty(ps, :reservoir)
+        @test hasproperty(ps, :states_modifiers)
+        @test hasproperty(ps, :readout)
+
+        @test hasproperty(st, :reservoir)
+        @test hasproperty(st, :states_modifiers)
+        @test hasproperty(st, :readout)
+
+        @test ps.readout.weight isa AbstractArray
+    end
+
+    @testset "forward pass: vector input" begin
+        ngrc = NGRC(3, 2; num_delays = 1, features = (square_feature,),
+            include_input = True())
+
+        ps, st = setup(rng, ngrc)
+
+        x = rand(Float32, 3)
+        y, st2 = ngrc(x, ps, st)
+
+        @test size(y) == (2,)
+        @test propertynames(st2) == propertynames(st)
+    end
+
+    @testset "forward pass: matrix input via collectstates" begin
+        ngrc = NGRC(3, 2; num_delays = 1, features = (square_feature,),
+            include_input = True())
+
+        ps, st = setup(rng, ngrc)
+
+        X = rand(Float32, 3, 10)
+        states, st2 = collectstates(ngrc, X, ps, st)
+
+        @test size(states, 2) == size(X, 2)
+        @test size(states, 1) == ngrc.readout.in_dims
+
+        @test propertynames(st2) == propertynames(st)
+    end
+
+    @testset "simple 1D linear system learning" begin
+        # x_{t+1} = a * x_t with a = 0.8
+        a = 0.8f0
+        T = 200
+        x = zeros(Float32, T)
+        x[1] = 1.0f0
+        for t in 1:(T - 1)
+            x[t + 1] = a * x[t]
+        end
+
+        X_in = reshape(x[1:(end - 1)], 1, :)
+        Y_out = reshape(x[2:end], 1, :)
+
+        ngrc = NGRC(1, 1; num_delays = 0, stride = 1, features = (),
+            include_input = True(), ro_dims = 1)
+
+        ps, st = setup(rng, ngrc)
+
+        ps_tr, st_tr = train!(ngrc, X_in, Y_out, ps, st;
+            train_method = StandardRidge(1e-6))
+
+        @test hasproperty(ps_tr, :readout)
+        w = ps_tr.readout.weight
+        @test size(w) == (1, 1)
+        @test isapprox(w[1, 1], a; atol = 0.05)
+    end
+end

test/esn/test_esn_hybrid.jl test/models/test_esn_hybrid.jltest/esn/test_esn_hybrid.jl renamed to test/models/test_esn_hybrid.jl

@@ -13,7 +13,13 @@ end
 end
 
 @testset "Echo State Networks" begin
-    @safetestset "ESN Initializers" include("esn/test_inits.jl")
-    @safetestset "ESN model" include("esn/test_esn.jl")
-    @safetestset "DeepESN model" include("esn/test_esn_deep.jl")
+    @safetestset "ESN Initializers" include("test_inits.jl")
+    @safetestset "ESN model" include("models/test_esn.jl")
+    @safetestset "DeepESN model" include("models/test_esn_deep.jl")
+    @safetestset "DelayESN model" include("models/test_esn_delay.jl")
+    @safetestset "HybridESN model" include("models/test_esn_hybrid.jl")
+end
+
+@testset "Next Generation Reservoir Computing" begin
+    @safetestset "NGRC model" include("models/test_ngrc.jl")
 end