fix: make docs and tests pass

Author: MartinuzziFrancesco

docs/pages.jl

@@ -5,14 +5,15 @@ pages = [
         "Chaos forecasting with an ESN" => "tutorials/lorenz_basic.md",
         #"Using Different Training Methods" => "esn_tutorials/different_training.md",
         "Deep Echo State Networks" => "tutorials/deep_esn.md",
-        "Hybrid Echo State Networks" => "tutorials/hybrid.md",
+        #"Hybrid Echo State Networks" => "tutorials/hybrid.md",
         "Reservoir Computing with Cellular Automata" => "tutorials/reca.md"],
     "API Documentation" => Any[
         "Layers" => "api/layers.md",
         "Models" => "api/models.md",
-        "States" => "api/states.md",
+        "Utilities" => "api/utils.md",
         "Train" => "api/train.md",
         "Predict" => "api/predict.md",
+        "States" => "api/states.md",
         "Initializers" => "api/inits.md"],
     "References" => "references.md"
 ]

docs/src/api/layers.md

@@ -2,21 +2,27 @@
 
 ## Base Layers
 
-```@doc
+```@docs
     ReservoirChain
     Collect
     StatefulLayer
-    LinearReadout
 ```
 
-## External Layers
+## Readout Layers
 
 ```@docs
+    LinearReadout
     SVMReadout
 ```
 
 ## Echo State Networks
 
-```@doc
+```@docs
     ESNCell
 ```
+
+## REservoir computing with cellualr automata
+
+```@docs
+    RECACell
+```

docs/src/api/models.md

@@ -25,5 +25,3 @@ The input encodings are the equivalent of the input matrices of the ESNs. These
 ```@docs
     RandomMapping
 ```
-
-The training and prediction follow the same workflow as the ESN. It is important to note that currently we were unable to find any papers using these models with a `Generative` approach for the prediction, so full support is given only to the `Predictive` method.

docs/src/api/utils.md

@@ -0,0 +1,5 @@
+# Utilities
+
+```@docs
+    collectstates
+```

docs/src/tutorials/lorenz_basic.md

@@ -133,8 +133,7 @@ ReservoirComputing.jl provides
 additional utilities functions for autoregressive forecasting:
 
 ```@example lorenz
-pred_length
-output, st = predict(esn, predict_len, ps, st; initialdata=test[:, 1])
+output, st = predict(esn, predict_len, ps, st; initialdata=test_data[:, 1])
 ```
 
 To inspect the results, they can easily be plotted using an external library.

docs/src/tutorials/scratch.md

@@ -9,7 +9,7 @@ utilities)
 ## Using provided layers: ReservoirChain, ESNCell, and LinearReadout
 
 The library provides a [`ReservoirChain`](@ref), which is virtually
-equivalent to Lux's [`Chain`](@extref). Passing layers, or functions,
+equivalent to Lux's `Chain`. Passing layers, or functions,
 to the chain will concatenate them, and will allow the flow of the input
 data through the model.
 
@@ -48,7 +48,7 @@ the weights of the input layer:
 
 ```@example scratch
 using Random
-Random.seed(43)
+Random.seed!(43)
 
 rng = MersenneTwister(17)
 ps_s, st_s = setup(rng, esn_scratch)

src/extensions/reca.jl

@@ -39,7 +39,7 @@ cell = RECACell(DCA(90), enc)
 
 rc = ReservoirChain(
     StatefulLayer(cell),
-    Readout(enc.states_size => in_dims; include_collect = true)
+    LinearReadout(enc.states_size => in_dims; include_collect = true)
 )
 ```
 
@@ -148,7 +148,7 @@ Construct a cellular–automata reservoir model.
 At each time step the input vector is randomly embedded into a Cellular
 Automaton (CA) lattice, the CA is evolved for `generations` steps, and the
 flattened evolution (excluding the initial row) is used as the reservoir state.
-A linear [`Readout`](@ref) maps these features to `out_dims`.
+A linear [`LinearReadout`](@ref) maps these features to `out_dims`.
 
 !!! note
     This constructor is only available when the `CellularAutomata.jl` package is

src/layers/basic.jl

@@ -11,7 +11,7 @@ abstract type AbstractReservoirTrainableLayer <: AbstractLuxLayer end
 Linear readout layer with optional bias and elementwise activation. Intended as
 the final, trainable mapping from collected features (e.g., reservoir state) to
 outputs. When `include_collect=true`, training will collect features immediately
-before this layer (logically inserting a [`Collect()`](@ref) right before it).
+before this layer (logically inserting a [`Collect`](@ref) right before it).
 
 ## Equation
 
@@ -29,7 +29,7 @@ before this layer (logically inserting a [`Collect()`](@ref) right before it).
 
 - `use_bias`: Include an additive bias vector `b`. Default: `false`.
 - `include_collect`: If `true` (default), training collects features immediately
-  before this layer (as if a [`Collect()`](@ref) were inserted right before it).
+  before this layer (as if a [`Collect`](@ref) were inserted right before it).
 
 ## Parameters
 
