Merge pull request #289 from SciML/fm/blockdiagonal

feat: add block_diagonal initializer

Author: MartinuzziFrancesco

docs/src/api/inits.md

@@ -3,46 +3,47 @@
 ## Input layers
 
 ```@docs
-    scaled_rand
-    weighted_init
-    minimal_init
-    weighted_minimal
     chebyshev_mapping
+    informed_init
     logistic_mapping
+    minimal_init
     modified_lm
-    informed_init
+    scaled_rand
+    weighted_init
+    weighted_minimal
 ```
 
 ## Reservoirs
 
 ```@docs
-    rand_sparse
-    pseudo_svd
+    block_diagonal
     chaotic_init
-    low_connectivity
+    cycle_jumps
     delay_line
     delay_line_backward
-    simple_cycle
-    cycle_jumps
     double_cycle
-    true_double_cycle
+    forward_connection
+    low_connectivity
+    pseudo_svd
+    rand_sparse
     selfloop_cycle
-    selfloop_feedback_cycle
     selfloop_delayline_backward
+    selfloop_feedback_cycle
     selfloop_forward_connection
-    forward_connection
+    simple_cycle
+    true_double_cycle
 ```
 
 ## Building functions
 
 ```@docs
-    scale_radius!
-    delay_line!
+    add_jumps!
     backward_connection!
-    simple_cycle!
+    delay_line!
     reverse_simple_cycle!
+    scale_radius!
     self_loop!
-    add_jumps!
+    simple_cycle!
 ```
 
 ## References

docs/src/refs.bib

