diff --git a/.gitignore b/.gitignore index 09c048f8..20e9a512 100644 --- a/.gitignore +++ b/.gitignore @@ -4,4 +4,5 @@ .DS_Store Manifest.toml /dev/ -docs/build \ No newline at end of file +docs/build +AGENTS.md \ No newline at end of file diff --git a/docs/pages.jl b/docs/pages.jl index a15e1791..10fd3650 100644 --- a/docs/pages.jl +++ b/docs/pages.jl @@ -5,6 +5,7 @@ pages = [ "Building a model from scratch" => "tutorials/scratch.md", "Chaos forecasting with an ESN" => "tutorials/lorenz_basic.md", "Continuous-time reservoirs from a SciMLProblem" => "tutorials/sciml_reservoir.md", + "Continuous ESN — forecasting Lorenz" => "tutorials/continuous_esn.md", "Fitting a Next Generation Reservoir Computer" => "tutorials/ngrc.md", "Deep Echo State Networks" => "tutorials/deep_esn.md", "Training Reservoir Computing Models" => "tutorials/train.md", diff --git a/docs/src/api/layers.md b/docs/src/api/layers.md index ddd1fb3f..48928f70 100644 --- a/docs/src/api/layers.md +++ b/docs/src/api/layers.md @@ -36,6 +36,12 @@ LIFESNCell ``` +## Continuous-time Layers + +```@docs + ContinuousESNCell +``` + ## Wrappers ```@docs diff --git a/docs/src/api/models.md b/docs/src/api/models.md index 97b947b5..b306c81c 100644 --- a/docs/src/api/models.md +++ b/docs/src/api/models.md @@ -20,6 +20,12 @@ ResESN ``` +## Continuous-time Echo State Networks + +```@docs + ContinuousESN +``` + ## Next generation reservoir computing ```@docs diff --git a/docs/src/refs.bib b/docs/src/refs.bib index 0f3946ab..9e15668c 100644 --- a/docs/src/refs.bib +++ b/docs/src/refs.bib @@ -557,3 +557,16 @@ @inproceedings{Cossu2025 month = oct, year = {2024} } + +@incollection{Lukosevicius2012, + title = {A Practical Guide to Applying Echo State Networks}, + url = {https://doi.org/10.1007/978-3-642-35289-8_36}, + doi = {10.1007/978-3-642-35289-8_36}, + booktitle = {Neural Networks: Tricks of the Trade, 2nd Edition}, + publisher = {Springer Berlin Heidelberg}, + author = {Luko{\v{s}}evi{\v{c}}ius, Mantas}, + editor = {Montavon, Gr{\'{e}}goire and Orr, Genevi{\`{e}}ve B. and M{\"{u}}ller, Klaus-Robert}, + year = {2012}, + pages = {659--686} +} + diff --git a/docs/src/tutorials/continuous_esn.md b/docs/src/tutorials/continuous_esn.md new file mode 100644 index 00000000..e99a1520 --- /dev/null +++ b/docs/src/tutorials/continuous_esn.md @@ -0,0 +1,119 @@ +# Continuous ESN: forecasting Lorenz + +[`ContinuousESN`](@ref) is a continuous-time Echo State Network that +implements the ODE of [Lukosevicius2012](@cite): + +```math +\dot{\mathbf{x}}(t) = -\mathbf{x}(t) + \tanh\!\left( + \mathbf{W}_{\text{in}}\,\mathbf{u}(t) + \mathbf{W}_r\,\mathbf{x}(t) + + \mathbf{b}\right) +``` + +This tutorial trains a `ContinuousESN` on Lorenz-63 data and rolls it +forward autoregressively to reproduce the attractor. The training and +prediction pipeline is the same as for [`ESN`](@ref). + +## Building a Lorenz dataset + +```@example continuous-esn-lorenz +using ReservoirComputing +using SciMLBase +using DataInterpolations +using OrdinaryDiffEqTsit5 +using Plots +using Random + +Random.seed!(42) +rng = MersenneTwister(17) + +function lorenz!(du, u, p, t) + du[1] = p[1] * (u[2] - u[1]) + du[2] = u[1] * (p[2] - u[3]) - u[2] + du[3] = u[1] * u[2] - p[3] * u[3] +end +data_prob = ODEProblem( + lorenz!, [1.0, 0.0, 0.0], (0.0, 40.0), [10.0, 28.0, 8 / 3] +) +data = Array(solve(data_prob, Tsit5(); saveat = 0.02)) + +shift, train_len, predict_len = 300, 1000, 250 +input_data = data[:, shift:(shift + train_len - 1)] +target_data = data[:, (shift + 1):(shift + train_len)] +test = data[:, (shift + train_len):(shift + train_len + predict_len - 1)] +``` + +## Constructing the `ContinuousESN` + +```@example continuous-esn-lorenz +N_res = 100 + +# Float64 initialisers so the reservoir, the solve, and the input all +# share a numeric type. Without these the cell would default to +# Float32 via `scaled_rand` / `rand_sparse` / `zeros32`. +init_input_f64(rng, d...) = scaled_rand(rng, Float64, d...) +init_reservoir_f64(rng, d...) = rand_sparse(rng, Float64, d...) +init_bias_f64(rng, d...) = zeros(Float64, d...) + +esn_train = ContinuousESN( + 3, N_res, 3, (0.0, Float64(train_len)), Tsit5(); + use_bias = true, + init_input = init_input_f64, + init_reservoir = init_reservoir_f64, + init_bias = init_bias_f64, + state_modifiers = (NLAT2(),), + reltol = 1.0e-6, abstol = 1.0e-8 +) +esn_pred = ContinuousESN( + 3, N_res, 3, (0.0, Float64(predict_len)), Tsit5(); + use_bias = true, + init_input = init_input_f64, + init_reservoir = init_reservoir_f64, + init_bias = init_bias_f64, + state_modifiers = (NLAT2(),), + reltol = 1.0e-6, abstol = 1.0e-8 +) + +ps, st = setup(rng, esn_train) +``` + +## Training + +```@example continuous-esn-lorenz +ps, st = train!(esn_train, input_data, target_data, ps, st) +``` + +## Autoregressive rollout + +```@example continuous-esn-lorenz +ps_pred, st_pred = setup(rng, esn_pred) +ps_pred = merge(ps_pred, (readout = ps.readout,)) +st_pred = merge(st_pred, (readout = st.readout,)) + +output, _ = predict( + esn_pred, predict_len, ps_pred, st_pred; initialdata = test[:, 1] +) + +plot( + transpose(output)[:, 1], transpose(output)[:, 2], + transpose(output)[:, 3]; label = "predicted" +) +plot!( + transpose(test)[:, 1], transpose(test)[:, 2], + transpose(test)[:, 3]; label = "actual" +) +``` + +The two trajectories agree on the early portion of the rollout before +chaotic divergence dominates — the same behaviour the discrete-ESN +tutorial produces. The point of the example is that nothing in the +training loop changes between discrete ESN, `SciMLProblemReservoir` +with hand-rolled equations, and `ContinuousESN`: the same `train!` / +`predict` pipeline drives all three. + +## When to reach for `ContinuousESN` vs `SciMLProblemReservoir` + +* `ContinuousESN` pre-bakes the continuous ESN ODE; use it when the + standard continuous ESN is what you want. +* [`SciMLProblemReservoir`](@ref) is the generic building block; use it + when the reservoir ODE is not the standard eq (5) — bespoke RHS, SDE, + DDE, or non-standard parameter layout. diff --git a/ext/RCODEReservoirExt.jl b/ext/RCODEReservoirExt.jl index 661f21cf..529ab3f7 100644 --- a/ext/RCODEReservoirExt.jl +++ b/ext/RCODEReservoirExt.jl @@ -1,37 +1,25 @@ module RCODEReservoirExt using DataInterpolations: ConstantInterpolation +using LinearAlgebra: mul! using LuxCore: apply -# `solve` and `remake` come from `SciMLBase`. The user picks the concrete -# solver type (e.g. `Tsit5()`) and loads its package separately -# (`OrdinaryDiffEqTsit5`, `OrdinaryDiffEq`, …); dispatch at solve time -# selects the right method via the type they passed in `res.args[1]`. We -# deliberately don't list a solver package as a weakdep trigger so users -# aren't forced to pull the full `OrdinaryDiffEq` meta-package in. -using SciMLBase: remake, solve, NullParameters +using Random: AbstractRNG +using SciMLBase: ODEProblem, remake, solve, NullParameters using ReservoirComputing: ReservoirComputing, AbstractReservoirComputer, AbstractSampler, AbstractSciMLProblemReservoir, + ContinuousESN, + ContinuousESNCell, + LinearReadout, TerminalStateSampling, - collectstates + _wrap_layers, + collectstates, + rand_sparse, + scaled_rand import ReservoirComputing: _collectstates, _predict -# --------------------------------------------------------------------------- -# Parameter assembly -# -# At solve time the extension must hand the ODE three things: -# (1) the interpolated input signal `u(t)` exposed as `p.input`, -# (2) the user's static parameters (if any) from `prob.p`, -# (3) any Lux-managed reservoir parameters from `ps.reservoir`. -# -# `_to_namedtuple` normalises `prob.p` so that nothing / `NullParameters` / a -# user `NamedTuple` all collapse into a single `NamedTuple` we can merge into. -# Anything else is rejected with a clear error — wrapping unknown payloads -# silently would hide bugs in user-defined ODEs. -# --------------------------------------------------------------------------- - _to_namedtuple(prob_p::NamedTuple) = prob_p _to_namedtuple(::NullParameters) = NamedTuple() _to_namedtuple(::Nothing) = NamedTuple() @@ -75,19 +63,6 @@ function _build_solve_params(prob_p, ps_reservoir, input_interp) return merge(merged, (input = input_interp,)) end -# --------------------------------------------------------------------------- -# Input signal construction -# -# `collectstates` sees a discrete `data::AbstractMatrix` and reconstructs the -# continuous-time input via linear interpolation between input columns. The -# grid mirrors the `saveat` grid so an input column and its corresponding -# state sample share the same time stamp. -# -# `_make_const_input_fn` is the closed-loop counterpart used inside the -# autoregressive `predict`: between two reservoir-output events the input is -# the previous output, held constant. -# --------------------------------------------------------------------------- - """ ZeroOrderHoldInterp(data, ts) @@ -135,34 +110,13 @@ function _make_input_fn(data::AbstractMatrix, ts::AbstractVector) end function _make_const_input_fn(u_vec::AbstractVector, t_lo, t_hi) - # `cache_parameters=true` is fine for vector u (autoregressive predict - # always holds the previous readout output constant over one sub-interval). return ConstantInterpolation([u_vec, u_vec], [t_lo, t_hi]; cache_parameters = true) end -# --------------------------------------------------------------------------- -# Samplers -# -# A sampler maps a continuous trajectory into the discrete state matrix the -# readout sees. `TerminalStateSampling` reads the solution exactly at the -# user-visible time grid (the same one we pass through `saveat`), so the -# result is just the columnar view of `sol.u`. -# --------------------------------------------------------------------------- - function _sample(::TerminalStateSampling, sol) return reduce(hcat, sol.u) end -# --------------------------------------------------------------------------- -# State-modifier composition -# -# The discrete fallback threads `states_modifiers` per reservoir step (see -# `_partial_apply` in `reservoircomputer.jl`). For the continuous path we -# evolve the trajectory first and then apply modifiers column-by-column to -# the sampled matrix. This keeps the per-sample semantics identical to the -# discrete code without contaminating the ODE right-hand side. -# --------------------------------------------------------------------------- - function _apply_modifiers_continuous( modifiers::Tuple, states_matrix::AbstractMatrix, ps_mods, st_mods ) @@ -188,28 +142,6 @@ function _apply_modifiers_continuous( return output, new_st end -# --------------------------------------------------------------------------- -# Continuous `_collectstates` -# -# Pipeline: -# 1. Split `res.tspan` into `n_samples` equal-width windows. -# 2. Place input column `k` at the *start* of window `k` (time -# `t0 + (k-1)Δt`) and request a sample at the *end* of window `k` -# (time `t0 + kΔt`). This alignment matches the discrete reservoir -# semantics — `states[:, k]` is the state after processing input `k` -# — and is what makes the Euler-equivalence test land without an -# off-by-one shift. -# 3. `remake` the user's problem with the locked `tspan` and the merged -# parameter pack (interpolated input injected as `p.input`). -# 4. `solve(...; saveat = sample_ts, save_everystep=false, dense=false)`. -# `res.kwargs` come last so user kwargs win on collision — the -# constructor already rejects the three protected keys, so they -# cannot collide in practice. -# 5. Push the trajectory through the sampler → raw state matrix. -# 6. Apply state modifiers → final state matrix matching the discrete -# `(state_dims, n_samples)` shape expected by the readout. -# --------------------------------------------------------------------------- - function _collectstates( res::AbstractSciMLProblemReservoir, rc::AbstractReservoirComputer, @@ -264,14 +196,6 @@ function _collectstates( return modified_states, newst end -# --------------------------------------------------------------------------- -# Teacher-forced `predict` -# -# Solve once over the whole tspan, then apply the readout column-by-column. -# Cheaper than the autoregressive path because the ODE never has to be -# restarted between samples. -# --------------------------------------------------------------------------- - function _predict( ::AbstractSciMLProblemReservoir, rc::AbstractReservoirComputer, @@ -293,22 +217,6 @@ function _predict( return outputs, merge(new_st, (readout = st_ro,)) end -# --------------------------------------------------------------------------- -# Autoregressive `predict` -# -# Split `tspan` into `steps` equal sub-intervals. For each sub-interval the -# input is the previous readout output, held constant via a -# `ConstantInterpolation`. After each sub-solve we: -# - sample the terminal state, -# - apply state modifiers (per-sample, consistent with the discrete loop), -# - apply the readout, -# - feed the output back as the next input. -# -# The initial reservoir state is `res.prob.u0`; users who want to continue -# from a previously computed trajectory should `remake(prob; u0 = …)` before -# constructing the reservoir. -# --------------------------------------------------------------------------- - function _predict( res::AbstractSciMLProblemReservoir, rc::AbstractReservoirComputer, @@ -389,4 +297,178 @@ function _predict( return outputs, newst end +function ReservoirComputing.ContinuousESN( + in_dims::Integer, res_dims::Integer, out_dims::Integer, + activation, tspan, args...; + use_bias::Bool = false, + init_bias = ReservoirComputing.zeros32, + init_reservoir = rand_sparse, + init_input = scaled_rand, + init_state = ReservoirComputing.randn32, + equations = ReservoirComputing._continuous_esn_rhs!, + state_modifiers = (), + readout_activation = identity, + kwargs... + ) + in_dims > 0 || throw(ArgumentError("in_dims must be positive, got $in_dims")) + res_dims > 0 || throw(ArgumentError("res_dims must be positive, got $res_dims")) + out_dims > 0 || throw(ArgumentError("out_dims must be positive, got $out_dims")) + length(tspan) == 2 || throw( + ArgumentError( + "tspan must be a length-2 tuple/pair (t0, t1), got length $(length(tspan))" + ) + ) + (isfinite(tspan[1]) && isfinite(tspan[2])) || throw( + ArgumentError("tspan endpoints must be finite, got $tspan") + ) + tspan[2] > tspan[1] || throw( + ArgumentError( + "ContinuousESN requires `tspan[2] > tspan[1]`, got tspan = $tspan" + ) + ) + ReservoirComputing._check_protected_kwargs(kwargs) + + cell = ContinuousESNCell( + activation, in_dims, res_dims, + init_bias, init_reservoir, init_input, init_state, + ReservoirComputing.static(use_bias), + equations, tspan, args, kwargs + ) + + mods_tuple = state_modifiers isa Tuple || state_modifiers isa AbstractVector ? + Tuple(state_modifiers) : (state_modifiers,) + mods = _wrap_layers(mods_tuple) + + readout = LinearReadout(res_dims => out_dims, readout_activation) + return ContinuousESN(cell, mods, readout) +end + +function ReservoirComputing.ContinuousESN( + in_dims::Integer, res_dims::Integer, out_dims::Integer, + tspan::Union{Tuple, Pair}, args...; kwargs... + ) + return ContinuousESN(in_dims, res_dims, out_dims, tanh, tspan, args...; kwargs...) +end + +function _collectstates( + cell::ContinuousESNCell, + rc::AbstractReservoirComputer, + data::AbstractMatrix, + ps::NamedTuple, + st::NamedTuple + ) + n_samples = size(data, 2) + n_samples ≥ 2 || throw( + ArgumentError( + "ContinuousESN collectstates needs at least 2 input columns " * + "to define a time grid; got $n_samples." + ) + ) + + t0, t1 = cell.tspan + t1 > t0 || throw( + ArgumentError( + "ContinuousESN requires `tspan[2] > tspan[1]`, got tspan = ($t0, $t1)." + ) + ) + + Δt = (t1 - t0) / n_samples + input_ts = collect(range(t0, t1 - Δt; length = n_samples)) + sample_ts = collect(range(t0 + Δt, t1; length = n_samples)) + + input_interp = _make_input_fn(data, input_ts) + solve_p = _build_solve_params(nothing, ps.reservoir, input_interp) + + # `u0` element type follows `ps.reservoir.input_matrix` so the solver + # state, the parameter pack, and the input signal share a numeric + # type. The user controls eltype through the `init_*` initialisers. + u0 = zeros(eltype(ps.reservoir.input_matrix), cell.out_dims) + prob = ODEProblem(cell.equations, u0, cell.tspan, solve_p) + + sol = solve( + prob, cell.args...; + saveat = sample_ts, + save_everystep = false, + dense = false, + cell.kwargs... + ) + + raw_states = _sample(TerminalStateSampling(), sol) + modified_states, st_mods = _apply_modifiers_continuous( + rc.states_modifiers, raw_states, ps.states_modifiers, st.states_modifiers + ) + + newst = ( + reservoir = st.reservoir, + states_modifiers = st_mods, + readout = st.readout, + ) + return modified_states, newst +end + +function _predict( + cell::ContinuousESNCell, + rc::AbstractReservoirComputer, + steps::Integer, + ps::NamedTuple, + st::NamedTuple; + initialdata::AbstractVector + ) + steps ≥ 1 || throw(ArgumentError("steps must be ≥ 1, got $steps")) + + t0, t1 = cell.tspan + t1 > t0 || throw( + ArgumentError( + "Autoregressive predict requires `tspan[2] > tspan[1]`, got " * + "tspan = ($t0, $t1)." + ) + ) + ts = collect(range(t0, t1; length = steps + 1)) + window_starts = @view ts[1:(end - 1)] + window_ends = @view ts[2:end] + + current_state = zeros(eltype(ps.reservoir.input_matrix), cell.out_dims) + current_input = initialdata + + st_mods = st.states_modifiers + st_ro = st.readout + + local outputs + for (step_idx, (t_lo, t_hi)) in enumerate(zip(window_starts, window_ends)) + input_fn = _make_const_input_fn(current_input, t_lo, t_hi) + solve_p = _build_solve_params(nothing, ps.reservoir, input_fn) + sub_prob = ODEProblem(cell.equations, current_state, (t_lo, t_hi), solve_p) + sol = solve( + sub_prob, cell.args...; + saveat = [t_hi], + save_everystep = false, + dense = false, + cell.kwargs... + ) + current_state = sol.u[end] + + if !isempty(rc.states_modifiers) + state_after_mods, st_mods = ReservoirComputing._apply_seq( + rc.states_modifiers, current_state, ps.states_modifiers, st_mods + ) + else + state_after_mods = current_state + end + + current_output, st_ro = apply(rc.readout, state_after_mods, ps.readout, st_ro) + if step_idx == 1 + outputs = similar(current_output, length(current_output), steps) + end + outputs[:, step_idx] .= current_output + current_input = current_output + end + + newst = ( + reservoir = st.reservoir, + states_modifiers = st_mods, + readout = st_ro, + ) + return outputs, newst +end + end # module diff --git a/src/ReservoirComputing.jl b/src/ReservoirComputing.jl index 3bec529a..71ff7eb4 100644 --- a/src/ReservoirComputing.jl +++ b/src/ReservoirComputing.jl @@ -25,6 +25,7 @@ include("reservoircomputer.jl") include("layers/basic.jl") include("layers/lux_layers.jl") include("layers/esn_cell.jl") +include("layers/continuous_esn_cell.jl") include("layers/additive_eiesn_cell.jl") include("layers/eiesn_cell.jl") include("layers/es2n_cell.jl") @@ -59,12 +60,14 @@ include("models/lifesn.jl") include("models/ngrc.jl") include("models/rmnesn.jl") include("models/rmnresesn.jl") +include("models/continuous_esn.jl") #extensions include("extensions/reca.jl") export ReservoirComputer -export AbstractSciMLProblemReservoir, SciMLProblemReservoir +export AbstractSciMLProblemReservoir, SciMLProblemReservoir, ContinuousESN export AbstractSampler, TerminalStateSampling +export ContinuousESNCell export AdditiveEIESNCell, EIESNCell, ES2NCell, ESNCell, EuSNCell, LIFESNCell, MemoryESNCell, MemoryResESNCell, ResESNCell, RMNCell export StatefulLayer, LinearReadout, ReservoirChain, Collect, collectstates, diff --git a/src/layers/continuous_esn_cell.jl b/src/layers/continuous_esn_cell.jl new file mode 100644 index 00000000..3ec23d34 --- /dev/null +++ b/src/layers/continuous_esn_cell.jl @@ -0,0 +1,110 @@ +@doc raw""" + ContinuousESNCell(in_dims => out_dims; + tspan, args = (), activation = tanh, use_bias = false, + init_bias = zeros32, init_reservoir = rand_sparse, + init_input = scaled_rand, init_state = zeros32, + equations = _continuous_esn_rhs!, kwargs...) + +Continuous-time Echo State Network cell +([Lukosevicius2012](@cite)). Integrates the ODE + +```math +\dot{\mathbf{x}}(t) = -\mathbf{x}(t) + \tanh\!\left( + \mathbf{W}_{\text{in}}\,\mathbf{u}(t) + + \mathbf{W}_r\,\mathbf{x}(t) + \mathbf{b}\right) +``` + +## Arguments + + - `in_dims`: Input dimension. + - `out_dims`: Reservoir (hidden state) dimension. + +## Keyword arguments + + - `tspan`: Integration interval `(t0, t1)`. Length-2, strictly + increasing, finite. + - `args`: Tuple of positional arguments forwarded to `solve`. The solver + algorithm (`Tsit5()`, `Euler()`, …) is the first element by convention. + Default: `()`. + - `activation`: Reservoir activation. Default: `tanh`. + - `use_bias`: Whether to include a bias term. Default: `false`. + - `init_reservoir`: Initialiser for `W_r`. Default: [`rand_sparse`](@ref). + - `init_input`: Initialiser for `W_in`. Default: [`scaled_rand`](@ref). + - `init_bias`: Initialiser for the bias. Only used if `use_bias=true`. + Default: `zeros32`. + - `init_state`: Initialiser for the initial hidden state. + Default: `zeros32`. + - `equations`: ODE right-hand side `(dx, x, p, t) -> nothing`. + Default: continuous-time leaky-integrator ESN ODE. + - `kwargs...`: Forwarded to `solve` as keyword arguments. + +## Parameters + + - `input_matrix :: (out_dims × in_dims)` — `W_in` + - `reservoir_matrix :: (out_dims × out_dims)` — `W_r` + - `bias :: (out_dims,)` — present only if `use_bias = true` +""" +@concrete struct ContinuousESNCell <: AbstractSciMLProblemReservoir + activation + in_dims <: IntegerType + out_dims <: IntegerType + init_bias + init_reservoir + init_input + init_state + use_bias <: StaticBool + equations + tspan + args + kwargs +end + +function ContinuousESNCell( + (in_dims, out_dims)::Pair{<:IntegerType, <:IntegerType}; + tspan, args = (), activation = tanh, use_bias::BoolType = False(), + init_bias = zeros32, init_reservoir = rand_sparse, + init_input = scaled_rand, init_state = zeros32, + equations = _continuous_esn_rhs!, kwargs... + ) + return ContinuousESNCell( + activation, in_dims, out_dims, + init_bias, init_reservoir, init_input, init_state, + static(use_bias), + equations, tspan, args, kwargs + ) +end + +function initialparameters(rng::AbstractRNG, cell::ContinuousESNCell) + ps = ( + input_matrix = cell.init_input(rng, cell.out_dims, cell.in_dims), + reservoir_matrix = cell.init_reservoir(rng, cell.out_dims, cell.out_dims), + ) + if has_bias(cell) + ps = merge(ps, (bias = cell.init_bias(rng, cell.out_dims),)) + end + return ps +end + +function initialstates(::AbstractRNG, ::ContinuousESNCell) + return NamedTuple() +end + +function _continuous_esn_rhs!(dx, x, p, t) + input_t = p.input(t) + mul!(dx, p.reservoir_matrix, x) + mul!(dx, p.input_matrix, input_t, true, true) + if haskey(p, :bias) + @. dx = -x + tanh(dx + p.bias) + else + @. dx = -x + tanh(dx) + end + return nothing +end + +function Base.show(io::IO, cell::ContinuousESNCell) + print(io, "ContinuousESNCell(", cell.in_dims, " => ", cell.out_dims) + print(io, ", tspan = ") + show(io, cell.tspan) + has_bias(cell) || print(io, ", use_bias = false") + return print(io, ")") +end diff --git a/src/layers/sciml_reservoir.jl b/src/layers/sciml_reservoir.jl index 1558588e..2800961c 100644 --- a/src/layers/sciml_reservoir.jl +++ b/src/layers/sciml_reservoir.jl @@ -110,9 +110,8 @@ function SciMLProblemReservoir(prob, sampler, tspan, args...; kwargs...) return SciMLProblemReservoir(prob, sampler, tspan, args, kwargs) end -# Empty parameters/state by default. Concrete CTESN/LSM subtypes that land -# later (PR3, PR5) override these to expose `ps.reservoir.ode_params` and -# `st.reservoir` caches. +# Empty parameters/state by default. Concrete subtypes (e.g. `ContinuousESN`) +# override these to expose reservoir matrices and any solver caches. function initialparameters(::AbstractRNG, ::AbstractSciMLProblemReservoir) return NamedTuple() end diff --git a/src/models/continuous_esn.jl b/src/models/continuous_esn.jl new file mode 100644 index 00000000..d5d95a15 --- /dev/null +++ b/src/models/continuous_esn.jl @@ -0,0 +1,127 @@ +@doc raw""" + ContinuousESN(in_dims, res_dims, out_dims, [activation,] tspan, args...; + use_bias = false, + init_reservoir = rand_sparse, init_input = scaled_rand, + init_bias = zeros32, init_state = randn32, + equations = _continuous_esn_rhs!, + state_modifiers = (), readout_activation = identity, + kwargs...) + +Continuous-time Echo State Network +([Lukosevicius2012](@cite)). Composes a [`ContinuousESNCell`](@ref), +optional `state_modifiers`, and a [`LinearReadout`](@ref). + +## Equations + +```math +\begin{aligned} + \dot{\mathbf{x}}(t) &= -\mathbf{x}(t) + \tanh\!\left( + \mathbf{W}_{\text{in}}\,\mathbf{u}(t) + \mathbf{W}_r\,\mathbf{x}(t) + + \mathbf{b}\right) \\ + \mathbf{z}(t) &= \mathrm{Mods}\!\left(\mathbf{x}(t)\right) \\ + \mathbf{y}(t) &= \rho\!\left(\mathbf{W}_{\text{out}}\,\mathbf{z}(t) + + \mathbf{b}_{\text{out}}\right) +\end{aligned} +``` + +## Arguments + + - `in_dims`: Input dimension. + - `res_dims`: Reservoir (hidden state) dimension. + - `out_dims`: Output dimension. + - `activation`: Reservoir activation. Default: `tanh`. + - `tspan`: Integration interval `(t0, t1)` for `collectstates`. + Length-2, strictly increasing, finite. + - `args...`: Forwarded to `solve` positionally. The solver algorithm + (`Tsit5()`, `Euler()`, …) is the first element by convention. + +## Keyword arguments + +Reservoir (passed to [`ContinuousESNCell`](@ref)): + + - `use_bias`: Whether the reservoir uses a bias term. Default: `false`. + - `init_reservoir`: Initialiser for `W_r`. Default: [`rand_sparse`](@ref). + - `init_input`: Initialiser for `W_in`. Default: [`scaled_rand`](@ref). + - `init_bias`: Initialiser for `b`. Default: `zeros32`. + - `init_state`: Initialiser for the initial hidden state. + Default: `randn32`. + - `equations`: ODE right-hand side function `(dx, x, p, t) -> nothing`. + Default is the continuous-time leaky-integrator ESN ODE. + +Composition: + + - `state_modifiers`: Layers applied to reservoir states before the + readout. Accepts a single layer, an `AbstractVector`, or a `Tuple`. + Default: empty `()`. + - `readout_activation`: Activation for the linear readout. + Default: `identity`. + +Solve metadata: + + - `kwargs...`: Forwarded to `solve`. The keys `saveat`, + `save_everystep`, and `dense` are reserved and rejected at + construction. + +## Parameters + + - `reservoir` — parameters of the internal [`ContinuousESNCell`](@ref): + - `input_matrix :: (res_dims × in_dims)` — `W_in` + - `reservoir_matrix :: (res_dims × res_dims)` — `W_r` + - `bias :: (res_dims,)` — present only if `use_bias = true` + - `states_modifiers` — parameters for each modifier layer (may be empty). + - `readout` — parameters of [`LinearReadout`](@ref): + - `weight :: (out_dims × res_dims)` — `W_out` + - `bias :: (out_dims,)` — `b_out` (if the readout uses bias) + +## States + + - `reservoir` — states for the internal [`ContinuousESNCell`](@ref). + - `states_modifiers` — states for each modifier layer (may be empty). + - `readout` — states for [`LinearReadout`](@ref). + +!!! note + The `RCODEReservoirExt` extension must be loaded for this constructor + to succeed. Load a solver package such as `OrdinaryDiffEqTsit5` + alongside `SciMLBase` and `DataInterpolations`. +""" +@concrete struct ContinuousESN <: + AbstractEchoStateNetwork{(:reservoir, :states_modifiers, :readout)} + reservoir + states_modifiers + readout +end + +function ContinuousESN(::Any, ::Any, ::Any, ::Any...; kwargs...) + return error( + "ContinuousESN requires the RCODEReservoirExt extension and an " * + "OrdinaryDiffEq solver package. Load `SciMLBase`, `DataInterpolations`, " * + "and a solver package (e.g. `OrdinaryDiffEqTsit5`) to enable it." + ) +end + +function Base.show(io::IO, esn::ContinuousESN) + print(io, "ContinuousESN(\n") + + print(io, " reservoir = ") + show(io, esn.reservoir) + print(io, ",\n") + + print(io, " state_modifiers = ") + if isempty(esn.states_modifiers) + print(io, "()") + else + print(io, "(") + for (idx, mod) in enumerate(esn.states_modifiers) + idx > 1 && print(io, ", ") + show(io, mod) + end + print(io, ")") + end + print(io, ",\n") + + print(io, " readout = ") + show(io, esn.readout) + print(io, "\n)") + + return +end diff --git a/test/test_continuous_esn.jl b/test/test_continuous_esn.jl new file mode 100644 index 00000000..07dca915 --- /dev/null +++ b/test/test_continuous_esn.jl @@ -0,0 +1,204 @@ +using Test +using Random +using LinearAlgebra +using Statistics +using ReservoirComputing +using OrdinaryDiffEq +using SciMLBase +using DataInterpolations + +@testset "ContinuousESN: construction + parameter shapes" begin + rng = MersenneTwister(0) + in_dim, res_dim, out_dim = 3, 50, 2 + esn = ContinuousESN(in_dim, res_dim, out_dim, (0.0, 5.0), Tsit5()) + + @test esn isa ContinuousESN + @test propertynames(esn) == (:reservoir, :states_modifiers, :readout) + @test esn.reservoir isa ContinuousESNCell + @test esn.readout isa LinearReadout + @test isempty(esn.states_modifiers) + + ps, st = setup(rng, esn) + @test size(ps.reservoir.input_matrix) == (res_dim, in_dim) + @test size(ps.reservoir.reservoir_matrix) == (res_dim, res_dim) + @test !haskey(ps.reservoir, :bias) + @test all(isfinite, ps.reservoir.input_matrix) + @test all(isfinite, ps.reservoir.reservoir_matrix) + @test size(ps.readout.weight) == (out_dim, res_dim) +end + +@testset "ContinuousESN: construction validation" begin + @test_throws ArgumentError ContinuousESN(0, 5, 2, (0.0, 1.0), Tsit5()) + @test_throws ArgumentError ContinuousESN(3, 0, 2, (0.0, 1.0), Tsit5()) + @test_throws ArgumentError ContinuousESN(3, 5, 0, (0.0, 1.0), Tsit5()) + @test_throws ArgumentError ContinuousESN(3, 5, 2, (0.0, 0.0), Tsit5()) + @test_throws ArgumentError ContinuousESN(3, 5, 2, (1.0, 0.0), Tsit5()) + @test_throws ArgumentError ContinuousESN(3, 5, 2, (0.0, Inf), Tsit5()) + @test_throws ArgumentError ContinuousESN(3, 5, 2, (-Inf, 1.0), Tsit5()) + @test_throws ArgumentError ContinuousESN(3, 5, 2, (0.0, NaN), Tsit5()) + @test_throws ArgumentError ContinuousESN(3, 5, 2, (1.0,), Tsit5()) + @test_throws ArgumentError ContinuousESN(3, 5, 2, (0.0, 1.0, 2.0), Tsit5()) + for badkw in (:saveat, :save_everystep, :dense) + @test_throws ArgumentError ContinuousESN( + 3, 5, 2, (0.0, 1.0), Tsit5(); (badkw => true,)... + ) + end +end + +@testset "ContinuousESN: bias toggle" begin + rng = MersenneTwister(1) + in_dim, res_dim, out_dim = 2, 10, 1 + esn_b = ContinuousESN( + in_dim, res_dim, out_dim, (0.0, 1.0), Tsit5(); use_bias = true + ) + ps, st = setup(rng, esn_b) + @test haskey(ps.reservoir, :bias) + @test length(ps.reservoir.bias) == res_dim + @test all(==(0), ps.reservoir.bias) +end + +@testset "ContinuousESN: bias path under Tsit5" begin + rng = MersenneTwister(101) + in_dim, res_dim, out_dim, T_steps = 2, 20, 1, 12 + nonzero_bias(rng, d...) = 0.5f0 .* randn(rng, Float32, d...) + esn_with_bias = ContinuousESN( + in_dim, res_dim, out_dim, (0.0, 2.0), Tsit5(); + use_bias = true, init_bias = nonzero_bias, + reltol = 1.0e-8, abstol = 1.0e-10 + ) + esn_no_bias = ContinuousESN( + in_dim, res_dim, out_dim, (0.0, 2.0), Tsit5(); + use_bias = false, reltol = 1.0e-8, abstol = 1.0e-10 + ) + ps_b, st_b = setup(MersenneTwister(0), esn_with_bias) + ps_n, st_n = setup(MersenneTwister(0), esn_no_bias) + data = randn(Float32, in_dim, T_steps) + + s_b, _ = collectstates(esn_with_bias, data, ps_b, st_b) + s_n, _ = collectstates(esn_no_bias, data, ps_n, st_n) + @test all(isfinite, s_b) + @test s_b != s_n +end + +@testset "ContinuousESN: forward (collectstates)" begin + rng = MersenneTwister(7) + in_dim, res_dim, out_dim, T_steps = 2, 16, 1, 25 + esn = ContinuousESN( + in_dim, res_dim, out_dim, (0.0, 3.0), Tsit5(); + reltol = 1.0e-8, abstol = 1.0e-10 + ) + ps, st = setup(rng, esn) + data = randn(Float32, in_dim, T_steps) + + s1, _ = collectstates(esn, data, ps, st) + s2, _ = collectstates(esn, data, ps, st) + @test size(s1) == (res_dim, T_steps) + @test all(isfinite, s1) + @test s1 ≈ s2 +end + +@testset "ContinuousESN: Euler equivalence with discrete leaky ESN" begin + dense_init_f64(rng, d...) = rand_sparse(rng, Float64, d...; sparsity = 0.5) + init_input_f64(rng, d...) = scaled_rand(rng, Float64, d...) + init_bias_f64(rng, d...) = zeros(Float64, d...) + for α in (1.0, 0.5) + rng = MersenneTwister(13) + in_dim, res_dim, out_dim, T_steps = 2, 16, 1, 12 + tspan = (0.0, T_steps * α) + + esn = ContinuousESN( + in_dim, res_dim, out_dim, tspan, Euler(); + use_bias = true, dt = α, + init_reservoir = dense_init_f64, + init_input = init_input_f64, + init_bias = init_bias_f64 + ) + ps, st = setup(MersenneTwister(0), esn) + + data = randn(rng, Float64, in_dim, T_steps) + cont_states, _ = collectstates(esn, data, ps, st) + + W_r = ps.reservoir.reservoir_matrix + W_in = ps.reservoir.input_matrix + b = ps.reservoir.bias + disc_states = zeros(Float64, res_dim, T_steps) + x = zeros(Float64, res_dim) + for (k, u_col) in enumerate(eachcol(data)) + x = (1 - α) .* x .+ α .* tanh.(W_r * x .+ W_in * u_col .+ b) + disc_states[:, k] = x + end + @test cont_states ≈ disc_states atol = 1.0e-10 + end +end + +@testset "ContinuousESN: teacher-forced predict" begin + rng = MersenneTwister(21) + in_dim, res_dim, out_dim, T_steps = 2, 12, 3, 10 + esn = ContinuousESN( + in_dim, res_dim, out_dim, (0.0, 2.0), Tsit5(); + reltol = 1.0e-8, abstol = 1.0e-10 + ) + ps, st = setup(rng, esn) + + data = randn(Float32, in_dim, T_steps) + tf1, _ = predict(esn, data, ps, st) + tf2, _ = predict(esn, data, ps, st) + @test size(tf1) == (out_dim, T_steps) + @test all(isfinite, tf1) + @test tf1 ≈ tf2 +end + +@testset "ContinuousESN: autoregressive predict" begin + rng = MersenneTwister(22) + # Autoregressive rollout feeds outputs back as inputs, so in_dim == out_dim. + dim, res_dim, steps = 3, 12, 5 + esn = ContinuousESN( + dim, res_dim, dim, (0.0, 1.0), Tsit5(); + reltol = 1.0e-8, abstol = 1.0e-10 + ) + ps, st = setup(rng, esn) + + init = randn(Float32, dim) + ar1, _ = predict(esn, steps, ps, st; initialdata = init) + ar2, _ = predict(esn, steps, ps, st; initialdata = init) + @test size(ar1) == (dim, steps) + @test all(isfinite, ar1) + @test ar1 ≈ ar2 +end + +@testset "ContinuousESN: state modifiers compose" begin + rng = MersenneTwister(31) + in_dim, res_dim, out_dim, T_steps = 2, 10, 1, 8 + esn_plain = ContinuousESN( + in_dim, res_dim, out_dim, (0.0, 1.0), Tsit5(); + reltol = 1.0e-8, abstol = 1.0e-10 + ) + esn_mod = ContinuousESN( + in_dim, res_dim, out_dim, (0.0, 1.0), Tsit5(); + reltol = 1.0e-8, abstol = 1.0e-10, + state_modifiers = (NLAT2(),) + ) + + ps_p, st_p = setup(MersenneTwister(0), esn_plain) + ps_m, st_m = setup(MersenneTwister(0), esn_mod) + data = randn(Float32, in_dim, T_steps) + + sp, _ = collectstates(esn_plain, data, ps_p, st_p) + sm, _ = collectstates(esn_mod, data, ps_m, st_m) + @test size(sm) == size(sp) + @test all(isfinite, sm) + @test sm != sp +end + +@testset "ContinuousESN: custom init eltype propagates" begin + rng = MersenneTwister(43) + init_input_f64(rng, d...) = scaled_rand(rng, Float64, d...) + init_res_f64(rng, d...) = rand_sparse(rng, Float64, d...) + esn = ContinuousESN( + 2, 32, 1, (0.0, 1.0), Tsit5(); + init_input = init_input_f64, init_reservoir = init_res_f64 + ) + ps, _ = setup(rng, esn) + @test eltype(ps.reservoir.input_matrix) == Float64 + @test eltype(ps.reservoir.reservoir_matrix) == Float64 +end