From 4b5ee22d8df9e4dbc648adad075edb5b93db6560 Mon Sep 17 00:00:00 2001 From: Saswat Susmoy Date: Wed, 17 Jun 2026 02:14:40 +0530 Subject: [PATCH 1/8] chore(ctesn): add discussion-stub for PR3 CTESN model MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Skeleton-only placeholder so the draft PR exists as a discussion vehicle. No include, no export, no tests — the body errors with a clear message directing the user to the (not-yet-implemented) extension constructor. Open design questions are tracked in the PR description. Stacked on PR #450; will be rebased onto master once that merges. --- src/extensions/ctesn.jl | 41 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) create mode 100644 src/extensions/ctesn.jl diff --git a/src/extensions/ctesn.jl b/src/extensions/ctesn.jl new file mode 100644 index 000000000..3660e8dad --- /dev/null +++ b/src/extensions/ctesn.jl @@ -0,0 +1,41 @@ +@doc raw""" + CTESN(in_dims, res_dims, out_dims; kwargs...) + +Continuous-Time Echo State Network ([Anantharaman2021](@cite)). **Work in progress +— PR3 of the SciML 2026 fellowship roadmap (#397).** This is a stub: the real +constructor lives in the `RCODEReservoirExt` package extension and will be +filled in once the design questions on the open pull request are resolved. + +`CTESN <: AbstractSciMLProblemReservoir`, so it reuses the continuous-time +`_collectstates` / `_predict` machinery landed in PR #450 (`add-ode-reservoir-ext`). +It specialises that machinery to a concrete reservoir ODE. + +The canonical reservoir ODE from [Anantharaman2021](@cite) eq. (3) is: + +```math +\dot{\mathbf{r}}(t) = \tanh\!\left(A\,\mathbf{r}(t) + W_{\text{hyb}}\,\mathbf{x}(t)\right) +``` + +with no leak term and no bias — `f = tanh`, `g = id` hardcoded by the paper. +The implementation may also expose an opt-in leaky form +`ṙ = -α r + tanh(W_in u(t) + W_r r + b)` (the continuous limit of the +discrete leaky ESN of [Lukosevicius2012](@cite)); see the PR thread for the +default-form decision. + +!!! note + This constructor errors unless the `RCODEReservoirExt` extension is loaded + (`SciMLBase` + `DataInterpolations`) **and** a concrete `OrdinaryDiffEq` + solver package is available for `args[1]`. + +## Status + +- PR3 work in progress on branch `add-ctesn`, stacked on top of PR #450 + (`add-ode-reservoir-ext`). Will be rebased onto `master` once PR #450 merges. +- Open design questions are in the pull-request description. +""" +CTESN(::Any...) = error( + "CTESN requires the RCODEReservoirExt extension and an OrdinaryDiffEq " * + "solver package. Load `SciMLBase`, `DataInterpolations`, and a solver " * + "package (e.g. `OrdinaryDiffEqTsit5`) to enable it. PR3 is a " * + "work-in-progress — the real constructor is not yet wired up." +) From b2f9883a0dca3abd2a9c578774f98de422e00853 Mon Sep 17 00:00:00 2001 From: Saswat Susmoy Date: Thu, 18 Jun 2026 01:21:23 +0530 Subject: [PATCH 2/8] =?UTF-8?q?refactor(stub):=20rename=20CTESN=20?= =?UTF-8?q?=E2=86=92=20ContinuousESN=20per=20PR3=20design=20discussion?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Francesco's call on #456: drop the CTESN name (which referred to the Anantharaman 2021 parametric-surrogate paper) and ship instead a thin wrapper around the Lukoševičius 2012 §3.2.6 eq (5) leaky-integrator continuous ESN. Stub file moves from `src/extensions/` to `src/models/` to sit alongside ESN as a semantic sibling; docstring rewritten to reflect the new scope (leaky-integrator default, parametric-surrogate CTESN explicitly out of scope). --- src/{extensions/ctesn.jl => models/continuous_esn.jl} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename src/{extensions/ctesn.jl => models/continuous_esn.jl} (100%) diff --git a/src/extensions/ctesn.jl b/src/models/continuous_esn.jl similarity index 100% rename from src/extensions/ctesn.jl rename to src/models/continuous_esn.jl From 63249085c2d0721f88390039f217ec1e8afae5fb Mon Sep 17 00:00:00 2001 From: Saswat Susmoy Date: Thu, 18 Jun 2026 01:22:29 +0530 Subject: [PATCH 3/8] docs(continuous-esn): rewrite stub docstring for ContinuousESN scope MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Prior commit captured only the file rename — the docstring content update was on disk but not staged. This commit folds in the actual docstring rewrite: references Lukoševičius 2012 §3.2.6 eq (5) as the canonical ODE, explains how the discrete leak rate α maps to the continuous Euler step, and explicitly distinguishes from the Anantharaman 2021 parametric-surrogate CTESN. --- src/models/continuous_esn.jl | 55 ++++++++++++++++++++---------------- 1 file changed, 31 insertions(+), 24 deletions(-) diff --git a/src/models/continuous_esn.jl b/src/models/continuous_esn.jl index 3660e8dad..a47fecaaa 100644 --- a/src/models/continuous_esn.jl +++ b/src/models/continuous_esn.jl @@ -1,26 +1,33 @@ @doc raw""" - CTESN(in_dims, res_dims, out_dims; kwargs...) + ContinuousESN(in_dims, res_dims, out_dims; kwargs...) -Continuous-Time Echo State Network ([Anantharaman2021](@cite)). **Work in progress -— PR3 of the SciML 2026 fellowship roadmap (#397).** This is a stub: the real -constructor lives in the `RCODEReservoirExt` package extension and will be -filled in once the design questions on the open pull request are resolved. +Continuous-time leaky-integrator Echo State Network ([Lukosevicius2012](@cite) +§3.2.6, eq 5). **Work in progress — PR3 of the SciML 2026 fellowship roadmap +(#397).** This is a stub; the real constructor lives in the `RCODEReservoirExt` +package extension and will be filled in once the design is settled. -`CTESN <: AbstractSciMLProblemReservoir`, so it reuses the continuous-time -`_collectstates` / `_predict` machinery landed in PR #450 (`add-ode-reservoir-ext`). -It specialises that machinery to a concrete reservoir ODE. - -The canonical reservoir ODE from [Anantharaman2021](@cite) eq. (3) is: +`ContinuousESN <: AbstractSciMLProblemReservoir`, so it reuses the +continuous-time `_collectstates` / `_predict` machinery landed in PR #450 +(`add-ode-reservoir-ext`) and pre-bakes the leaky-integrator reservoir ODE +of [Lukosevicius2012](@cite) eq (5): ```math -\dot{\mathbf{r}}(t) = \tanh\!\left(A\,\mathbf{r}(t) + W_{\text{hyb}}\,\mathbf{x}(t)\right) +\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) ``` -with no leak term and no bias — `f = tanh`, `g = id` hardcoded by the paper. -The implementation may also expose an opt-in leaky form -`ṙ = -α r + tanh(W_in u(t) + W_r r + b)` (the continuous limit of the -discrete leaky ESN of [Lukosevicius2012](@cite)); see the PR thread for the -default-form decision. +The leak rate `α` familiar from the discrete leaky ESN is, in the continuous +formulation, exactly the Euler discretisation step: take `Δt = α` and the +update collapses to `x(n+1) = (1-α)x(n) + α·tanh(W_in u(n+1) + W_r x(n) + b)`. +It is therefore controlled in `ContinuousESN` via the integration `tspan` and +the input-window count, not as a separate parameter of the ODE. + +This is *not* a port of the parametric-surrogate CTESN of +[Anantharaman2021](@cite); that model uses +`\dot{\mathbf{r}} = \tanh(A\,\mathbf{r} + W_{\text{hyb}}\,\mathbf{x}(t))` without +a decay term and is trained via RBF interpolation of `W_out(p)` over a +parameter space. It would land as a separate type in a future PR. !!! note This constructor errors unless the `RCODEReservoirExt` extension is loaded @@ -29,13 +36,13 @@ default-form decision. ## Status -- PR3 work in progress on branch `add-ctesn`, stacked on top of PR #450 - (`add-ode-reservoir-ext`). Will be rebased onto `master` once PR #450 merges. -- Open design questions are in the pull-request description. +- PR3 work in progress on branch `add-ctesn`, stacked on top of PR #450. + Will be rebased onto `master` once #450 merges. +- Open design questions tracked in the pull-request description. """ -CTESN(::Any...) = error( - "CTESN requires the RCODEReservoirExt extension and an OrdinaryDiffEq " * - "solver package. Load `SciMLBase`, `DataInterpolations`, and a solver " * - "package (e.g. `OrdinaryDiffEqTsit5`) to enable it. PR3 is a " * - "work-in-progress — the real constructor is not yet wired up." +ContinuousESN(::Any...) = error( + "ContinuousESN requires the RCODEReservoirExt extension and an " * + "OrdinaryDiffEq solver package. Load `SciMLBase`, `DataInterpolations`, " * + "and a solver package (e.g. `OrdinaryDiffEqTsit5`) to enable it. PR3 is " * + "a work-in-progress — the real constructor is not yet wired up." ) From db4cac9913128a938814f224650b4f63afb3fe16 Mon Sep 17 00:00:00 2001 From: Saswat Susmoy Date: Tue, 23 Jun 2026 03:35:30 +0530 Subject: [PATCH 4/8] =?UTF-8?q?feat(models):=20implement=20ContinuousESN?= =?UTF-8?q?=20=E2=80=94=20leaky-integrator=20continuous=20ESN?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `ContinuousESN <: AbstractSciMLProblemReservoir`. Thin wrapper that pre-bakes the Lukoševičius 2012 §3.2.6 eq (5) ODE ẋ(t) = α · (-x(t) + tanh(W_in·u(t) + W_r·x(t) + b)) with `leak_coefficient = α`. Forward-Euler discretisation at Δt = 1 collapses to the discrete leaky ESN exactly. Struct lives in src/models/, real constructor + in-place RHS (mul!, fused tanh broadcast, compile-time bias branch) in `RCODEReservoirExt`. Validates `tspan`, dim positivity, leak positivity, finite endpoints, length-2 tspan, and Inf-on-cast at setup. Test suite covers shape, construction errors, bias toggle, bias under Tsit5, forward determinism, Euler equivalence vs discrete leaky ESN at α ∈ {1.0, 0.3}, both predict modes, modifier composition, and T-kwarg propagation — 43 assertions across 10 testsets. Tutorial added under docs/src/tutorials/continuous_esn.md (Lorenz forecasting walkthrough). Adds Lukoševičius 2012 and Anantharaman 2021 to refs.bib. Drops the stale PR3/PR5 stub comment in PR2's sciml_reservoir.jl now that the first concrete subtype has landed. --- docs/pages.jl | 1 + docs/src/api/layers.md | 1 + docs/src/refs.bib | 20 ++ docs/src/tutorials/continuous_esn.md | 151 +++++++++++++++ ext/RCODEReservoirExt.jl | 107 ++++++++++- src/ReservoirComputing.jl | 3 +- src/layers/sciml_reservoir.jl | 5 +- src/models/continuous_esn.jl | 152 +++++++++++---- test/test_continuous_esn.jl | 275 +++++++++++++++++++++++++++ 9 files changed, 676 insertions(+), 39 deletions(-) create mode 100644 docs/src/tutorials/continuous_esn.md create mode 100644 test/test_continuous_esn.jl diff --git a/docs/pages.jl b/docs/pages.jl index a15e17918..10fd36500 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 ddd1fb3f3..c446b1ab4 100644 --- a/docs/src/api/layers.md +++ b/docs/src/api/layers.md @@ -47,6 +47,7 @@ ```@docs AbstractSciMLProblemReservoir SciMLProblemReservoir + ContinuousESN AbstractSampler TerminalStateSampling ``` diff --git a/docs/src/refs.bib b/docs/src/refs.bib index 0f3946ab0..f895bf0a8 100644 --- a/docs/src/refs.bib +++ b/docs/src/refs.bib @@ -557,3 +557,23 @@ @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} +} + +@article{Anantharaman2021, + title = {Accelerating Simulation of Stiff Nonlinear Systems using Continuous-Time Echo State Networks}, + url = {https://arxiv.org/abs/2010.04004}, + author = {Anantharaman, Ranjan and Ma, Yingbo and Gowda, Shashi and Laughman, Christopher and Shah, Viral and Edelman, Alan and Rackauckas, Christopher}, + year = {2021}, + journal = {arXiv:2010.04004} +} diff --git a/docs/src/tutorials/continuous_esn.md b/docs/src/tutorials/continuous_esn.md new file mode 100644 index 000000000..5b725dfda --- /dev/null +++ b/docs/src/tutorials/continuous_esn.md @@ -0,0 +1,151 @@ +# Continuous ESN: forecasting Lorenz + +[`ContinuousESN`](@ref) is a thin wrapper around +[`SciMLProblemReservoir`](@ref) that pre-bakes the leaky-integrator +continuous Echo State Network ODE of [Lukosevicius2012](@cite) §3.2.6 +eq (5): + +```math +\dot{\mathbf{x}}(t) = \alpha \left( + -\mathbf{x}(t) + \tanh\!\left( + \mathbf{W}_{\text{in}}\,\mathbf{u}(t) + \mathbf{W}_r\,\mathbf{x}(t) + + \mathbf{b} + \right) +\right) +``` + +with leaking rate `α`. Forward-Euler discretisation at step `Δt = 1` +recovers the discrete leaky ESN `x(n+1) = (1-α) x(n) + α tanh(...)` +exactly. The reservoir matrices `W_r`, `W_in`, and optional bias `b` +live in `ps.reservoir` and are constructed by `setup(rng, rc)`; the +ODE solver and any solve-time keyword arguments are captured at +construction. Under the hood the same `RCODEReservoirExt` extension +that powers `SciMLProblemReservoir` also runs `ContinuousESN`. + +This tutorial walks through training a `ContinuousESN` on Lorenz-63 +data and rolling it forward autoregressively to reproduce the +attractor. + +## Loading the extension + +```julia +using ReservoirComputing +using SciMLBase +using DataInterpolations +using OrdinaryDiffEqTsit5 +``` + +`SciMLBase` provides `solve` / `remake`, `DataInterpolations` is used +internally for the per-window input signal in autoregressive mode, and +`OrdinaryDiffEqTsit5` supplies the concrete `Tsit5()` solver type. + +## 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` + +The constructor mirrors `SciMLProblemReservoir`: the integration +`tspan` and the solver positional argument are captured at +construction time, exactly as in DiffEqFlux's `NeuralODE`. + +```@example continuous-esn-lorenz +N_res = 100 +ce_train = ContinuousESN( + 3, N_res, (0.0, Float64(train_len)), Tsit5(); + leak_coefficient = 1.0, use_bias = true, T = Float64, + reltol = 1.0e-6, abstol = 1.0e-8 +) +ce_pred = ContinuousESN( + 3, N_res, (0.0, Float64(predict_len)), Tsit5(); + leak_coefficient = 1.0, use_bias = true, T = Float64, + reltol = 1.0e-6, abstol = 1.0e-8 +) + +rc_train = ReservoirComputer(ce_train, (NLAT2(),), LinearReadout(N_res => 3)) +rc_pred = ReservoirComputer(ce_pred, (NLAT2(),), LinearReadout(N_res => 3)) + +ps, st = setup(rng, rc_train) +``` + +## Training + +`train!` is sampler-agnostic: it routes through the continuous +`_collectstates` provided by the extension and fits a linear readout +on the collected states. + +```@example continuous-esn-lorenz +ps, st = train!(rc_train, input_data, target_data, ps, st) +``` + +## Autoregressive rollout + +`predict(rc, steps, ps, st; initialdata)` splits the predict-time +`tspan` into `steps` equal sub-intervals; on each sub-interval the +previous readout output is held constant as the input. The default +initial reservoir state is `prob.u0` (zeros) — if you want to continue +from the trained reservoir's terminal state, `remake` the prob before +constructing the predict reservoir. + +```@example continuous-esn-lorenz +ps_pred, st_pred = setup(rng, rc_pred) +ps_pred = merge(ps_pred, (readout = ps.readout,)) +st_pred = merge(st_pred, (readout = st.readout,)) + +output, _ = predict( + rc_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 eq (5); reach for it when the standard + leaky-integrator continuous ESN is what you want and you'd otherwise + be hand-rolling the same RHS. +* [`SciMLProblemReservoir`](@ref) is the generic building block; reach + for it when the reservoir ODE is *not* eq (5) — bespoke RHS, SDE, + DDE, or non-standard parameter layout. + +Both go through the same `_collectstates` / `_predict` dispatch and +the same protected-`saveat` discipline. diff --git a/ext/RCODEReservoirExt.jl b/ext/RCODEReservoirExt.jl index 661f21cfa..c8acbf9eb 100644 --- a/ext/RCODEReservoirExt.jl +++ b/ext/RCODEReservoirExt.jl @@ -1,21 +1,26 @@ module RCODEReservoirExt using DataInterpolations: ConstantInterpolation +using LinearAlgebra: mul! using LuxCore: apply +using Random: AbstractRNG # `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 SciMLBase: ODEProblem, remake, solve, NullParameters using ReservoirComputing: ReservoirComputing, AbstractReservoirComputer, AbstractSampler, AbstractSciMLProblemReservoir, + ContinuousESN, TerminalStateSampling, - collectstates + collectstates, + rand_sparse, + scaled_rand import ReservoirComputing: _collectstates, _predict # --------------------------------------------------------------------------- @@ -389,4 +394,102 @@ function _predict( return outputs, newst end +# --------------------------------------------------------------------------- +# ContinuousESN +# +# Concrete subtype of `AbstractSciMLProblemReservoir` that pre-bakes the +# leaky-integrator continuous ESN ODE of Lukoševičius 2012 §3.2.6 eq (5): +# +# ẋ(t) = α · (-x(t) + tanh(W_in·u(t) + W_r·x(t) + b)) +# +# `W_r`, `W_in`, `b`, and the scalar `leak_coefficient = α` live in +# `ps.reservoir` and are injected into the solve `p` by +# `_build_solve_params`. The RHS is in-place (`mul!` for both matvecs, +# fused tanh broadcast) so it stays allocation-free on the solver hot +# path. +# --------------------------------------------------------------------------- + +function _continuous_esn_rhs!(dx, x, p, t) + # dx ← W_r · x + mul!(dx, p.W_r, x) + # dx ← dx + W_in · u(t) + u_t = p.input(t) + mul!(dx, p.W_in, u_t, true, true) + # dx ← α · (-x + tanh(dx + b?)) + # + # `haskey` on a NamedTuple resolves at compile time (the keys are + # part of the type), so the two branches don't add runtime cost. + if haskey(p, :b) + @. dx = p.leak_coefficient * (-x + tanh(dx + p.b)) + else + @. dx = p.leak_coefficient * (-x + tanh(dx)) + end + return nothing +end + +# Type-aware default bias initialiser — `(rng, T, dims...) -> zeros(T, dims...)`. +# `WeightInitializers.zeros32` is Float32-only and ignores `T`, so we ship a +# small `T`-respecting default and document the expected signature in the +# constructor docstring. +_typed_zeros_init(::AbstractRNG, ::Type{T}, dims::Integer...) where {T} = + zeros(T, dims...) + +function ReservoirComputing.ContinuousESN( + in_dims::Integer, res_dims::Integer, tspan, args...; + leak_coefficient::Real = 1.0, + use_bias::Bool = false, + init_reservoir = rand_sparse, + init_input = scaled_rand, + init_bias = _typed_zeros_init, + T::Type{<:AbstractFloat} = Float32, + 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")) + leak_coefficient > 0 || throw( + ArgumentError("leak_coefficient must be positive, got $leak_coefficient") + ) + 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( + "SciMLProblemReservoir requires `tspan[2] > tspan[1]`, got tspan = $tspan" + ) + ) + ReservoirComputing._check_protected_kwargs(kwargs) + + # Zero initial reservoir state — Lukoševičius 2012 starts from x(0) = 0 + # by convention. Users continuing from a previously trained terminal + # state should `remake(rc.reservoir.prob; u0 = ...)` themselves. + u0 = zeros(T, res_dims) + + # `prob.p = nothing` — all reservoir parameters live in `ps.reservoir` + # and get merged into the solve `p` by `_build_solve_params`. + prob = ODEProblem(_continuous_esn_rhs!, u0, tspan) + + return ContinuousESN( + prob, + TerminalStateSampling(), + tspan, + args, + kwargs, + in_dims, + res_dims, + leak_coefficient, + use_bias, + init_reservoir, + init_input, + init_bias, + T + ) +end + end # module diff --git a/src/ReservoirComputing.jl b/src/ReservoirComputing.jl index 5d4e37318..20d348d8d 100644 --- a/src/ReservoirComputing.jl +++ b/src/ReservoirComputing.jl @@ -59,11 +59,12 @@ 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 ESNCell, ES2NCell, EuSNCell, MemoryESNCell, MemoryResESNCell, RMNCell export StatefulLayer, LinearReadout, ReservoirChain, Collect, collectstates, diff --git a/src/layers/sciml_reservoir.jl b/src/layers/sciml_reservoir.jl index 1558588e2..2800961c3 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 index a47fecaaa..d833b9be9 100644 --- a/src/models/continuous_esn.jl +++ b/src/models/continuous_esn.jl @@ -1,48 +1,134 @@ @doc raw""" - ContinuousESN(in_dims, res_dims, out_dims; kwargs...) + ContinuousESN(in_dims, res_dims, tspan, args...; kwargs...) Continuous-time leaky-integrator Echo State Network ([Lukosevicius2012](@cite) -§3.2.6, eq 5). **Work in progress — PR3 of the SciML 2026 fellowship roadmap -(#397).** This is a stub; the real constructor lives in the `RCODEReservoirExt` -package extension and will be filled in once the design is settled. - -`ContinuousESN <: AbstractSciMLProblemReservoir`, so it reuses the -continuous-time `_collectstates` / `_predict` machinery landed in PR #450 -(`add-ode-reservoir-ext`) and pre-bakes the leaky-integrator reservoir ODE -of [Lukosevicius2012](@cite) eq (5): +§3.2.6, eq 5). `ContinuousESN <: AbstractSciMLProblemReservoir`, so it reuses +the continuous-time `_collectstates` / `_predict` machinery in the +`RCODEReservoirExt` package extension and pre-bakes the leaky-integrator +reservoir 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} +\dot{\mathbf{x}}(t) = \alpha \left( + -\mathbf{x}(t) + \tanh\!\left( + \mathbf{W}_{\text{in}}\,\mathbf{u}(t) + \mathbf{W}_r\,\mathbf{x}(t) + \mathbf{b} + \right) \right) ``` -The leak rate `α` familiar from the discrete leaky ESN is, in the continuous -formulation, exactly the Euler discretisation step: take `Δt = α` and the -update collapses to `x(n+1) = (1-α)x(n) + α·tanh(W_in u(n+1) + W_r x(n) + b)`. -It is therefore controlled in `ContinuousESN` via the integration `tspan` and -the input-window count, not as a separate parameter of the ODE. +with reservoir matrix `W_r`, input matrix `W_in`, optional bias `b`, and +scalar leaking rate `α` exposed as `leak_coefficient`. Forward-Euler +discretisation at step `Δt = 1` collapses the ODE to the discrete leaky +update `x(n+1) = (1-α)·x(n) + α·tanh(W_in·u(n) + W_r·x(n) + b)`. This is *not* a port of the parametric-surrogate CTESN of -[Anantharaman2021](@cite); that model uses -`\dot{\mathbf{r}} = \tanh(A\,\mathbf{r} + W_{\text{hyb}}\,\mathbf{x}(t))` without +[Anantharaman2021](@cite); that model uses `ṙ = tanh(A·r + W_hyb·u)` without a decay term and is trained via RBF interpolation of `W_out(p)` over a -parameter space. It would land as a separate type in a future PR. +parameter space. A separate type would land for it in a future PR. -!!! note - This constructor errors unless the `RCODEReservoirExt` extension is loaded - (`SciMLBase` + `DataInterpolations`) **and** a concrete `OrdinaryDiffEq` - solver package is available for `args[1]`. +## Arguments + + - `in_dims`: Input dimension. + - `res_dims`: Reservoir (hidden state) dimension. + - `tspan`: Integration interval for `collectstates`. Must be a length-2, + strictly-increasing, finite tuple/pair. Mirrors the DiffEqFlux + `NeuralODE` convention. + - `args...`: Positional arguments forwarded to `solve`. The solver + algorithm (e.g. `Tsit5()`, `Euler()`) is the first element by + convention. + +## Keyword arguments + + - `leak_coefficient`: Leaking rate `α`. Default: `1.0`. Pairing `Δt = 1` + with `leak_coefficient = α` recovers the discrete leaky ESN under + forward Euler. + - `use_bias`: Whether to include a bias term `b`. Default: `false`. + - `init_reservoir`: Initializer for `W_r`. Default: + [`rand_sparse`](@ref). Receives `(rng, T, res_dims, res_dims)`. + - `init_input`: Initializer for `W_in`. Default: [`scaled_rand`](@ref). + Receives `(rng, T, res_dims, in_dims)`. + - `init_bias`: Initializer for `b`. Used only when `use_bias = true`. + Default: a type-aware zeros initialiser. Receives `(rng, T, res_dims)`. + - `T`: Element type for the reservoir matrices and initial state. + Default: `Float32`. + - `kwargs...`: Forwarded to `solve`. The three keys `saveat`, + `save_everystep`, and `dense` are owned by the extension helper and + rejected at construction. + +## Continuous-time stability note -## Status +`scale_radius!` and the default reservoir initialisers control the +discrete-time spectral radius `ρ(W_r)`. For the α-factored ODE used here, +a log-norm contraction argument gives the sufficient condition +`‖W_r‖_2 < 1` (operator 2-norm bound), independent of `α`: the leaking +rate scales the rate of decay but not the asymptotic bound. The condition +is strictly stronger than `ρ(W_r) < 1` because `‖·‖_2 ≥ ρ(·)`. Published +continuous ESN work treats `ρ` as an empirical hyperparameter — adjust +accordingly. -- PR3 work in progress on branch `add-ctesn`, stacked on top of PR #450. - Will be rebased onto `master` once #450 merges. -- Open design questions tracked in the pull-request description. +!!! note + This constructor errors unless `RCODEReservoirExt` is loaded. Load + `SciMLBase`, `DataInterpolations`, and a concrete OrdinaryDiffEq + solver package (e.g. `OrdinaryDiffEqTsit5`, `OrdinaryDiffEq`) to + enable it. """ -ContinuousESN(::Any...) = error( - "ContinuousESN requires the RCODEReservoirExt extension and an " * - "OrdinaryDiffEq solver package. Load `SciMLBase`, `DataInterpolations`, " * - "and a solver package (e.g. `OrdinaryDiffEqTsit5`) to enable it. PR3 is " * - "a work-in-progress — the real constructor is not yet wired up." -) +@concrete struct ContinuousESN <: AbstractSciMLProblemReservoir + prob + sampler + tspan + args + kwargs + in_dims + res_dims + leak_coefficient + use_bias + init_reservoir + init_input + init_bias + matrix_type +end + +# Public factory — the real implementation lives in `RCODEReservoirExt`. +# Loading `SciMLBase` + `DataInterpolations` adds the typed method that +# actually builds the `ODEProblem`; without the extension, this fallback +# fires and produces a clear error. +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, ce::ContinuousESN) + print(io, "ContinuousESN(") + print(io, "in_dims = ", ce.in_dims) + print(io, ", res_dims = ", ce.res_dims) + print(io, ", leak_coefficient = ", ce.leak_coefficient) + print(io, ", use_bias = ", ce.use_bias) + print(io, ", tspan = ") + show(io, ce.tspan) + print(io, ")") + return +end + +function initialparameters(rng::AbstractRNG, ce::ContinuousESN) + T = ce.matrix_type + W_r = ce.init_reservoir(rng, T, ce.res_dims, ce.res_dims) + W_in = ce.init_input(rng, T, ce.res_dims, ce.in_dims) + leak_coefficient = T(ce.leak_coefficient) + isfinite(leak_coefficient) || throw( + ArgumentError( + "leak_coefficient converted to $T overflowed to $leak_coefficient; " * + "pick a smaller value or widen `T`." + ) + ) + ps = (W_r = W_r, W_in = W_in, leak_coefficient = leak_coefficient) + if ce.use_bias + ps = merge(ps, (b = ce.init_bias(rng, T, ce.res_dims),)) + end + return ps +end + +function initialstates(::AbstractRNG, ::ContinuousESN) + return NamedTuple() +end diff --git a/test/test_continuous_esn.jl b/test/test_continuous_esn.jl new file mode 100644 index 000000000..de59d717b --- /dev/null +++ b/test/test_continuous_esn.jl @@ -0,0 +1,275 @@ +using Test +using Random +using LinearAlgebra +using Statistics +using ReservoirComputing +using OrdinaryDiffEq +using SciMLBase +using DataInterpolations + +# --------------------------------------------------------------------------- +# 1. Construction: shape + parameter wiring +# +# `ContinuousESN(in_dims, res_dims, tspan, solver)` should build a reservoir +# whose `initialparameters` yields the canonical eq-(5) trio (`W_r`, +# `W_in`, `inv_leak`) at the correct shapes, with no bias by default. +# --------------------------------------------------------------------------- + +@testset "ContinuousESN: construction + parameter shapes" begin + rng = MersenneTwister(0) + in_dim, res_dim = 3, 50 + ce = ContinuousESN(in_dim, res_dim, (0.0, 5.0), Tsit5()) + + rc = ReservoirComputer(ce, LinearReadout(res_dim => 2)) + ps, st = setup(rng, rc) + + @test ce isa AbstractSciMLProblemReservoir + @test size(ps.reservoir.W_r) == (res_dim, res_dim) + @test size(ps.reservoir.W_in) == (res_dim, in_dim) + @test ps.reservoir.leak_coefficient == 1.0f0 + @test !haskey(ps.reservoir, :b) + @test all(isfinite, ps.reservoir.W_r) + @test all(isfinite, ps.reservoir.W_in) +end + +# --------------------------------------------------------------------------- +# 2. Construction validation: argument errors +# +# Positive `in_dims`, `res_dims`, `leak_coefficient`, and strictly +# increasing `tspan` are required. Protected `solve` kwargs are rejected +# at construction (inherits `SciMLProblemReservoir`'s _PROTECTED_SOLVE_KWARGS). +# --------------------------------------------------------------------------- + +@testset "ContinuousESN: construction validation" begin + @test_throws ArgumentError ContinuousESN(0, 5, (0.0, 1.0), Tsit5()) + @test_throws ArgumentError ContinuousESN(3, 0, (0.0, 1.0), Tsit5()) + @test_throws ArgumentError ContinuousESN(3, 5, (0.0, 0.0), Tsit5()) + @test_throws ArgumentError ContinuousESN(3, 5, (1.0, 0.0), Tsit5()) + @test_throws ArgumentError ContinuousESN( + 3, 5, (0.0, 1.0), Tsit5(); leak_coefficient = 0.0 + ) + # Non-finite endpoints + @test_throws ArgumentError ContinuousESN(3, 5, (0.0, Inf), Tsit5()) + @test_throws ArgumentError ContinuousESN(3, 5, (-Inf, 1.0), Tsit5()) + @test_throws ArgumentError ContinuousESN(3, 5, (0.0, NaN), Tsit5()) + # Wrong-length tspan + @test_throws ArgumentError ContinuousESN(3, 5, (1.0,), Tsit5()) + @test_throws ArgumentError ContinuousESN(3, 5, (0.0, 1.0, 2.0), Tsit5()) + # leak_coefficient that overflows `T` is caught at `setup` time, not + # construction (the construction check sees the unconverted value). + # Use a dense init to avoid `rand_sparse`'s NaN-on-tiny-matrix edge. + let dense_init = (rng, T, d...) -> rand_sparse(rng, T, d...; sparsity = 0.5) + ce_overflow = ContinuousESN( + 3, 16, (0.0, 1.0), Tsit5(); + leak_coefficient = 1.0e50, T = Float32, init_reservoir = dense_init + ) + rc_overflow = ReservoirComputer(ce_overflow, LinearReadout(16 => 1)) + @test_throws ArgumentError setup(MersenneTwister(0), rc_overflow) + end + for badkw in (:saveat, :save_everystep, :dense) + @test_throws ArgumentError ContinuousESN( + 3, 5, (0.0, 1.0), Tsit5(); (badkw => true,)... + ) + end +end + +# --------------------------------------------------------------------------- +# 3. Bias toggle +# +# `use_bias = true` should add a `b` parameter of length `res_dims`. +# --------------------------------------------------------------------------- + +@testset "ContinuousESN: bias toggle" begin + rng = MersenneTwister(1) + in_dim, res_dim = 2, 10 + ce_b = ContinuousESN(in_dim, res_dim, (0.0, 1.0), Tsit5(); use_bias = true) + rc = ReservoirComputer(ce_b, LinearReadout(res_dim => 1)) + ps, st = setup(rng, rc) + @test haskey(ps.reservoir, :b) + @test length(ps.reservoir.b) == res_dim + # Default init is the type-aware zeros initialiser. + @test all(==(0), ps.reservoir.b) +end + +# --------------------------------------------------------------------------- +# 3b. Bias actually exercised under an adaptive solver +# +# The `haskey(:b)` branch in `_continuous_esn_rhs!` is solver-agnostic in +# theory, but only the Euler-equivalence test (which uses `Euler`) hits +# the bias-on path under integration. Pair `use_bias = true` with a +# nonzero bias initialiser and `Tsit5` to confirm: (a) the bias does +# perturb states relative to the no-bias baseline, and (b) the result +# stays finite. +# --------------------------------------------------------------------------- + +@testset "ContinuousESN: bias path under Tsit5" begin + rng = MersenneTwister(101) + in_dim, res_dim, T_steps = 2, 20, 12 + nonzero_bias(rng, T, d...) = T(0.5) .* randn(rng, T, d...) + ce_with_bias = ContinuousESN( + in_dim, res_dim, (0.0, 2.0), Tsit5(); + use_bias = true, init_bias = nonzero_bias, + reltol = 1.0e-8, abstol = 1.0e-10 + ) + ce_no_bias = ContinuousESN( + in_dim, res_dim, (0.0, 2.0), Tsit5(); + use_bias = false, reltol = 1.0e-8, abstol = 1.0e-10 + ) + rc_b = ReservoirComputer(ce_with_bias, LinearReadout(res_dim => 1)) + rc_n = ReservoirComputer(ce_no_bias, LinearReadout(res_dim => 1)) + ps_b, st_b = setup(MersenneTwister(0), rc_b) + ps_n, st_n = setup(MersenneTwister(0), rc_n) + data = randn(Float32, in_dim, T_steps) + + s_b, _ = collectstates(rc_b, data, ps_b, st_b) + s_n, _ = collectstates(rc_n, data, ps_n, st_n) + @test all(isfinite, s_b) + @test s_b != s_n +end + +# --------------------------------------------------------------------------- +# 4. Forward pass: shape + finiteness + determinism +# --------------------------------------------------------------------------- + +@testset "ContinuousESN: forward (collectstates)" begin + rng = MersenneTwister(7) + in_dim, res_dim, T_steps = 2, 16, 25 + ce = ContinuousESN( + in_dim, res_dim, (0.0, 3.0), Tsit5(); + reltol = 1.0e-8, abstol = 1.0e-10 + ) + rc = ReservoirComputer(ce, LinearReadout(res_dim => 1)) + ps, st = setup(rng, rc) + data = randn(Float32, in_dim, T_steps) + + s1, _ = collectstates(rc, data, ps, st) + s2, _ = collectstates(rc, data, ps, st) + @test size(s1) == (res_dim, T_steps) + @test all(isfinite, s1) + @test s1 ≈ s2 +end + +# --------------------------------------------------------------------------- +# 5. Euler equivalence with discrete leaky ESN +# +# Solving `ẋ = α(-x + tanh(W_r·x + W_in·u(t) + b))` with explicit Euler at +# step `Δt = 1` and `leak_coefficient = α` collapses algebraically to the +# discrete leaky ESN update `x_{k+1} = (1-α)·x_k + α·tanh(...)`. With the +# corrected window alignment (input at window start, sample at window +# end), continuous and discrete trajectories agree exactly. +# +# Covers both vanilla (α = 1) and leaky (α = 0.3) regimes. +# --------------------------------------------------------------------------- + +@testset "ContinuousESN: Euler equivalence with discrete leaky ESN" begin + # Dense reservoir init avoids `rand_sparse`'s spectral-scaling edge case + # on very small matrices (8×8 at 10% density can yield a singular matrix + # whose `eigvals` are NaN, tripping `check_inf_nan`). + dense_init(rng, T, d...) = rand_sparse(rng, T, d...; sparsity = 0.5) + for α in (1.0, 0.3) + rng = MersenneTwister(13) + in_dim, res_dim, T_steps = 2, 16, 12 + + ce = ContinuousESN( + in_dim, res_dim, (0.0, Float64(T_steps)), Euler(); + leak_coefficient = α, use_bias = true, T = Float64, dt = 1.0, + init_reservoir = dense_init + ) + rc = ReservoirComputer(ce, LinearReadout(res_dim => 1)) + ps, st = setup(MersenneTwister(0), rc) + + data = randn(rng, in_dim, T_steps) + cont_states, _ = collectstates(rc, data, ps, st) + + W_r = ps.reservoir.W_r + W_in = ps.reservoir.W_in + b = ps.reservoir.b + 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 + +# --------------------------------------------------------------------------- +# 6. Teacher-forced predict + autoregressive predict shape/determinism +# --------------------------------------------------------------------------- + +@testset "ContinuousESN: teacher-forced predict" begin + rng = MersenneTwister(21) + in_dim, res_dim, out_dim, T_steps = 2, 12, 3, 10 + ce = ContinuousESN( + in_dim, res_dim, (0.0, 2.0), Tsit5(); + reltol = 1.0e-8, abstol = 1.0e-10 + ) + rc = ReservoirComputer(ce, LinearReadout(res_dim => out_dim)) + ps, st = setup(rng, rc) + + data = randn(Float32, in_dim, T_steps) + tf1, _ = predict(rc, data, ps, st) + tf2, _ = predict(rc, 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 + ce = ContinuousESN( + dim, res_dim, (0.0, 1.0), Tsit5(); + reltol = 1.0e-8, abstol = 1.0e-10 + ) + rc = ReservoirComputer(ce, LinearReadout(res_dim => dim)) + ps, st = setup(rng, rc) + + init = randn(Float32, dim) + ar1, _ = predict(rc, steps, ps, st; initialdata = init) + ar2, _ = predict(rc, steps, ps, st; initialdata = init) + @test size(ar1) == (dim, steps) + @test all(isfinite, ar1) + @test ar1 ≈ ar2 +end + +# --------------------------------------------------------------------------- +# 7. Compatibility with state modifiers +# --------------------------------------------------------------------------- + +@testset "ContinuousESN: state modifiers compose" begin + rng = MersenneTwister(31) + in_dim, res_dim, T_steps = 2, 10, 8 + ce = ContinuousESN( + in_dim, res_dim, (0.0, 1.0), Tsit5(); + reltol = 1.0e-8, abstol = 1.0e-10 + ) + rc_plain = ReservoirComputer(ce, LinearReadout(res_dim => 1)) + rc_mod = ReservoirComputer(ce, (NLAT2(),), LinearReadout(res_dim => 1)) + + ps_p, st_p = setup(MersenneTwister(0), rc_plain) + ps_m, st_m = setup(MersenneTwister(0), rc_mod) + data = randn(Float32, in_dim, T_steps) + + sp, _ = collectstates(rc_plain, data, ps_p, st_p) + sm, _ = collectstates(rc_mod, data, ps_m, st_m) + @test size(sm) == size(sp) + @test all(isfinite, sm) + @test sm != sp +end + +# --------------------------------------------------------------------------- +# 8. Eltype propagates through `T` kwarg +# --------------------------------------------------------------------------- + +@testset "ContinuousESN: T kwarg controls matrix eltype" begin + rng = MersenneTwister(43) + ce64 = ContinuousESN(2, 32, (0.0, 1.0), Tsit5(); T = Float64) + rc = ReservoirComputer(ce64, LinearReadout(32 => 1)) + ps, _ = setup(rng, rc) + @test eltype(ps.reservoir.W_r) == Float64 + @test eltype(ps.reservoir.W_in) == Float64 + @test typeof(ps.reservoir.leak_coefficient) == Float64 +end From 7a7a84b5a4c63db64de4fdbf869e2ae9c2688903 Mon Sep 17 00:00:00 2001 From: Saswat Susmoy Date: Thu, 25 Jun 2026 23:22:26 +0530 Subject: [PATCH 5/8] docs(continuous-esn): address review on PR #456 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Move `ContinuousESN` from layers.md to models.md — it is a model (sibling of `ESN`/`EuSN`), not a continuous-time reservoir building block; the `SciMLProblemReservoir` family stays under layers. - Drop the standalone "Loading the extension" block in the tutorial; the `@example continuous-esn-lorenz` block immediately below already carries the `using` lines, so the prose block was redundant. --- docs/src/api/layers.md | 1 - docs/src/api/models.md | 6 ++++++ docs/src/tutorials/continuous_esn.md | 18 ++++-------------- 3 files changed, 10 insertions(+), 15 deletions(-) diff --git a/docs/src/api/layers.md b/docs/src/api/layers.md index c446b1ab4..ddd1fb3f3 100644 --- a/docs/src/api/layers.md +++ b/docs/src/api/layers.md @@ -47,7 +47,6 @@ ```@docs AbstractSciMLProblemReservoir SciMLProblemReservoir - ContinuousESN AbstractSampler TerminalStateSampling ``` diff --git a/docs/src/api/models.md b/docs/src/api/models.md index 97b947b5e..b306c81c8 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/tutorials/continuous_esn.md b/docs/src/tutorials/continuous_esn.md index 5b725dfda..b00418b4a 100644 --- a/docs/src/tutorials/continuous_esn.md +++ b/docs/src/tutorials/continuous_esn.md @@ -24,20 +24,10 @@ that powers `SciMLProblemReservoir` also runs `ContinuousESN`. This tutorial walks through training a `ContinuousESN` on Lorenz-63 data and rolling it forward autoregressively to reproduce the -attractor. - -## Loading the extension - -```julia -using ReservoirComputing -using SciMLBase -using DataInterpolations -using OrdinaryDiffEqTsit5 -``` - -`SciMLBase` provides `solve` / `remake`, `DataInterpolations` is used -internally for the per-window input signal in autoregressive mode, and -`OrdinaryDiffEqTsit5` supplies the concrete `Tsit5()` solver type. +attractor. `SciMLBase` provides `solve` / `remake`, `DataInterpolations` +backs the per-window input signal in autoregressive mode, and an +OrdinaryDiffEq solver package (e.g. `OrdinaryDiffEqTsit5`) supplies the +concrete solver type. ## Building a Lorenz dataset From d9b5fda98ba3eb14be84e81ae13ef91526e672ef Mon Sep 17 00:00:00 2001 From: Saswat Susmoy Date: Tue, 30 Jun 2026 21:24:02 +0530 Subject: [PATCH 6/8] refactor(continuous-esn): introduce ContinuousESNCell, slim model to 3-field ESN shape MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address Francesco's review on #456: `ContinuousESN` now shares its struct shape with `ESN` — three fields `(reservoir, states_modifiers, readout)` — and the substance moves into a new `ContinuousESNCell` that mirrors `ESNCell`'s field layout plus an `equations` field carrying the ODE right-hand side. * `src/layers/continuous_esn_cell.jl` (NEW): `ContinuousESNCell <: AbstractSciMLProblemReservoir` so the existing continuous dispatch in `RCODEReservoirExt` fires on `rc.reservoir`. `initialparameters` returns the canonical `(input_matrix, reservoir_matrix, [bias])` trio. Default `_continuous_esn_rhs!` implements Lukoševičius 2012 §3.2.6 eq (5) verbatim — no α in the RHS; α emerges only when the ODE is forward-Euler discretised at step `Δt = α`. Only the `@concrete` inner constructor is exposed; the convenience builder lives in the extension's `ContinuousESN(...)` constructor (avoids a Pair-vs-activation ambiguity with the auto-generated inner ctor that Aqua flagged). * `src/models/continuous_esn.jl`: 13-field flat struct → 3-field container. Stub constructor errors without the extension; real one in `RCODEReservoirExt`. * `ext/RCODEReservoirExt.jl`: drop the bulk constructor + the old α-bearing RHS + the `_typed_zeros_init` helper. Add a slim `ContinuousESN(in_dims, res_dims, out_dims, [activation,] tspan, args...; use_bias, init_*, equations, state_modifiers, ...)` that validates inputs, builds the cell, and wraps it with `state_modifiers` + `LinearReadout`. Add `_collectstates(::ContinuousESNCell, ...)` (builds the `ODEProblem` on the fly from `cell.equations`, sized by `cell.out_dims`, with `ps.reservoir` merged into the solve `p`) and a parallel `_predict(::ContinuousESNCell, ..., steps; initialdata)` for the autoregressive path. Teacher-forced `_predict` inherits the generic `AbstractSciMLProblemReservoir` method unchanged. Rename `u_t` to `input_t` per the same review thread. * Field-list note: Francesco's sketch listed only the ESNCell-like fields. The cell additionally carries `tspan`, `args`, `kwargs` because it IS the dispatch target for the continuous path — without those fields `_collectstates` would have nothing to feed `solve`. Documented as an open question in the PR reply. * `test/test_continuous_esn.jl`: rewritten for the new field shape (`ps.reservoir.input_matrix` / `reservoir_matrix` / `bias`). Euler equivalence test moved from "α in the RHS, dt=1" to "no α in the RHS, dt=α with tspan=(0, T_steps·α)"; α values restricted to `{1.0, 0.5}` so the requested saveat grid aligns bit-for-bit with the Euler step boundaries (non-FP-exact α introduces a sub-step offset that the Euler interpolant smooths over). 9 testsets, 45 assertions, all green at `atol=1e-10` for the equivalence test. * Docs: add `ContinuousESNCell` to the layers API page and rewrite the Lorenz tutorial without the dropped `leak_coefficient` / `T` kwargs. --- docs/src/api/layers.md | 1 + docs/src/tutorials/continuous_esn.md | 95 ++++++---- ext/RCODEReservoirExt.jl | 256 ++++++++++++++++++------- src/ReservoirComputing.jl | 3 +- src/layers/continuous_esn_cell.jl | 125 +++++++++++++ src/models/continuous_esn.jl | 198 ++++++++++---------- test/test_continuous_esn.jl | 267 +++++++++++++-------------- 7 files changed, 607 insertions(+), 338 deletions(-) create mode 100644 src/layers/continuous_esn_cell.jl diff --git a/docs/src/api/layers.md b/docs/src/api/layers.md index ddd1fb3f3..dd9f28dd6 100644 --- a/docs/src/api/layers.md +++ b/docs/src/api/layers.md @@ -22,6 +22,7 @@ ```@docs AdditiveEIESNCell + ContinuousESNCell EIESNCell ES2NCell ESNCell diff --git a/docs/src/tutorials/continuous_esn.md b/docs/src/tutorials/continuous_esn.md index b00418b4a..d326d580e 100644 --- a/docs/src/tutorials/continuous_esn.md +++ b/docs/src/tutorials/continuous_esn.md @@ -1,33 +1,36 @@ # Continuous ESN: forecasting Lorenz -[`ContinuousESN`](@ref) is a thin wrapper around -[`SciMLProblemReservoir`](@ref) that pre-bakes the leaky-integrator +[`ContinuousESN`](@ref) is a thin convenience wrapper around a +[`ContinuousESNCell`](@ref) that pre-bakes the leaky-integrator continuous Echo State Network ODE of [Lukosevicius2012](@cite) §3.2.6 eq (5): ```math -\dot{\mathbf{x}}(t) = \alpha \left( - -\mathbf{x}(t) + \tanh\!\left( - \mathbf{W}_{\text{in}}\,\mathbf{u}(t) + \mathbf{W}_r\,\mathbf{x}(t) - + \mathbf{b} - \right) -\right) +\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) ``` -with leaking rate `α`. Forward-Euler discretisation at step `Δt = 1` -recovers the discrete leaky ESN `x(n+1) = (1-α) x(n) + α tanh(...)` -exactly. The reservoir matrices `W_r`, `W_in`, and optional bias `b` -live in `ps.reservoir` and are constructed by `setup(rng, rc)`; the -ODE solver and any solve-time keyword arguments are captured at -construction. Under the hood the same `RCODEReservoirExt` extension -that powers `SciMLProblemReservoir` also runs `ContinuousESN`. +No leaking-rate term `α` appears in the ODE — `α` emerges only when the +ODE is forward-Euler discretised with step `Δt = α`, recovering the +discrete leaky ESN update `x(n+1) = (1-α)·x(n) + α·tanh(…)` exactly. To +target an effective leak rate `α`, choose `tspan = (0, n_samples · α)` +so the per-window width matches the desired step. + +`ContinuousESN` shares its struct shape with [`ESN`](@ref): three +fields `(reservoir, states_modifiers, readout)`. The reservoir +matrices `W_in`, `W_r`, and optional bias `b` live in +`ps.reservoir` as `input_matrix`, `reservoir_matrix`, and `bias`, and +are constructed by `setup(rng, esn)`. The ODE solver and any +solve-time keyword arguments are captured at construction. Under the +hood the `RCODEReservoirExt` package extension runs the integration. This tutorial walks through training a `ContinuousESN` on Lorenz-63 data and rolling it forward autoregressively to reproduce the -attractor. `SciMLBase` provides `solve` / `remake`, `DataInterpolations` -backs the per-window input signal in autoregressive mode, and an -OrdinaryDiffEq solver package (e.g. `OrdinaryDiffEqTsit5`) supplies the -concrete solver type. +attractor. `SciMLBase` provides `solve` / `remake`, +`DataInterpolations` backs the per-window input signal in +autoregressive mode, and an OrdinaryDiffEq solver package +(e.g. `OrdinaryDiffEqTsit5`) supplies the concrete solver type. ## Building a Lorenz dataset @@ -60,27 +63,41 @@ test = data[:, (shift + train_len):(shift + train_len + predict_len - 1)] ## Constructing the `ContinuousESN` -The constructor mirrors `SciMLProblemReservoir`: the integration -`tspan` and the solver positional argument are captured at -construction time, exactly as in DiffEqFlux's `NeuralODE`. +The constructor signature mirrors `ESN`: `(in_dims, res_dims, out_dims, +[activation,] tspan, args...; kwargs...)`. The integration `tspan` and +the solver positional argument are captured at construction, exactly +as in DiffEqFlux's `NeuralODE`. ```@example continuous-esn-lorenz N_res = 100 -ce_train = ContinuousESN( - 3, N_res, (0.0, Float64(train_len)), Tsit5(); - leak_coefficient = 1.0, use_bias = true, T = Float64, + +# 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 ) -ce_pred = ContinuousESN( - 3, N_res, (0.0, Float64(predict_len)), Tsit5(); - leak_coefficient = 1.0, use_bias = true, T = Float64, +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 ) -rc_train = ReservoirComputer(ce_train, (NLAT2(),), LinearReadout(N_res => 3)) -rc_pred = ReservoirComputer(ce_pred, (NLAT2(),), LinearReadout(N_res => 3)) - -ps, st = setup(rng, rc_train) +ps, st = setup(rng, esn_train) ``` ## Training @@ -90,25 +107,25 @@ ps, st = setup(rng, rc_train) on the collected states. ```@example continuous-esn-lorenz -ps, st = train!(rc_train, input_data, target_data, ps, st) +ps, st = train!(esn_train, input_data, target_data, ps, st) ``` ## Autoregressive rollout -`predict(rc, steps, ps, st; initialdata)` splits the predict-time +`predict(esn, steps, ps, st; initialdata)` splits the predict-time `tspan` into `steps` equal sub-intervals; on each sub-interval the previous readout output is held constant as the input. The default -initial reservoir state is `prob.u0` (zeros) — if you want to continue -from the trained reservoir's terminal state, `remake` the prob before -constructing the predict reservoir. +initial reservoir state is zeros sized to `res_dims` — if you want to +continue from the trained reservoir's terminal state, pass a custom +`equations` closure that captures it. ```@example continuous-esn-lorenz -ps_pred, st_pred = setup(rng, rc_pred) +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( - rc_pred, predict_len, ps_pred, st_pred; initialdata = test[:, 1] + esn_pred, predict_len, ps_pred, st_pred; initialdata = test[:, 1] ) plot( diff --git a/ext/RCODEReservoirExt.jl b/ext/RCODEReservoirExt.jl index c8acbf9eb..5f0da9736 100644 --- a/ext/RCODEReservoirExt.jl +++ b/ext/RCODEReservoirExt.jl @@ -17,7 +17,10 @@ using ReservoirComputing: ReservoirComputing, AbstractSampler, AbstractSciMLProblemReservoir, ContinuousESN, + ContinuousESNCell, + LinearReadout, TerminalStateSampling, + _wrap_layers, collectstates, rand_sparse, scaled_rand @@ -395,101 +398,214 @@ function _predict( end # --------------------------------------------------------------------------- -# ContinuousESN +# ContinuousESN — convenience constructor # -# Concrete subtype of `AbstractSciMLProblemReservoir` that pre-bakes the -# leaky-integrator continuous ESN ODE of Lukoševičius 2012 §3.2.6 eq (5): -# -# ẋ(t) = α · (-x(t) + tanh(W_in·u(t) + W_r·x(t) + b)) -# -# `W_r`, `W_in`, `b`, and the scalar `leak_coefficient = α` live in -# `ps.reservoir` and are injected into the solve `p` by -# `_build_solve_params`. The RHS is in-place (`mul!` for both matvecs, -# fused tanh broadcast) so it stays allocation-free on the solver hot -# path. +# Builds the 3-field `(reservoir, states_modifiers, readout)` model whose +# `reservoir` is a `ContinuousESNCell`. The cell carries Lukoševičius 2012 +# §3.2.6 eq (5) as its `equations` field. Each of `collectstates`, +# `predict(rc, data, ...)`, and `predict(rc, steps, ...; initialdata)` +# dispatches on `rc.reservoir::ContinuousESNCell` below, building the +# `ODEProblem` lazily so the reservoir weights stay in `ps.reservoir` and +# can be re-initialised by `setup(rng, rc)` like every other ESN-family +# model. # --------------------------------------------------------------------------- -function _continuous_esn_rhs!(dx, x, p, t) - # dx ← W_r · x - mul!(dx, p.W_r, x) - # dx ← dx + W_in · u(t) - u_t = p.input(t) - mul!(dx, p.W_in, u_t, true, true) - # dx ← α · (-x + tanh(dx + b?)) - # - # `haskey` on a NamedTuple resolves at compile time (the keys are - # part of the type), so the two branches don't add runtime cost. - if haskey(p, :b) - @. dx = p.leak_coefficient * (-x + tanh(dx + p.b)) - else - @. dx = p.leak_coefficient * (-x + tanh(dx)) - end - return nothing -end - -# Type-aware default bias initialiser — `(rng, T, dims...) -> zeros(T, dims...)`. -# `WeightInitializers.zeros32` is Float32-only and ignores `T`, so we ship a -# small `T`-respecting default and document the expected signature in the -# constructor docstring. -_typed_zeros_init(::AbstractRNG, ::Type{T}, dims::Integer...) where {T} = - zeros(T, dims...) - function ReservoirComputing.ContinuousESN( - in_dims::Integer, res_dims::Integer, tspan, args...; - leak_coefficient::Real = 1.0, + 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_bias = _typed_zeros_init, - T::Type{<:AbstractFloat} = Float32, + 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")) - leak_coefficient > 0 || throw( - ArgumentError("leak_coefficient must be positive, got $leak_coefficient") - ) + 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" - ) + ArgumentError("tspan endpoints must be finite, got $tspan") ) tspan[2] > tspan[1] || throw( ArgumentError( - "SciMLProblemReservoir requires `tspan[2] > tspan[1]`, got tspan = $tspan" + "ContinuousESN requires `tspan[2] > tspan[1]`, got tspan = $tspan" ) ) ReservoirComputing._check_protected_kwargs(kwargs) - # Zero initial reservoir state — Lukoševičius 2012 starts from x(0) = 0 - # by convention. Users continuing from a previously trained terminal - # state should `remake(rc.reservoir.prob; u0 = ...)` themselves. - u0 = zeros(T, res_dims) - - # `prob.p = nothing` — all reservoir parameters live in `ps.reservoir` - # and get merged into the solve `p` by `_build_solve_params`. - prob = ODEProblem(_continuous_esn_rhs!, u0, tspan) - - return ContinuousESN( - prob, - TerminalStateSampling(), - tspan, - args, - kwargs, - in_dims, - res_dims, - leak_coefficient, - use_bias, - init_reservoir, - init_input, - init_bias, - T + 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 + +# Allow omitting `activation` (defaults to `tanh`) by routing +# `ContinuousESN(in, res, out, tspan, args...)` through the five-arg form. +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 + +# --------------------------------------------------------------------------- +# Continuous `_collectstates` for `ContinuousESNCell` +# +# Builds the `ODEProblem` on the fly from `cell.equations` and an initial +# state of zeros sized to `cell.out_dims`. The weight matrices come from +# `ps.reservoir` (`input_matrix`, `reservoir_matrix`, optional `bias`) and +# are merged into the solve `p` by `_build_solve_params`. Sampler is fixed +# to `TerminalStateSampling` since eq (5) only exposes a single +# point-state per window. +# --------------------------------------------------------------------------- + +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 + +# --------------------------------------------------------------------------- +# Autoregressive `_predict` for `ContinuousESNCell` +# +# Same control flow as the generic `AbstractSciMLProblemReservoir` path, +# but the `ODEProblem` is built on demand from `cell.equations` rather +# than pulled from a pre-built `res.prob`. Initial reservoir state is +# zeros sized to `cell.out_dims`; users who want to continue from a +# previously trained terminal state should pass it in via a custom +# `equations` closure or re-run `collectstates` first. +# --------------------------------------------------------------------------- + +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 20d348d8d..327678334 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") @@ -66,7 +67,7 @@ include("extensions/reca.jl") export ReservoirComputer export AbstractSciMLProblemReservoir, SciMLProblemReservoir, ContinuousESN export AbstractSampler, TerminalStateSampling -export ESNCell, ES2NCell, EuSNCell, MemoryESNCell, MemoryResESNCell, RMNCell +export ContinuousESNCell, ESNCell, ES2NCell, EuSNCell, MemoryESNCell, MemoryResESNCell, RMNCell export StatefulLayer, LinearReadout, ReservoirChain, Collect, collectstates, DelayLayer, NonlinearFeaturesLayer export AdditiveEIESNCell, EIESNCell, ES2NCell, ESNCell, EuSNCell, LIFESNCell, ResESNCell diff --git a/src/layers/continuous_esn_cell.jl b/src/layers/continuous_esn_cell.jl new file mode 100644 index 000000000..8439cc572 --- /dev/null +++ b/src/layers/continuous_esn_cell.jl @@ -0,0 +1,125 @@ +@doc raw""" + ContinuousESNCell(activation, in_dims, out_dims, + init_bias, init_reservoir, init_input, init_state, + use_bias, equations, tspan, args, kwargs) + +Continuous-time leaky-integrator ESN cell ([Lukosevicius2012](@cite) §3.2.6, +eq 5). `ContinuousESNCell <: AbstractSciMLProblemReservoir` so it dispatches +into the continuous `collectstates` / `predict` pipeline provided by the +`RCODEReservoirExt` package extension. Field layout mirrors [`ESNCell`](@ref) +plus the ODE right-hand side (`equations`) and the solve metadata +(`tspan`, `args`, `kwargs`) needed by the extension's solver. + +The default `equations` value [`_continuous_esn_rhs!`](@ref) integrates +Lukoševičius 2012 §3.2.6 eq (5): + +```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) +``` + +No leaking-rate term `α` appears in the ODE — `α` emerges only when the +ODE is forward-Euler discretised with step `Δt = α`, recovering the +discrete leaky update `x(n+1) = (1-α)·x(n) + α·tanh(...)`. + +This type is normally constructed by the [`ContinuousESN`](@ref) +convenience constructor in the `RCODEReservoirExt` extension rather than +directly. Driving the inner constructor by hand requires placing every +field in the order above and pre-static-ing `use_bias`. + +## Fields + + - `activation`: Reservoir nonlinearity. Documented as `tanh` by the + `ContinuousESN` constructor; only enters the cell via a custom + `equations` closure. + - `in_dims`, `out_dims`: Input and reservoir dimensions. + - `init_input`, `init_reservoir`, `init_bias`, `init_state`: + Initialisers (same convention as [`ESNCell`](@ref)). + - `use_bias::StaticBool`: Static-bool flag controlling the optional + `bias` parameter. + - `equations`: ODE right-hand side `(dx, x, p, t) -> nothing`. The + extension injects `p.input(t)` (interpolated input signal) and merges + `ps.reservoir` into `p`. + - `tspan`: Integration interval `(t0, t1)` for `collectstates`. + - `args`, `kwargs`: Forwarded to `solve` positionally / by keyword. + +## 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 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 + +@doc raw""" + _continuous_esn_rhs!(dx, x, p, t) + +Default right-hand side for [`ContinuousESNCell`](@ref). Implements +Lukoševičius 2012 §3.2.6 eq (5) in place: + +```math +\dot{\mathbf{x}} = -\mathbf{x} + \tanh\!\left( + \mathbf{W}_{\text{in}}\,\mathbf{u}(t) + + \mathbf{W}_r\,\mathbf{x} + \mathbf{b}\right) +``` + +`p` is the merged parameter pack assembled by the `RCODEReservoirExt` +helper. It carries: + * `p.input` — interpolated input signal injected by the extension. + * `p.input_matrix` (`W_in`), `p.reservoir_matrix` (`W_r`) — reservoir + matrices pulled from `ps.reservoir`. + * `p.bias` (optional) — bias vector, present iff `use_bias=True()`. + +The RHS uses `mul!` for both matrix products and a fused tanh broadcast +to stay allocation-free on the solver hot path. The bias branch is +selected at compile time via `haskey`, since the keys of a `NamedTuple` +are part of its type. +""" +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/models/continuous_esn.jl b/src/models/continuous_esn.jl index d833b9be9..fb9c7d149 100644 --- a/src/models/continuous_esn.jl +++ b/src/models/continuous_esn.jl @@ -1,96 +1,113 @@ @doc raw""" - ContinuousESN(in_dims, res_dims, tspan, args...; kwargs...) + ContinuousESN(in_dims, res_dims, out_dims, activation = tanh, + tspan, args...; + use_bias = false, + init_reservoir = rand_sparse, init_input = scaled_rand, + init_bias = zeros32, init_state = randn32, + state_modifiers = (), readout_activation = identity, + equations = _continuous_esn_rhs!, + kwargs...) Continuous-time leaky-integrator Echo State Network ([Lukosevicius2012](@cite) -§3.2.6, eq 5). `ContinuousESN <: AbstractSciMLProblemReservoir`, so it reuses -the continuous-time `_collectstates` / `_predict` machinery in the -`RCODEReservoirExt` package extension and pre-bakes the leaky-integrator -reservoir ODE +§3.2.6, eq 5). Thin convenience wrapper that composes: + + 1) a [`ContinuousESNCell`](@ref) carrying the reservoir matrices and the + ODE right-hand side, + 2) zero or more `state_modifiers` applied to the sampled reservoir + states, and + 3) a [`LinearReadout`](@ref) mapping the modified states to outputs. + +Structurally identical to [`ESN`](@ref) — three fields +`(reservoir, states_modifiers, readout)` — so it slots into the same +training and prediction pipeline. The reservoir field is the +`ContinuousESNCell` rather than a `StatefulLayer(ESNCell)`, which routes +`collectstates` / `predict` through the continuous-time +`RCODEReservoirExt` extension. + +## Equations ```math -\dot{\mathbf{x}}(t) = \alpha \left( - -\mathbf{x}(t) + \tanh\!\left( - \mathbf{W}_{\text{in}}\,\mathbf{u}(t) + \mathbf{W}_r\,\mathbf{x}(t) + \mathbf{b} - \right) -\right) +\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} ``` -with reservoir matrix `W_r`, input matrix `W_in`, optional bias `b`, and -scalar leaking rate `α` exposed as `leak_coefficient`. Forward-Euler -discretisation at step `Δt = 1` collapses the ODE to the discrete leaky -update `x(n+1) = (1-α)·x(n) + α·tanh(W_in·u(n) + W_r·x(n) + b)`. - -This is *not* a port of the parametric-surrogate CTESN of -[Anantharaman2021](@cite); that model uses `ṙ = tanh(A·r + W_hyb·u)` without -a decay term and is trained via RBF interpolation of `W_out(p)` over a -parameter space. A separate type would land for it in a future PR. +No leaking-rate term `α` appears in the ODE. `α` emerges only when the +ODE is forward-Euler discretised with step `Δt = α`, recovering the +discrete leaky update of [`ESN`](@ref). To change the effective leak, +adjust `tspan` so the per-input-window width matches the desired `Δt`. ## Arguments - `in_dims`: Input dimension. - `res_dims`: Reservoir (hidden state) dimension. - - `tspan`: Integration interval for `collectstates`. Must be a length-2, - strictly-increasing, finite tuple/pair. Mirrors the DiffEqFlux - `NeuralODE` convention. - - `args...`: Positional arguments forwarded to `solve`. The solver - algorithm (e.g. `Tsit5()`, `Euler()`) is the first element by - convention. + - `out_dims`: Output dimension. + - `activation`: Reservoir activation passed to the cell. Default: `tanh`. + - `tspan`: Integration interval `(t0, t1)`. Length-2, strictly + increasing, finite. + - `args...`: Forwarded to `solve` positionally. The solver algorithm + (`Tsit5()`, `Euler()`, …) is the first element by convention. ## Keyword arguments - - `leak_coefficient`: Leaking rate `α`. Default: `1.0`. Pairing `Δt = 1` - with `leak_coefficient = α` recovers the discrete leaky ESN under - forward Euler. - - `use_bias`: Whether to include a bias term `b`. Default: `false`. - - `init_reservoir`: Initializer for `W_r`. Default: - [`rand_sparse`](@ref). Receives `(rng, T, res_dims, res_dims)`. - - `init_input`: Initializer for `W_in`. Default: [`scaled_rand`](@ref). - Receives `(rng, T, res_dims, in_dims)`. - - `init_bias`: Initializer for `b`. Used only when `use_bias = true`. - Default: a type-aware zeros initialiser. Receives `(rng, T, res_dims)`. - - `T`: Element type for the reservoir matrices and initial state. - Default: `Float32`. +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. Default: + [`_continuous_esn_rhs!`](@ref). + +Composition: + + - `state_modifiers`: A layer or collection of layers applied to the + sampled 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 three keys `saveat`, - `save_everystep`, and `dense` are owned by the extension helper and - rejected at construction. + `save_everystep`, and `dense` are rejected at construction. -## Continuous-time stability note +## Parameters -`scale_radius!` and the default reservoir initialisers control the -discrete-time spectral radius `ρ(W_r)`. For the α-factored ODE used here, -a log-norm contraction argument gives the sufficient condition -`‖W_r‖_2 < 1` (operator 2-norm bound), independent of `α`: the leaking -rate scales the rate of decay but not the asymptotic bound. The condition -is strictly stronger than `ρ(W_r) < 1` because `‖·‖_2 ≥ ρ(·)`. Published -continuous ESN work treats `ρ` as an empirical hyperparameter — adjust -accordingly. + - `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` — a `Tuple` with parameters for each modifier + layer (may be empty). + - `readout` — parameters of [`LinearReadout`](@ref), typically: + - `weight :: (out_dims × res_dims)` — `W_out` + - `bias :: (out_dims,)` — `b_out` (if the readout uses bias) !!! note This constructor errors unless `RCODEReservoirExt` is loaded. Load - `SciMLBase`, `DataInterpolations`, and a concrete OrdinaryDiffEq + `SciMLBase`, `DataInterpolations`, and a concrete `OrdinaryDiffEq` solver package (e.g. `OrdinaryDiffEqTsit5`, `OrdinaryDiffEq`) to enable it. """ -@concrete struct ContinuousESN <: AbstractSciMLProblemReservoir - prob - sampler - tspan - args - kwargs - in_dims - res_dims - leak_coefficient - use_bias - init_reservoir - init_input - init_bias - matrix_type +@concrete struct ContinuousESN <: + AbstractEchoStateNetwork{(:reservoir, :states_modifiers, :readout)} + reservoir + states_modifiers + readout end # Public factory — the real implementation lives in `RCODEReservoirExt`. # Loading `SciMLBase` + `DataInterpolations` adds the typed method that -# actually builds the `ODEProblem`; without the extension, this fallback -# fires and produces a clear error. +# actually builds the cell; without the extension, this fallback fires +# and produces a clear error. function ContinuousESN(::Any, ::Any, ::Any, ::Any...; kwargs...) return error( "ContinuousESN requires the RCODEReservoirExt extension and an " * @@ -99,36 +116,29 @@ function ContinuousESN(::Any, ::Any, ::Any, ::Any...; kwargs...) ) end -function Base.show(io::IO, ce::ContinuousESN) - print(io, "ContinuousESN(") - print(io, "in_dims = ", ce.in_dims) - print(io, ", res_dims = ", ce.res_dims) - print(io, ", leak_coefficient = ", ce.leak_coefficient) - print(io, ", use_bias = ", ce.use_bias) - print(io, ", tspan = ") - show(io, ce.tspan) - print(io, ")") - return -end +function Base.show(io::IO, esn::ContinuousESN) + print(io, "ContinuousESN(\n") -function initialparameters(rng::AbstractRNG, ce::ContinuousESN) - T = ce.matrix_type - W_r = ce.init_reservoir(rng, T, ce.res_dims, ce.res_dims) - W_in = ce.init_input(rng, T, ce.res_dims, ce.in_dims) - leak_coefficient = T(ce.leak_coefficient) - isfinite(leak_coefficient) || throw( - ArgumentError( - "leak_coefficient converted to $T overflowed to $leak_coefficient; " * - "pick a smaller value or widen `T`." - ) - ) - ps = (W_r = W_r, W_in = W_in, leak_coefficient = leak_coefficient) - if ce.use_bias - ps = merge(ps, (b = ce.init_bias(rng, T, ce.res_dims),)) + 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 - return ps -end + print(io, ",\n") -function initialstates(::AbstractRNG, ::ContinuousESN) - return NamedTuple() + 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 index de59d717b..e6afd53e4 100644 --- a/test/test_continuous_esn.jl +++ b/test/test_continuous_esn.jl @@ -4,71 +4,63 @@ using LinearAlgebra using Statistics using ReservoirComputing using OrdinaryDiffEq +using OrdinaryDiffEqLowOrderRK: Euler using SciMLBase using DataInterpolations # --------------------------------------------------------------------------- # 1. Construction: shape + parameter wiring # -# `ContinuousESN(in_dims, res_dims, tspan, solver)` should build a reservoir -# whose `initialparameters` yields the canonical eq-(5) trio (`W_r`, -# `W_in`, `inv_leak`) at the correct shapes, with no bias by default. +# `ContinuousESN(in_dims, res_dims, out_dims, tspan, solver)` builds a +# 3-field `(reservoir, states_modifiers, readout)` model. The reservoir +# is a `ContinuousESNCell`; `ps.reservoir` exposes the canonical ESN-cell +# parameter trio (`input_matrix`, `reservoir_matrix`, optional `bias`). # --------------------------------------------------------------------------- @testset "ContinuousESN: construction + parameter shapes" begin rng = MersenneTwister(0) - in_dim, res_dim = 3, 50 - ce = ContinuousESN(in_dim, res_dim, (0.0, 5.0), Tsit5()) + in_dim, res_dim, out_dim = 3, 50, 2 + esn = ContinuousESN(in_dim, res_dim, out_dim, (0.0, 5.0), Tsit5()) - rc = ReservoirComputer(ce, LinearReadout(res_dim => 2)) - ps, st = setup(rng, rc) + @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) - @test ce isa AbstractSciMLProblemReservoir - @test size(ps.reservoir.W_r) == (res_dim, res_dim) - @test size(ps.reservoir.W_in) == (res_dim, in_dim) - @test ps.reservoir.leak_coefficient == 1.0f0 - @test !haskey(ps.reservoir, :b) - @test all(isfinite, ps.reservoir.W_r) - @test all(isfinite, ps.reservoir.W_in) + 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 # --------------------------------------------------------------------------- # 2. Construction validation: argument errors # -# Positive `in_dims`, `res_dims`, `leak_coefficient`, and strictly -# increasing `tspan` are required. Protected `solve` kwargs are rejected -# at construction (inherits `SciMLProblemReservoir`'s _PROTECTED_SOLVE_KWARGS). +# Positive `in_dims`, `res_dims`, `out_dims`, and a finite, strictly +# increasing length-2 `tspan` are required. Protected `solve` kwargs are +# rejected at construction. # --------------------------------------------------------------------------- @testset "ContinuousESN: construction validation" begin - @test_throws ArgumentError ContinuousESN(0, 5, (0.0, 1.0), Tsit5()) - @test_throws ArgumentError ContinuousESN(3, 0, (0.0, 1.0), Tsit5()) - @test_throws ArgumentError ContinuousESN(3, 5, (0.0, 0.0), Tsit5()) - @test_throws ArgumentError ContinuousESN(3, 5, (1.0, 0.0), Tsit5()) - @test_throws ArgumentError ContinuousESN( - 3, 5, (0.0, 1.0), Tsit5(); leak_coefficient = 0.0 - ) + @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()) # Non-finite endpoints - @test_throws ArgumentError ContinuousESN(3, 5, (0.0, Inf), Tsit5()) - @test_throws ArgumentError ContinuousESN(3, 5, (-Inf, 1.0), Tsit5()) - @test_throws ArgumentError ContinuousESN(3, 5, (0.0, NaN), 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()) # Wrong-length tspan - @test_throws ArgumentError ContinuousESN(3, 5, (1.0,), Tsit5()) - @test_throws ArgumentError ContinuousESN(3, 5, (0.0, 1.0, 2.0), Tsit5()) - # leak_coefficient that overflows `T` is caught at `setup` time, not - # construction (the construction check sees the unconverted value). - # Use a dense init to avoid `rand_sparse`'s NaN-on-tiny-matrix edge. - let dense_init = (rng, T, d...) -> rand_sparse(rng, T, d...; sparsity = 0.5) - ce_overflow = ContinuousESN( - 3, 16, (0.0, 1.0), Tsit5(); - leak_coefficient = 1.0e50, T = Float32, init_reservoir = dense_init - ) - rc_overflow = ReservoirComputer(ce_overflow, LinearReadout(16 => 1)) - @test_throws ArgumentError setup(MersenneTwister(0), rc_overflow) - end + @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, (0.0, 1.0), Tsit5(); (badkw => true,)... + 3, 5, 2, (0.0, 1.0), Tsit5(); (badkw => true,)... ) end end @@ -76,118 +68,118 @@ end # --------------------------------------------------------------------------- # 3. Bias toggle # -# `use_bias = true` should add a `b` parameter of length `res_dims`. +# `use_bias = true` adds a `bias` parameter of length `res_dims`. # --------------------------------------------------------------------------- @testset "ContinuousESN: bias toggle" begin rng = MersenneTwister(1) - in_dim, res_dim = 2, 10 - ce_b = ContinuousESN(in_dim, res_dim, (0.0, 1.0), Tsit5(); use_bias = true) - rc = ReservoirComputer(ce_b, LinearReadout(res_dim => 1)) - ps, st = setup(rng, rc) - @test haskey(ps.reservoir, :b) - @test length(ps.reservoir.b) == res_dim - # Default init is the type-aware zeros initialiser. - @test all(==(0), ps.reservoir.b) + 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 + # Default init is `zeros32`. + @test all(==(0), ps.reservoir.bias) end # --------------------------------------------------------------------------- -# 3b. Bias actually exercised under an adaptive solver +# 4. Bias actually exercised under an adaptive solver # -# The `haskey(:b)` branch in `_continuous_esn_rhs!` is solver-agnostic in -# theory, but only the Euler-equivalence test (which uses `Euler`) hits -# the bias-on path under integration. Pair `use_bias = true` with a -# nonzero bias initialiser and `Tsit5` to confirm: (a) the bias does -# perturb states relative to the no-bias baseline, and (b) the result -# stays finite. +# Pair `use_bias = true` with a nonzero bias initialiser and Tsit5 to +# confirm: (a) the bias does perturb states relative to the no-bias +# baseline, and (b) the result stays finite. # --------------------------------------------------------------------------- @testset "ContinuousESN: bias path under Tsit5" begin rng = MersenneTwister(101) - in_dim, res_dim, T_steps = 2, 20, 12 - nonzero_bias(rng, T, d...) = T(0.5) .* randn(rng, T, d...) - ce_with_bias = ContinuousESN( - in_dim, res_dim, (0.0, 2.0), Tsit5(); + 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 ) - ce_no_bias = ContinuousESN( - in_dim, res_dim, (0.0, 2.0), Tsit5(); + 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 ) - rc_b = ReservoirComputer(ce_with_bias, LinearReadout(res_dim => 1)) - rc_n = ReservoirComputer(ce_no_bias, LinearReadout(res_dim => 1)) - ps_b, st_b = setup(MersenneTwister(0), rc_b) - ps_n, st_n = setup(MersenneTwister(0), rc_n) + 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(rc_b, data, ps_b, st_b) - s_n, _ = collectstates(rc_n, data, ps_n, st_n) + 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 # --------------------------------------------------------------------------- -# 4. Forward pass: shape + finiteness + determinism +# 5. Forward pass: shape + finiteness + determinism # --------------------------------------------------------------------------- @testset "ContinuousESN: forward (collectstates)" begin rng = MersenneTwister(7) - in_dim, res_dim, T_steps = 2, 16, 25 - ce = ContinuousESN( - in_dim, res_dim, (0.0, 3.0), Tsit5(); + 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 ) - rc = ReservoirComputer(ce, LinearReadout(res_dim => 1)) - ps, st = setup(rng, rc) + ps, st = setup(rng, esn) data = randn(Float32, in_dim, T_steps) - s1, _ = collectstates(rc, data, ps, st) - s2, _ = collectstates(rc, data, ps, st) + 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 # --------------------------------------------------------------------------- -# 5. Euler equivalence with discrete leaky ESN +# 6. Euler equivalence with discrete leaky ESN # -# Solving `ẋ = α(-x + tanh(W_r·x + W_in·u(t) + b))` with explicit Euler at -# step `Δt = 1` and `leak_coefficient = α` collapses algebraically to the -# discrete leaky ESN update `x_{k+1} = (1-α)·x_k + α·tanh(...)`. With the -# corrected window alignment (input at window start, sample at window -# end), continuous and discrete trajectories agree exactly. -# -# Covers both vanilla (α = 1) and leaky (α = 0.3) regimes. +# Eq (5) is `ẋ = -x + tanh(W_r·x + W_in·u(t) + b)`. Forward-Euler with +# step `Δt = α` gives `x_{k+1} = x_k + α·(-x_k + tanh(...)) = +# (1-α)·x_k + α·tanh(...)`, i.e. the discrete leaky ESN update at leak +# rate α. To match `T_steps` discrete steps we set `tspan = (0, T_steps·α)` +# so the per-window width `Δt = α` matches the Euler step `dt = α`. # --------------------------------------------------------------------------- @testset "ContinuousESN: Euler equivalence with discrete leaky ESN" begin - # Dense reservoir init avoids `rand_sparse`'s spectral-scaling edge case - # on very small matrices (8×8 at 10% density can yield a singular matrix - # whose `eigvals` are NaN, tripping `check_inf_nan`). - dense_init(rng, T, d...) = rand_sparse(rng, T, d...; sparsity = 0.5) - for α in (1.0, 0.3) + # `α` values that are exactly representable in Float64 so that + # `tspan = (0, T_steps · α)` and the integrator step `dt = α` align + # bit-for-bit with the requested `saveat` grid. Non-FP-exact `α` + # (e.g. 0.3) introduces a sub-step offset between saveat and the + # integrator's step boundaries, which the linear interpolant inside + # Euler smooths over — equivalence then holds only approximately. + 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, T_steps = 2, 16, 12 + in_dim, res_dim, out_dim, T_steps = 2, 16, 1, 12 + tspan = (0.0, T_steps * α) - ce = ContinuousESN( - in_dim, res_dim, (0.0, Float64(T_steps)), Euler(); - leak_coefficient = α, use_bias = true, T = Float64, dt = 1.0, - init_reservoir = dense_init + 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 ) - rc = ReservoirComputer(ce, LinearReadout(res_dim => 1)) - ps, st = setup(MersenneTwister(0), rc) + ps, st = setup(MersenneTwister(0), esn) - data = randn(rng, in_dim, T_steps) - cont_states, _ = collectstates(rc, data, ps, st) + data = randn(rng, Float64, in_dim, T_steps) + cont_states, _ = collectstates(esn, data, ps, st) - W_r = ps.reservoir.W_r - W_in = ps.reservoir.W_in - b = ps.reservoir.b + 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) + 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 @@ -195,22 +187,21 @@ end end # --------------------------------------------------------------------------- -# 6. Teacher-forced predict + autoregressive predict shape/determinism +# 7. Teacher-forced predict + autoregressive predict shape/determinism # --------------------------------------------------------------------------- @testset "ContinuousESN: teacher-forced predict" begin rng = MersenneTwister(21) in_dim, res_dim, out_dim, T_steps = 2, 12, 3, 10 - ce = ContinuousESN( - in_dim, res_dim, (0.0, 2.0), Tsit5(); + esn = ContinuousESN( + in_dim, res_dim, out_dim, (0.0, 2.0), Tsit5(); reltol = 1.0e-8, abstol = 1.0e-10 ) - rc = ReservoirComputer(ce, LinearReadout(res_dim => out_dim)) - ps, st = setup(rng, rc) + ps, st = setup(rng, esn) data = randn(Float32, in_dim, T_steps) - tf1, _ = predict(rc, data, ps, st) - tf2, _ = predict(rc, data, ps, st) + 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 @@ -220,56 +211,64 @@ end rng = MersenneTwister(22) # Autoregressive rollout feeds outputs back as inputs, so in_dim == out_dim. dim, res_dim, steps = 3, 12, 5 - ce = ContinuousESN( - dim, res_dim, (0.0, 1.0), Tsit5(); + esn = ContinuousESN( + dim, res_dim, dim, (0.0, 1.0), Tsit5(); reltol = 1.0e-8, abstol = 1.0e-10 ) - rc = ReservoirComputer(ce, LinearReadout(res_dim => dim)) - ps, st = setup(rng, rc) + ps, st = setup(rng, esn) init = randn(Float32, dim) - ar1, _ = predict(rc, steps, ps, st; initialdata = init) - ar2, _ = predict(rc, steps, ps, st; initialdata = init) + 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 # --------------------------------------------------------------------------- -# 7. Compatibility with state modifiers +# 8. Compatibility with state modifiers # --------------------------------------------------------------------------- @testset "ContinuousESN: state modifiers compose" begin rng = MersenneTwister(31) - in_dim, res_dim, T_steps = 2, 10, 8 - ce = ContinuousESN( - in_dim, res_dim, (0.0, 1.0), Tsit5(); + 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 ) - rc_plain = ReservoirComputer(ce, LinearReadout(res_dim => 1)) - rc_mod = ReservoirComputer(ce, (NLAT2(),), LinearReadout(res_dim => 1)) + 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), rc_plain) - ps_m, st_m = setup(MersenneTwister(0), rc_mod) + 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(rc_plain, data, ps_p, st_p) - sm, _ = collectstates(rc_mod, data, ps_m, st_m) + 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 # --------------------------------------------------------------------------- -# 8. Eltype propagates through `T` kwarg +# 9. Custom init eltype propagates +# +# Without a `T` kwarg, the eltype is whatever the `init_*` initialisers +# produce. Passing Float64 initialisers gives Float64 matrices. # --------------------------------------------------------------------------- -@testset "ContinuousESN: T kwarg controls matrix eltype" begin +@testset "ContinuousESN: custom init eltype propagates" begin rng = MersenneTwister(43) - ce64 = ContinuousESN(2, 32, (0.0, 1.0), Tsit5(); T = Float64) - rc = ReservoirComputer(ce64, LinearReadout(32 => 1)) - ps, _ = setup(rng, rc) - @test eltype(ps.reservoir.W_r) == Float64 - @test eltype(ps.reservoir.W_in) == Float64 - @test typeof(ps.reservoir.leak_coefficient) == Float64 + 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 From 181deb8355aabc8405b19b0ef45b48cf032249e1 Mon Sep 17 00:00:00 2001 From: Saswat Susmoy Date: Wed, 1 Jul 2026 21:41:43 +0530 Subject: [PATCH 7/8] style(continuous-esn): address review feedback on docstrings, comments, and cell constructor Clean up implementation details from docstrings and comments across the ContinuousESN model, cell, tutorial, and extension. Add a Pair-signature convenience constructor for ContinuousESNCell matching the ESNCell pattern. Move ContinuousESNCell to a new Continuous-time Layers subsection in the API docs and drop the unused Anantharaman2021 citation. --- .gitignore | 3 +- docs/src/api/layers.md | 7 +- docs/src/refs.bib | 7 -- docs/src/tutorials/continuous_esn.md | 57 ++--------- ext/RCODEReservoirExt.jl | 137 --------------------------- src/layers/continuous_esn_cell.jl | 98 +++++++++---------- src/models/continuous_esn.jl | 83 +++++++--------- test/test_continuous_esn.jl | 69 -------------- 8 files changed, 99 insertions(+), 362 deletions(-) diff --git a/.gitignore b/.gitignore index 09c048f85..20e9a5126 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/src/api/layers.md b/docs/src/api/layers.md index dd9f28dd6..48928f704 100644 --- a/docs/src/api/layers.md +++ b/docs/src/api/layers.md @@ -22,7 +22,6 @@ ```@docs AdditiveEIESNCell - ContinuousESNCell EIESNCell ES2NCell ESNCell @@ -37,6 +36,12 @@ LIFESNCell ``` +## Continuous-time Layers + +```@docs + ContinuousESNCell +``` + ## Wrappers ```@docs diff --git a/docs/src/refs.bib b/docs/src/refs.bib index f895bf0a8..9e15668c2 100644 --- a/docs/src/refs.bib +++ b/docs/src/refs.bib @@ -570,10 +570,3 @@ @incollection{Lukosevicius2012 pages = {659--686} } -@article{Anantharaman2021, - title = {Accelerating Simulation of Stiff Nonlinear Systems using Continuous-Time Echo State Networks}, - url = {https://arxiv.org/abs/2010.04004}, - author = {Anantharaman, Ranjan and Ma, Yingbo and Gowda, Shashi and Laughman, Christopher and Shah, Viral and Edelman, Alan and Rackauckas, Christopher}, - year = {2021}, - journal = {arXiv:2010.04004} -} diff --git a/docs/src/tutorials/continuous_esn.md b/docs/src/tutorials/continuous_esn.md index d326d580e..e99a15204 100644 --- a/docs/src/tutorials/continuous_esn.md +++ b/docs/src/tutorials/continuous_esn.md @@ -1,9 +1,7 @@ # Continuous ESN: forecasting Lorenz -[`ContinuousESN`](@ref) is a thin convenience wrapper around a -[`ContinuousESNCell`](@ref) that pre-bakes the leaky-integrator -continuous Echo State Network ODE of [Lukosevicius2012](@cite) §3.2.6 -eq (5): +[`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( @@ -11,26 +9,9 @@ eq (5): + \mathbf{b}\right) ``` -No leaking-rate term `α` appears in the ODE — `α` emerges only when the -ODE is forward-Euler discretised with step `Δt = α`, recovering the -discrete leaky ESN update `x(n+1) = (1-α)·x(n) + α·tanh(…)` exactly. To -target an effective leak rate `α`, choose `tspan = (0, n_samples · α)` -so the per-window width matches the desired step. - -`ContinuousESN` shares its struct shape with [`ESN`](@ref): three -fields `(reservoir, states_modifiers, readout)`. The reservoir -matrices `W_in`, `W_r`, and optional bias `b` live in -`ps.reservoir` as `input_matrix`, `reservoir_matrix`, and `bias`, and -are constructed by `setup(rng, esn)`. The ODE solver and any -solve-time keyword arguments are captured at construction. Under the -hood the `RCODEReservoirExt` package extension runs the integration. - -This tutorial walks through training a `ContinuousESN` on Lorenz-63 -data and rolling it forward autoregressively to reproduce the -attractor. `SciMLBase` provides `solve` / `remake`, -`DataInterpolations` backs the per-window input signal in -autoregressive mode, and an OrdinaryDiffEq solver package -(e.g. `OrdinaryDiffEqTsit5`) supplies the concrete solver type. +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 @@ -63,11 +44,6 @@ test = data[:, (shift + train_len):(shift + train_len + predict_len - 1)] ## Constructing the `ContinuousESN` -The constructor signature mirrors `ESN`: `(in_dims, res_dims, out_dims, -[activation,] tspan, args...; kwargs...)`. The integration `tspan` and -the solver positional argument are captured at construction, exactly -as in DiffEqFlux's `NeuralODE`. - ```@example continuous-esn-lorenz N_res = 100 @@ -102,23 +78,12 @@ ps, st = setup(rng, esn_train) ## Training -`train!` is sampler-agnostic: it routes through the continuous -`_collectstates` provided by the extension and fits a linear readout -on the collected states. - ```@example continuous-esn-lorenz ps, st = train!(esn_train, input_data, target_data, ps, st) ``` ## Autoregressive rollout -`predict(esn, steps, ps, st; initialdata)` splits the predict-time -`tspan` into `steps` equal sub-intervals; on each sub-interval the -previous readout output is held constant as the input. The default -initial reservoir state is zeros sized to `res_dims` — if you want to -continue from the trained reservoir's terminal state, pass a custom -`equations` closure that captures it. - ```@example continuous-esn-lorenz ps_pred, st_pred = setup(rng, esn_pred) ps_pred = merge(ps_pred, (readout = ps.readout,)) @@ -147,12 +112,8 @@ with hand-rolled equations, and `ContinuousESN`: the same `train!` / ## When to reach for `ContinuousESN` vs `SciMLProblemReservoir` -* `ContinuousESN` pre-bakes eq (5); reach for it when the standard - leaky-integrator continuous ESN is what you want and you'd otherwise - be hand-rolling the same RHS. -* [`SciMLProblemReservoir`](@ref) is the generic building block; reach - for it when the reservoir ODE is *not* eq (5) — bespoke RHS, SDE, +* `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. - -Both go through the same `_collectstates` / `_predict` dispatch and -the same protected-`saveat` discipline. diff --git a/ext/RCODEReservoirExt.jl b/ext/RCODEReservoirExt.jl index 5f0da9736..529ab3f78 100644 --- a/ext/RCODEReservoirExt.jl +++ b/ext/RCODEReservoirExt.jl @@ -4,12 +4,6 @@ using DataInterpolations: ConstantInterpolation using LinearAlgebra: mul! using LuxCore: apply using Random: AbstractRNG -# `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: ODEProblem, remake, solve, NullParameters using ReservoirComputing: ReservoirComputing, @@ -26,20 +20,6 @@ using ReservoirComputing: ReservoirComputing, 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() @@ -83,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) @@ -143,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 ) @@ -196,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, @@ -272,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, @@ -301,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, @@ -397,19 +297,6 @@ function _predict( return outputs, newst end -# --------------------------------------------------------------------------- -# ContinuousESN — convenience constructor -# -# Builds the 3-field `(reservoir, states_modifiers, readout)` model whose -# `reservoir` is a `ContinuousESNCell`. The cell carries Lukoševičius 2012 -# §3.2.6 eq (5) as its `equations` field. Each of `collectstates`, -# `predict(rc, data, ...)`, and `predict(rc, steps, ...; initialdata)` -# dispatches on `rc.reservoir::ContinuousESNCell` below, building the -# `ODEProblem` lazily so the reservoir weights stay in `ps.reservoir` and -# can be re-initialised by `setup(rng, rc)` like every other ESN-family -# model. -# --------------------------------------------------------------------------- - function ReservoirComputing.ContinuousESN( in_dims::Integer, res_dims::Integer, out_dims::Integer, activation, tspan, args...; @@ -456,8 +343,6 @@ function ReservoirComputing.ContinuousESN( return ContinuousESN(cell, mods, readout) end -# Allow omitting `activation` (defaults to `tanh`) by routing -# `ContinuousESN(in, res, out, tspan, args...)` through the five-arg form. function ReservoirComputing.ContinuousESN( in_dims::Integer, res_dims::Integer, out_dims::Integer, tspan::Union{Tuple, Pair}, args...; kwargs... @@ -465,17 +350,6 @@ function ReservoirComputing.ContinuousESN( return ContinuousESN(in_dims, res_dims, out_dims, tanh, tspan, args...; kwargs...) end -# --------------------------------------------------------------------------- -# Continuous `_collectstates` for `ContinuousESNCell` -# -# Builds the `ODEProblem` on the fly from `cell.equations` and an initial -# state of zeros sized to `cell.out_dims`. The weight matrices come from -# `ps.reservoir` (`input_matrix`, `reservoir_matrix`, optional `bias`) and -# are merged into the solve `p` by `_build_solve_params`. Sampler is fixed -# to `TerminalStateSampling` since eq (5) only exposes a single -# point-state per window. -# --------------------------------------------------------------------------- - function _collectstates( cell::ContinuousESNCell, rc::AbstractReservoirComputer, @@ -532,17 +406,6 @@ function _collectstates( return modified_states, newst end -# --------------------------------------------------------------------------- -# Autoregressive `_predict` for `ContinuousESNCell` -# -# Same control flow as the generic `AbstractSciMLProblemReservoir` path, -# but the `ODEProblem` is built on demand from `cell.equations` rather -# than pulled from a pre-built `res.prob`. Initial reservoir state is -# zeros sized to `cell.out_dims`; users who want to continue from a -# previously trained terminal state should pass it in via a custom -# `equations` closure or re-run `collectstates` first. -# --------------------------------------------------------------------------- - function _predict( cell::ContinuousESNCell, rc::AbstractReservoirComputer, diff --git a/src/layers/continuous_esn_cell.jl b/src/layers/continuous_esn_cell.jl index 8439cc572..96647037a 100644 --- a/src/layers/continuous_esn_cell.jl +++ b/src/layers/continuous_esn_cell.jl @@ -1,17 +1,12 @@ @doc raw""" - ContinuousESNCell(activation, in_dims, out_dims, - init_bias, init_reservoir, init_input, init_state, - use_bias, equations, tspan, args, kwargs) + 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 leaky-integrator ESN cell ([Lukosevicius2012](@cite) §3.2.6, -eq 5). `ContinuousESNCell <: AbstractSciMLProblemReservoir` so it dispatches -into the continuous `collectstates` / `predict` pipeline provided by the -`RCODEReservoirExt` package extension. Field layout mirrors [`ESNCell`](@ref) -plus the ODE right-hand side (`equations`) and the solve metadata -(`tspan`, `args`, `kwargs`) needed by the extension's solver. - -The default `equations` value [`_continuous_esn_rhs!`](@ref) integrates -Lukoševičius 2012 §3.2.6 eq (5): +Continuous-time Echo State Network cell +([Lukosevicius2012](@cite)). Integrates the ODE ```math \dot{\mathbf{x}}(t) = -\mathbf{x}(t) + \tanh\!\left( @@ -19,36 +14,35 @@ Lukoševičius 2012 §3.2.6 eq (5): + \mathbf{W}_r\,\mathbf{x}(t) + \mathbf{b}\right) ``` -No leaking-rate term `α` appears in the ODE — `α` emerges only when the -ODE is forward-Euler discretised with step `Δt = α`, recovering the -discrete leaky update `x(n+1) = (1-α)·x(n) + α·tanh(...)`. - -This type is normally constructed by the [`ContinuousESN`](@ref) -convenience constructor in the `RCODEReservoirExt` extension rather than -directly. Driving the inner constructor by hand requires placing every -field in the order above and pre-static-ing `use_bias`. - -## Fields - - - `activation`: Reservoir nonlinearity. Documented as `tanh` by the - `ContinuousESN` constructor; only enters the cell via a custom - `equations` closure. - - `in_dims`, `out_dims`: Input and reservoir dimensions. - - `init_input`, `init_reservoir`, `init_bias`, `init_state`: - Initialisers (same convention as [`ESNCell`](@ref)). - - `use_bias::StaticBool`: Static-bool flag controlling the optional - `bias` parameter. - - `equations`: ODE right-hand side `(dx, x, p, t) -> nothing`. The - extension injects `p.input(t)` (interpolated input signal) and merges - `ps.reservoir` into `p`. - - `tspan`: Integration interval `(t0, t1)` for `collectstates`. - - `args`, `kwargs`: Forwarded to `solve` positionally / by keyword. +## 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_esn_rhs!`](@ref). + - `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()` + - `bias :: (out_dims,)` — present only if `use_bias = true` """ @concrete struct ContinuousESNCell <: AbstractSciMLProblemReservoir activation @@ -65,6 +59,21 @@ field in the order above and pre-static-ing `use_bias`. 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), @@ -84,7 +93,7 @@ end _continuous_esn_rhs!(dx, x, p, t) Default right-hand side for [`ContinuousESNCell`](@ref). Implements -Lukoševičius 2012 §3.2.6 eq (5) in place: +[Lukosevicius2012](@cite) §3.2.6 eq (5): ```math \dot{\mathbf{x}} = -\mathbf{x} + \tanh\!\left( @@ -92,17 +101,8 @@ Lukoševičius 2012 §3.2.6 eq (5) in place: + \mathbf{W}_r\,\mathbf{x} + \mathbf{b}\right) ``` -`p` is the merged parameter pack assembled by the `RCODEReservoirExt` -helper. It carries: - * `p.input` — interpolated input signal injected by the extension. - * `p.input_matrix` (`W_in`), `p.reservoir_matrix` (`W_r`) — reservoir - matrices pulled from `ps.reservoir`. - * `p.bias` (optional) — bias vector, present iff `use_bias=True()`. - -The RHS uses `mul!` for both matrix products and a fused tanh broadcast -to stay allocation-free on the solver hot path. The bias branch is -selected at compile time via `haskey`, since the keys of a `NamedTuple` -are part of its type. +`p.input(t)` provides the interpolated input signal. The bias term is +included only when `p.bias` is present. """ function _continuous_esn_rhs!(dx, x, p, t) input_t = p.input(t) diff --git a/src/models/continuous_esn.jl b/src/models/continuous_esn.jl index fb9c7d149..22a111be7 100644 --- a/src/models/continuous_esn.jl +++ b/src/models/continuous_esn.jl @@ -1,28 +1,15 @@ @doc raw""" - ContinuousESN(in_dims, res_dims, out_dims, activation = tanh, - tspan, args...; - use_bias = false, - init_reservoir = rand_sparse, init_input = scaled_rand, - init_bias = zeros32, init_state = randn32, - state_modifiers = (), readout_activation = identity, - equations = _continuous_esn_rhs!, - kwargs...) - -Continuous-time leaky-integrator Echo State Network ([Lukosevicius2012](@cite) -§3.2.6, eq 5). Thin convenience wrapper that composes: - - 1) a [`ContinuousESNCell`](@ref) carrying the reservoir matrices and the - ODE right-hand side, - 2) zero or more `state_modifiers` applied to the sampled reservoir - states, and - 3) a [`LinearReadout`](@ref) mapping the modified states to outputs. - -Structurally identical to [`ESN`](@ref) — three fields -`(reservoir, states_modifiers, readout)` — so it slots into the same -training and prediction pipeline. The reservoir field is the -`ContinuousESNCell` rather than a `StatefulLayer(ESNCell)`, which routes -`collectstates` / `predict` through the continuous-time -`RCODEReservoirExt` extension. + 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 @@ -37,19 +24,14 @@ training and prediction pipeline. The reservoir field is the \end{aligned} ``` -No leaking-rate term `α` appears in the ODE. `α` emerges only when the -ODE is forward-Euler discretised with step `Δt = α`, recovering the -discrete leaky update of [`ESN`](@ref). To change the effective leak, -adjust `tspan` so the per-input-window width matches the desired `Δt`. - ## Arguments - `in_dims`: Input dimension. - `res_dims`: Reservoir (hidden state) dimension. - `out_dims`: Output dimension. - - `activation`: Reservoir activation passed to the cell. Default: `tanh`. - - `tspan`: Integration interval `(t0, t1)`. Length-2, strictly - increasing, finite. + - `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. @@ -68,16 +50,17 @@ Reservoir (passed to [`ContinuousESNCell`](@ref)): Composition: - - `state_modifiers`: A layer or collection of layers applied to the - sampled 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`. + - `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 three keys `saveat`, - `save_everystep`, and `dense` are rejected at construction. + - `kwargs...`: Forwarded to `solve`. The keys `saveat`, + `save_everystep`, and `dense` are reserved and rejected at + construction. ## Parameters @@ -85,17 +68,21 @@ Solve metadata: - `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` — a `Tuple` with parameters for each modifier - layer (may be empty). - - `readout` — parameters of [`LinearReadout`](@ref), typically: + - `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 - This constructor errors unless `RCODEReservoirExt` is loaded. Load - `SciMLBase`, `DataInterpolations`, and a concrete `OrdinaryDiffEq` - solver package (e.g. `OrdinaryDiffEqTsit5`, `OrdinaryDiffEq`) to - enable it. + 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)} @@ -104,10 +91,6 @@ Solve metadata: readout end -# Public factory — the real implementation lives in `RCODEReservoirExt`. -# Loading `SciMLBase` + `DataInterpolations` adds the typed method that -# actually builds the cell; without the extension, this fallback fires -# and produces a clear error. function ContinuousESN(::Any, ::Any, ::Any, ::Any...; kwargs...) return error( "ContinuousESN requires the RCODEReservoirExt extension and an " * diff --git a/test/test_continuous_esn.jl b/test/test_continuous_esn.jl index e6afd53e4..4c6bcbfcf 100644 --- a/test/test_continuous_esn.jl +++ b/test/test_continuous_esn.jl @@ -8,15 +8,6 @@ using OrdinaryDiffEqLowOrderRK: Euler using SciMLBase using DataInterpolations -# --------------------------------------------------------------------------- -# 1. Construction: shape + parameter wiring -# -# `ContinuousESN(in_dims, res_dims, out_dims, tspan, solver)` builds a -# 3-field `(reservoir, states_modifiers, readout)` model. The reservoir -# is a `ContinuousESNCell`; `ps.reservoir` exposes the canonical ESN-cell -# parameter trio (`input_matrix`, `reservoir_matrix`, optional `bias`). -# --------------------------------------------------------------------------- - @testset "ContinuousESN: construction + parameter shapes" begin rng = MersenneTwister(0) in_dim, res_dim, out_dim = 3, 50, 2 @@ -37,25 +28,15 @@ using DataInterpolations @test size(ps.readout.weight) == (out_dim, res_dim) end -# --------------------------------------------------------------------------- -# 2. Construction validation: argument errors -# -# Positive `in_dims`, `res_dims`, `out_dims`, and a finite, strictly -# increasing length-2 `tspan` are required. Protected `solve` kwargs are -# rejected at construction. -# --------------------------------------------------------------------------- - @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()) - # Non-finite endpoints @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()) - # Wrong-length tspan @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) @@ -65,12 +46,6 @@ end end end -# --------------------------------------------------------------------------- -# 3. Bias toggle -# -# `use_bias = true` adds a `bias` parameter of length `res_dims`. -# --------------------------------------------------------------------------- - @testset "ContinuousESN: bias toggle" begin rng = MersenneTwister(1) in_dim, res_dim, out_dim = 2, 10, 1 @@ -80,18 +55,9 @@ end ps, st = setup(rng, esn_b) @test haskey(ps.reservoir, :bias) @test length(ps.reservoir.bias) == res_dim - # Default init is `zeros32`. @test all(==(0), ps.reservoir.bias) end -# --------------------------------------------------------------------------- -# 4. Bias actually exercised under an adaptive solver -# -# Pair `use_bias = true` with a nonzero bias initialiser and Tsit5 to -# confirm: (a) the bias does perturb states relative to the no-bias -# baseline, and (b) the result stays finite. -# --------------------------------------------------------------------------- - @testset "ContinuousESN: bias path under Tsit5" begin rng = MersenneTwister(101) in_dim, res_dim, out_dim, T_steps = 2, 20, 1, 12 @@ -115,10 +81,6 @@ end @test s_b != s_n end -# --------------------------------------------------------------------------- -# 5. Forward pass: shape + finiteness + determinism -# --------------------------------------------------------------------------- - @testset "ContinuousESN: forward (collectstates)" begin rng = MersenneTwister(7) in_dim, res_dim, out_dim, T_steps = 2, 16, 1, 25 @@ -136,23 +98,7 @@ end @test s1 ≈ s2 end -# --------------------------------------------------------------------------- -# 6. Euler equivalence with discrete leaky ESN -# -# Eq (5) is `ẋ = -x + tanh(W_r·x + W_in·u(t) + b)`. Forward-Euler with -# step `Δt = α` gives `x_{k+1} = x_k + α·(-x_k + tanh(...)) = -# (1-α)·x_k + α·tanh(...)`, i.e. the discrete leaky ESN update at leak -# rate α. To match `T_steps` discrete steps we set `tspan = (0, T_steps·α)` -# so the per-window width `Δt = α` matches the Euler step `dt = α`. -# --------------------------------------------------------------------------- - @testset "ContinuousESN: Euler equivalence with discrete leaky ESN" begin - # `α` values that are exactly representable in Float64 so that - # `tspan = (0, T_steps · α)` and the integrator step `dt = α` align - # bit-for-bit with the requested `saveat` grid. Non-FP-exact `α` - # (e.g. 0.3) introduces a sub-step offset between saveat and the - # integrator's step boundaries, which the linear interpolant inside - # Euler smooths over — equivalence then holds only approximately. 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...) @@ -186,10 +132,6 @@ end end end -# --------------------------------------------------------------------------- -# 7. Teacher-forced predict + autoregressive predict shape/determinism -# --------------------------------------------------------------------------- - @testset "ContinuousESN: teacher-forced predict" begin rng = MersenneTwister(21) in_dim, res_dim, out_dim, T_steps = 2, 12, 3, 10 @@ -225,10 +167,6 @@ end @test ar1 ≈ ar2 end -# --------------------------------------------------------------------------- -# 8. Compatibility with state modifiers -# --------------------------------------------------------------------------- - @testset "ContinuousESN: state modifiers compose" begin rng = MersenneTwister(31) in_dim, res_dim, out_dim, T_steps = 2, 10, 1, 8 @@ -253,13 +191,6 @@ end @test sm != sp end -# --------------------------------------------------------------------------- -# 9. Custom init eltype propagates -# -# Without a `T` kwarg, the eltype is whatever the `init_*` initialisers -# produce. Passing Float64 initialisers gives Float64 matrices. -# --------------------------------------------------------------------------- - @testset "ContinuousESN: custom init eltype propagates" begin rng = MersenneTwister(43) init_input_f64(rng, d...) = scaled_rand(rng, Float64, d...) From 20d619875a3934d36daf1ea72ab8cf8e66d54d42 Mon Sep 17 00:00:00 2001 From: Saswat Susmoy Date: Tue, 7 Jul 2026 10:24:19 +0530 Subject: [PATCH 8/8] fix: remove broken cross-refs and missing dep in ContinuousESN CI MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Drop — Euler is re-exported by OrdinaryDiffEq and the sub-package is not a declared dependency. - Remove block from internal — it was never listed in any @docs block, causing Documenter missing_docs error. - Replace with descriptive text in ContinuousESNCell and ContinuousESN docstrings. --- src/layers/continuous_esn_cell.jl | 17 +---------------- src/models/continuous_esn.jl | 4 ++-- test/test_continuous_esn.jl | 1 - 3 files changed, 3 insertions(+), 19 deletions(-) diff --git a/src/layers/continuous_esn_cell.jl b/src/layers/continuous_esn_cell.jl index 96647037a..3ec23d34b 100644 --- a/src/layers/continuous_esn_cell.jl +++ b/src/layers/continuous_esn_cell.jl @@ -35,7 +35,7 @@ Continuous-time Echo State Network cell - `init_state`: Initialiser for the initial hidden state. Default: `zeros32`. - `equations`: ODE right-hand side `(dx, x, p, t) -> nothing`. - Default: [`_continuous_esn_rhs!`](@ref). + Default: continuous-time leaky-integrator ESN ODE. - `kwargs...`: Forwarded to `solve` as keyword arguments. ## Parameters @@ -89,21 +89,6 @@ function initialstates(::AbstractRNG, ::ContinuousESNCell) return NamedTuple() end -@doc raw""" - _continuous_esn_rhs!(dx, x, p, t) - -Default right-hand side for [`ContinuousESNCell`](@ref). Implements -[Lukosevicius2012](@cite) §3.2.6 eq (5): - -```math -\dot{\mathbf{x}} = -\mathbf{x} + \tanh\!\left( - \mathbf{W}_{\text{in}}\,\mathbf{u}(t) - + \mathbf{W}_r\,\mathbf{x} + \mathbf{b}\right) -``` - -`p.input(t)` provides the interpolated input signal. The bias term is -included only when `p.bias` is present. -""" function _continuous_esn_rhs!(dx, x, p, t) input_t = p.input(t) mul!(dx, p.reservoir_matrix, x) diff --git a/src/models/continuous_esn.jl b/src/models/continuous_esn.jl index 22a111be7..d5d95a152 100644 --- a/src/models/continuous_esn.jl +++ b/src/models/continuous_esn.jl @@ -45,8 +45,8 @@ Reservoir (passed to [`ContinuousESNCell`](@ref)): - `init_bias`: Initialiser for `b`. Default: `zeros32`. - `init_state`: Initialiser for the initial hidden state. Default: `randn32`. - - `equations`: ODE right-hand side. Default: - [`_continuous_esn_rhs!`](@ref). + - `equations`: ODE right-hand side function `(dx, x, p, t) -> nothing`. + Default is the continuous-time leaky-integrator ESN ODE. Composition: diff --git a/test/test_continuous_esn.jl b/test/test_continuous_esn.jl index 4c6bcbfcf..07dca9157 100644 --- a/test/test_continuous_esn.jl +++ b/test/test_continuous_esn.jl @@ -4,7 +4,6 @@ using LinearAlgebra using Statistics using ReservoirComputing using OrdinaryDiffEq -using OrdinaryDiffEqLowOrderRK: Euler using SciMLBase using DataInterpolations