@@ -42,9 +42,11 @@ before this layer (logically inserting a [`Collect()`](@ref) right before it).
 
 ## Notes
 
-- In ESN workflows, readout weights are typically set via ridge regression in
-  `train!(...)`. Therefore, how `Readout` gets initialized is of no consequence.
-- If you set `include_collect=false`, make sure a [`Collect()`](@ref) appears earlier in the chain.
+- In ESN workflows, readout weights are typically replaced via ridge regression in
+  [`train!`](@ref). Therefore, how `LinearReadout` gets initialized is of no consequence.
+  Additionally, the dimesions will also not be taken into account, as [`train!`](@ref)
+  will replace the weights.
+- If you set `include_collect=false`, make sure a [`Collect`](@ref) appears earlier in the chain.
   Otherwise training may operate on the post-readout signal,
   which is usually unintended.
 """
@@ -106,9 +108,9 @@ end
     Collect()
 
 Marker layer that passes data through unchanged but marks a feature
-checkpoint for [`collectstates`](@ref). At each time step, whenever a `Collect()` is
+checkpoint for [`collectstates`](@ref). At each time step, whenever a `Collect` is
 encountered in the chain, the current vector is recorded as part of the feature
-vector used to train the readout. If multiple `Collect()` layers exist, their
+vector used to train the readout. If multiple `Collect` layers exist, their
 vectors are concatenated with `vcat` in order of appearance.
 
 ## Arguments
@@ -137,13 +139,13 @@ vectors are concatenated with `vcat` in order of appearance.
 
 ## Notes
 
-- When used with a single `Collect()` before a [`LinearReadout`](@ref), training uses exactly
+- When used with a single `Collect` before a [`LinearReadout`](@ref), training uses exactly
   the tensor right before the readout (e.g., the reservoir state).
-- With **multiple** `Collect()` layers (e.g., after different submodules), the
+- With **multiple** `Collect` layers (e.g., after different submodules), the
   per-step features are `vcat`-ed in chain order to form one feature vector.
 - If the readout is constructed with `include_collect=true`, an *implicit*
   collection point is assumed immediately before the readout. Use an explicit
-  `Collect()` only when you want to control where/what is collected (or to stack
+  `Collect` only when you want to control where/what is collected (or to stack
   multiple features).
 
   ```julia
@@ -167,9 +169,9 @@ Base.show(io::IO, cl::Collect) = print(io, "Collection point of states")
     collectstates(rc, data, ps, st)
 
 Run the sequence `data` once through the reservoir chain `rc`, advancing the
-model state over time, and collect feature vectors at every [`Collect()`](@ref) layer.
-If more than one [`Collect()`](ref) is encountered in a step, their vectors are
-concatenated with `vcat` in order of appearance. If no [`Collect()`](@ref) is seen
+model state over time, and collect feature vectors at every [`Collect`](@ref) layer.
+If more than one [`Collect`](@ref) is encountered in a step, their vectors are
+concatenated with `vcat` in order of appearance. If no [`Collect`](@ref) is seen
 in a step, the feature defaults to the final vector exiting the chain for
 that time step.
 
@@ -180,15 +182,15 @@ that time step.
 
 ## Arguments
 
-- `rc`: A [`ReservoirChain`](@ref) (or compatible [`AbstractLuxLayer`](@extref) with `.layers`).
+- `rc`: A [`ReservoirChain`](@ref) (or compatible `AbstractLuxLayer` with `.layers`).
 - `data`: Input sequence of shape `(in_dims, T)`, where columns are time steps.
 - `ps`, `st`: Current parameters and state for `rc`.
 
 ## Returns
 
 - `states`: Reservoir states, i.e. a feature matrix with one column per
   time step. The feature dimension `n_features` equals the vertical concatenation
-  of all vectors captured at [`Collect()`](@ref) layers in that step.
+  of all vectors captured at [`Collect`](@ref) layers in that step.
 - `st`: Updated model states.
 
 """

src/layers/esn_cell.jl

@@ -25,13 +25,13 @@ Echo State Network (ESN) recurrent cell with optional leaky integration.
 
   - `use_bias`: Whether to include a bias term. Default: `false`.
   - `init_bias`: Initializer for the bias. Used only if `use_bias=true`.
-      Default is [`rand32`](@extref).
+      Default is `rand32`.
   - `init_reservoir`: Initializer for the reservoir matrix `W_res`.
     Default is [`rand_sparse`](@ref).
   - `init_input`: Initializer for the input matrix `W_in`.
     Default is [`scaled_rand`](@ref).
   - `init_state`: Initializer for the hidden state when an external
-    state is not provided. Default is [`randn32`](@extref).
+    state is not provided. Default is `randn32`.
   - `leak_coefficient`: Leak rate `α ∈ (0,1]`. Default: `1.0`.
 
 ## Inputs

src/layers/lux_layers.jl