@@ -326,4 +326,18 @@ @article{Herteux2020
   author = {Herteux,  Joschka and R\"{a}th,  Christoph},
   year = {2020},
   month = dec 
+}
+
+@article{Ma2023,
+  title = {Efficient forecasting of chaotic systems with block-diagonal and binary reservoir computing},
+  volume = {33},
+  ISSN = {1089-7682},
+  url = {http://dx.doi.org/10.1063/5.0151290},
+  DOI = {10.1063/5.0151290},
+  number = {6},
+  journal = {Chaos: An Interdisciplinary Journal of Nonlinear Science},
+  publisher = {AIP Publishing},
+  author = {Ma,  Haochun and Prosperino,  Davide and Haluszczynski,  Alexander and R\"{a}th,  Christoph},
+  year = {2023},
+  month = jun 
 }

src/ReservoirComputing.jl

@@ -37,15 +37,14 @@ include("reca/reca_input_encodings.jl")
 export NLADefault, NLAT1, NLAT2, NLAT3, PartialSquare, ExtendedSquare
 export StandardStates, ExtendedStates, PaddedStates, PaddedExtendedStates
 export StandardRidge
-export scaled_rand, weighted_init, informed_init, minimal_init, chebyshev_mapping,
-       logistic_mapping, modified_lm, weighted_minimal
-export rand_sparse, delay_line, delay_line_backward, cycle_jumps,
-       simple_cycle, pseudo_svd, chaotic_init, low_connectivity, double_cycle,
-       selfloop_cycle, selfloop_feedback_cycle, selfloop_delayline_backward,
-       selfloop_forward_connection, forward_connection, true_double_cycle
-export scale_radius!, delay_line!, backward_connection!, simple_cycle!,
-       reverse_simple_cycle!,
-       add_jumps!, self_loop!
+export chebyshev_mapping, informed_init, logistic_mapping, minimal_init,
+       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
+export add_jumps!, backward_connection!, delay_line!, reverse_simple_cycle!,
+       scale_radius!, self_loop!, simple_cycle!
 export RNN, MRNN, GRU, GRUParams, FullyGated, Minimal
 export train
 export ESN, HybridESN, KnowledgeModel, DeepESN

src/esn/esn_inits.jl

@@ -1881,14 +1881,106 @@ function forward_connection(rng::AbstractRNG, ::Type{T}, dims::Integer...;
     return return_init_as(Val(return_sparse), reservoir_matrix)
 end
 
+@doc raw"""
+    block_diagonal([rng], [T], dims...;
+        weight=1, block_size=1,
+        return_sparse=false)
+
+Creates a block‐diagonal matrix consisting of square blocks of size
+`block_size` along the main diagonal [Ma2023](@cite).  
+Each block may be filled with  
+  - a single scalar
+  - a vector of per‐block weights (length = number of blocks)
+
+# Equations
+
+```math
+W_{i,j} =
+\begin{cases}
+    w_b, & \text{if }\left\lfloor\frac{i-1}{s}\right\rfloor = \left\lfloor\frac{j-1}{s}\right\rfloor = b,\; 
+           s = \text{block\_size},\; b=0,\dots,nb-1, \\
+    0,   & \text{otherwise,}
+\end{cases}
+```
+
+# Arguments
+
+  - `rng`: Random number generator. Default is `Utils.default_rng()`.  
+  - `T`: Element type of the matrix. Default is `Float32`.  
+  - `dims`: Dimensions of the output matrix (must be two-dimensional).
+
+# Keyword arguments
+
+  - `weight`:  
+    - scalar: every block is filled with that value
+    - vector: length = number of blocks, one constant per block
+    Default is `1.0`.
+  - `block_size`: Size\(s\) of each square block on the diagonal. Default is `1.0`.
+  - `return_sparse`: If `true`, returns the matrix as sparse.
+    SparseArrays.jl must be lodead.
+    Default is `false`.
+
+# Examples
+
+```jldoctest
+# 4×4 with two 2×2 blocks of 1.0
+julia> W1 = block_diagonal(4, 4; block_size=2)
+4×4 Matrix{Float32}:
+ 1.0  1.0  0.0  0.0
+ 1.0  1.0  0.0  0.0
+ 0.0  0.0  1.0  1.0
+ 0.0  0.0  1.0  1.0
+
+# per-block weights [0.5, 2.0]
+julia> W2 = block_diagonal(4, 4; block_size=2, weight=[0.5, 2.0])
+4×4 Matrix{Float32}:
+ 0.5  0.5  0.0  0.0
+ 0.5  0.5  0.0  0.0
+ 0.0  0.0  2.0  2.0
+ 0.0  0.0  2.0  2.0
+```
+"""
+function block_diagonal(rng::AbstractRNG, ::Type{T}, dims::Integer...;
+        weight::Union{Number, AbstractVector} = T(1),
+        block_size::Integer = 1,
+        return_sparse::Bool = false) where {T <: Number}
+    throw_sparse_error(return_sparse)
+    check_res_size(dims...)
+    n_rows, n_cols = dims
+    total = min(n_rows, n_cols)
+    num_blocks = fld(total, block_size)
+    remainder = total - num_blocks * block_size
+    if remainder != 0
+        @warn "\n
+        With block_size=$block_size on a $n_rows×$n_cols matrix,
+        only $num_blocks block(s) of size $block_size fit,
+        leaving $remainder row(s)/column(s) unused.
+        \n"
+    end
+    weights = isa(weight, AbstractVector) ? T.(weight) : fill(T(weight), num_blocks)
+    @assert length(weights)==num_blocks "
+      weight vector must have length = number of blocks
+  "
+    reservoir_matrix = DeviceAgnostic.zeros(rng, T, n_rows, n_cols)
+    for block in 1:num_blocks
+        row_start = (block - 1) * block_size + 1
+        row_end = row_start + block_size - 1
+        col_start = (block - 1) * block_size + 1
+        col_end = col_start + block_size - 1
+        @inbounds reservoir_matrix[row_start:row_end, col_start:col_end] .= weights[block]
+    end
+
+    return return_init_as(Val(return_sparse), reservoir_matrix)
+end
+
 ### fallbacks
 #fallbacks for initializers #eventually to remove once migrated to WeightInitializers.jl
 for initializer in (:rand_sparse, :delay_line, :delay_line_backward, :cycle_jumps,
     :simple_cycle, :pseudo_svd, :chaotic_init, :scaled_rand, :weighted_init,
     :weighted_minimal, :informed_init, :minimal_init, :chebyshev_mapping,
     :logistic_mapping, :modified_lm, :low_connectivity, :double_cycle, :selfloop_cycle,
     :selfloop_feedback_cycle, :selfloop_delayline_backward, :selfloop_forward_connection,
-    :forward_connection, :true_double_cycle)
+    :forward_connection, :true_double_cycle, :block_diagonal)
     @eval begin
         function ($initializer)(dims::Integer...; kwargs...)
             return $initializer(Utils.default_rng(), Float32, dims...; kwargs...)

test/esn/test_inits.jl

@@ -19,31 +19,32 @@ end
 
 ft = [Float16, Float32, Float64]
 reservoir_inits = [
-    rand_sparse,
+    block_diagonal,
+    chaotic_init,
+    cycle_jumps,
     delay_line,
     delay_line_backward,
-    cycle_jumps,
-    simple_cycle,
-    pseudo_svd,
-    chaotic_init,
-    low_connectivity,
     double_cycle,
+    forward_connection,
+    low_connectivity,
+    pseudo_svd,
+    rand_sparse,
     selfloop_cycle,
-    selfloop_feedback_cycle,
     selfloop_delayline_backward,
+    selfloop_feedback_cycle,
     selfloop_forward_connection,
-    forward_connection,
+    simple_cycle,
     true_double_cycle
 ]
 input_inits = [
-    scaled_rand,
-    weighted_init,
-    weighted_minimal,
-    minimal_init,
-    minimal_init(; sampling_type = :irrational_sample!),
     chebyshev_mapping,
     logistic_mapping,
-    modified_lm(; factor = 4)
+    minimal_init,
+    minimal_init(; sampling_type = :irrational_sample!),
+    modified_lm(; factor = 4),
+    scaled_rand,
+    weighted_init,
+    weighted_minimal
 ]
 
 @testset "Reservoir Initializers" begin