@@ -8,6 +8,22 @@ end
 Base.show(io::IO, wf::WrappedFunction) = print(io, "WrappedFunction(", wf.func, ")")
 
 # adapted from lux layers/recurrent StatefulRecurrentCell
+@doc raw"""
+    StatefulLayer(cell::AbstractReservoirRecurrentCell)
+
+A lightweight wrapper that makes a recurrent cell carry its imput state to the
+next step.
+
+## Arguments
+
+- `cell`: Any `AbstractReservoirRecurrentCell` (e.g. [`ESNCell`](@ref)).
+
+## States
+
+- `cell`: internal states for the wrapped `cell` (e.g., RNG replicas, etc.).
+- `carry`: the per-sequence hidden state; initialized to `nothing`.
+
+"""
 @concrete struct StatefulLayer <: AbstractLuxWrapperLayer{:cell}
     cell <: AbstractReservoirRecurrentCell
 end
@@ -29,15 +45,76 @@ function applyrecurrentcell(sl::AbstractReservoirRecurrentCell, inp, ps, st, ::N
     return apply(sl, inp, ps, st)
 end
 
-###build the ReservoirChain
+@doc raw"""
+    ReservoirChain(layers...; name=nothing)
+    ReservoirChain(xs::AbstractVector; name=nothing)
+    ReservoirChain(nt::NamedTuple; name=nothing)
+    ReservoirChain(; name=nothing, kwargs...)
 
-#abstract type RCLayer <: AbstractLuxLayer end
-#abstract type RCContainerLayer <: AbstractLuxContainerLayer end
+A lightweight, Lux-compatible container that composes a sequence of layers
+and executes them in order. The implementation of `ReservoirChain` is
+equivalent to Lux's own `Chain`.
 
-"""
-    ReservoirChain(layers...)
+## Construction
+
+You can build a chain from:
+
+  - **Positional layers:** `ReservoirChain(l1, l2, ...)`
+  - **A vector of layers:** `ReservoirChain([l1, l2, ...])`
+  - **A named tuple of layers:** `ReservoirChain((; layer_a=l1, layer_b=l2))`
+  - **Keywords (sugar for a named tuple):** `ReservoirChain(; layer_a=l1, layer_b=l2)`
+
+In all cases, function objects are automatically wrapped via `WrappedFunction`
+so they can participate like regular layers. If a [`LinearReadout`](@ref) with
+`include_collect=true` is present, the chain automatically inserts a [`Collect`](@ref)
+layer immediately before that readout.
+
+Use `name` to optionally tag the chain instance.
+
+## Inputs
+
+`(x, ps, st)` where:
+
+  - `x`: input to the first layer.
+  - `ps`: parameters as a named tuple with the same fields and order as the chain's layers.
+  - `st`: states as a named tuple with the same fields and order as the chain's layers.
+
+The call `(c::ReservoirChain)(x, ps, st)` forwards `x` through each layer:
+`(x, ps_i, st_i) -> (x_next, st_i′)` and returns the final output and the
+updated states for every layer.
+
+## Returns
+
+  - `(y, st′)` where `y` is the output of the last layer and `st′` is a named
+    tuple collecting the updated states for each layer.
+
+## Parameters
+
+  - A `NamedTuple` whose fields correspond 1:1 with the layers. Each field
+    holds the parameters for that layer.
+  - Field names are generated as `:layer_1, :layer_2, ...` when constructed
+    positionally, or preserved when you pass a `NamedTuple`/keyword constructor.
+
+## States
+
+  - A `NamedTuple` whose fields correspond 1:1 with the layers. Each field
+    holds the state for that layer.
+
+## Layer access & indexing
+
+  - `c[i]`: get the *i*-th layer (1-based).
+  - `c[indices]`: return a new `ReservoirChain` formed by selecting a subset of layers.
+  - `getproperty(c, :layer_k)`: access layer `k` by its generated/explicit name.
+  - `length(c)`, `firstindex(c)`, `lastindex(c)`: standard collection interfaces.
+
+## Notes
+
+  - **Function wrapping:** Any plain `Function` in the constructor is wrapped as
+    `WrappedFunction(f)`. Non-layer, non-function objects will error.
+  - **Auto-collect for readouts:** When a [`LinearReadout`](@ref) has
+    `include_collect=true`, the constructor expands it to `(Collect(), readout)`
+    so that downstream tooling can capture features consistently.
 
-A simple container that holds a sequence of layers
 """
 @concrete struct ReservoirChain <: AbstractLuxWrapperLayer{:layers}
     layers <: NamedTuple
@@ -82,7 +159,7 @@ end
 (c::ReservoirChain)(x, ps, st::NamedTuple) = applychain(c.layers, x, ps, st)
 
 @generated function applychain(
-    layers::NamedTuple{fields}, x, ps, st::NamedTuple{fields}
+        layers::NamedTuple{fields}, x, ps, st::NamedTuple{fields}
 ) where {fields}
     @assert isa(fields, NTuple{<:Any, Symbol})
     N = length(fields)