feat(models): ContinuousESN

[Saswatsusmoy]

Summary

Third deliverable of the SciML Summer Fellowship 2026 continuous-time track (#397), stacked on #446 / #450.

Adds a continuous-time Echo State Network model on the SciML continuous-reservoir path:

• ContinuousESN — 3-field model (reservoir, states_modifiers, readout), sibling of ESN
• ContinuousESNCell — reservoir cell implementing the continuous leaky-integrator ODE
• RCODEReservoirExt — real constructor + _collectstates / _predict dispatch
• Lorenz forecasting tutorial, API docs, and tests

Canonical ODE is Lukoševičius 2012 §3.2.6 eq. (5):

$$\dot{\mathbf{x}}(t) = -\mathbf{x}(t) + \tanh\!\bigl( \mathbf{W}_{\mathrm{in}}\,\mathbf{u}(t) + \mathbf{W}_r\,\mathbf{x}(t) + \mathbf{b}\bigr)$$

Forward-Euler discretisation at step size α recovers the discrete leaky ESN. The discrete leak coefficient is not a free parameter of the continuous RHS; it appears only via the integration step.

Out of scope: the Anantharaman 2021 parametric-surrogate CTESN (ṙ = tanh(Ar + W_hyb x) with RBF-interpolated W_out(p)). This PR is the continuous leaky-integrator model, not that paper's training procedure.

What feature report was addressed?

Part of #397 (continuous-time ESN model + docs; Mackey-Glass deferred / Lorenz tutorial as validation).

Design decisions (settled in review)

Topic	Decision
Name	ContinuousESN / ContinuousESNCell (not CTESN)
Shape	Mirror ESN: 3-field model; ODE substance lives in the cell
Default ODE	Lukoševičius §3.2.6 eq. (5) — no α in the continuous RHS
Placement	src/models/continuous_esn.jl + src/layers/continuous_esn_cell.jl
Validation	Unit tests (incl. Euler ↔ discrete leaky ESN equivalence) + Lorenz attractor tutorial
Warmup / cold u0 on autoregressive predict	Deferred to a follow-up

Usage sketch

using ReservoirComputing, OrdinaryDiffEq, Random

esn = ContinuousESN(3, 300, 3, (0.0, 100.0), Tsit5();
    use_bias = true)
ps, st = setup(MersenneTwister(42), esn)
# train! / predict as with other models

Requires the continuous-time extension (load OrdinaryDiffEq / related SciML stack so RCODEReservoirExt activates).

Checklist

• Appropriate tests were added
• Any code changes were done in a way that does not break public API
• All documentation related to code changes were updated
• The new code follows the
  contributor guidelines, in particular the SciML Style Guide and
  COLPRAC.
• Any new documentation only uses public API

Test plan

• Construction / shape / tspan validation
• Bias on/off (incl. under Tsit5)
• Forward determinism
• Euler equivalence vs discrete leaky ESN at α ∈ {1.0, 0.5}
• Teacher-forced and autoregressive predict
• State modifiers compose with the continuous path
• Lorenz tutorial (docs/src/tutorials/continuous_esn.md)
• API docs for model + cell (models.md / layers.md)
• CI green at merge

Files

Path	Role
src/layers/continuous_esn_cell.jl	ContinuousESNCell
src/models/continuous_esn.jl	ContinuousESN model + stub without extension
ext/RCODEReservoirExt.jl	Constructor, RHS, continuous collect/predict
test/test_continuous_esn.jl	Construction, bias, determinism, Euler equivalence, predict modes
docs/src/tutorials/continuous_esn.md	Lorenz walkthrough
docs/src/api/{models,layers}.md	@docs entries
docs/src/refs.bib	Lukoševičius 2012

Follow-ups (not in this PR)

• Warmup / terminal-state seeding for autoregressive continuous predict (discussed in thread)
• Anantharaman-style parametric-surrogate CTESN
• Stiff-system benchmarks (Robertson, etc.)

Additional context

PR3 of the continuous-time stack after #446 (core types) and #450 (RCODEReservoirExt). Design discussion lived in this PR thread; implementation was refactored mid-review to the final 3-field + ContinuousESNCell shape.

[Saswatsusmoy]

Thanks @MartinuzziFrancesco — both the §3.2.6 pointer and the "skip the struct, run it raw" suggestion turned out to land in interesting places.

On §3.2.6 (Lukoševičius 2012, "Leaking Rate")

The relevant ODE is eq (5):

$$\dot{\mathbf{x}} = -\mathbf{x} + \tanh\!\left(\mathbf{W}^{\text{in}}[1;\mathbf{u}] + \mathbf{W}\,\mathbf{x}\right)$$

— no α in the continuous form. The leaking rate appears only when you Euler-discretise: take Δx/Δt ≈ ẋ with step Δt = α and you recover the discrete leaky update x(n+1) = (1-α)x(n) + α·tanh(W_in[1;u] + Wx(n)). So α is the discretisation step, not a separate parameter of the ODE. Equivalently you can fold it into eq (5) as a LHS time constant α·ẋ = … if you want to keep Δt ≡ 1. The section also notes this form is preferred over alternatives "because it guarantees that x(n) never goes outside the (-1, 1) interval" — bounded state by construction, no extra contraction argument needed. Practical advice: tune α to the timescale of u/y_target; can be per-neuron for multi-timescale tasks.

This reframes PR3 quite a bit. The Anantharaman CTESN paper uses ṙ = tanh(Ar + W_hyb x) with no decay term and trains a parametric surrogate via RBF interpolation of W_out(p) — a different approach aimed at stiff systems. The "continuous ESN" §3.2.6 points at is just eq (5), which is exactly what the Lorenz eye-test in PR #450's tutorial implements verbatim. So PR3 collapses from "build a CTESN model mirroring the paper" → "thin convenience constructor around SciMLProblemReservoir that pre-bakes eq (5)". Several of the open questions in the original PR body fall away.

Smoke test — eq (5) raw on Lorenz, no CTESN struct

Ran the tutorial pattern scaled to N_res = 300, train = 5000, predict = 1250 samples (dt = 0.02, σ = 10, ρ = 28, β = 8/3). Spectral scaling ‖W_r‖₂ = 0.9 (continuous-ESP-sufficient bound), ‖W_in‖ ≈ 0.1, ‖b‖ ≈ 0.05, MersenneTwister(17), Tsit5 at reltol = 1e-6. Happy to inline the full script as a follow-up comment if useful.

First pass, train NRMSE = 0.0004 (readout fits beautifully) but autoregressive rollout NRMSE was ~1.5 at all horizons. Tracked it down to the autoregressive _predict path in ext/RCODEReservoirExt.jl:

current_state = res.prob.u0   # → zeros(N) every time

The trained reservoir's terminal state never makes it into the predict-time u0. So the rollout starts cold and produces nothing useful regardless of how well the readout was trained. After manually re-running collectstates, grabbing the last column, and rebuilding the reservoir with remake(prob; u0 = terminal_state), the numbers became:

Horizon (Lyapunov times)	Steps	NRMSE
1 t_λ	55	0.157
2 t_λ	110	0.118
3 t_λ	166	0.112
4 t_λ	221	0.210
6 t_λ	331	0.729
8 t_λ	442	0.981

So eq (5) forecasts Lorenz at N=300 to ~3 Lyapunov times under NRMSE ≈ 0.11 before chaotic divergence dominates. Not quite our PR3 target of NRMSE < 0.1, but within striking distance of it before any hyperparameter tuning (spectral radius vs spectral norm, leak/time-constant, washout length, regularisation). I think the < 0.1 target is reachable.

What I'd like your read on:

1. The warm-up hole. Is this expected for the current PR feat: continuous-time reservoirs via RCODEReservoirExt #450 predict semantics — i.e., users are meant to remake(prob; u0 = …) themselves before calling predict? Or is this something that should be wired in (e.g. predict(rc, steps, ps, st; initialdata, warmup_data = nothing) that internally runs collectstates on warmup_data and seeds u0)? The current behaviour is correct but really easy to footgun on, as I just demonstrated.

2. PR3 reframe. Given §3.2.6, would you rather PR3 be:

   • (a) Drop the "CTESN" name entirely and ship a LeakyContinuousESN(in_dims, res_dims, out_dims; α, …) thin wrapper that pre-bakes eq (5) — semantically a sibling of ESN. The Anantharaman parametric-surrogate flavour would be a future PR with its own type.
   • (b) Keep the CTESN name but make it the eq-(5) wrapper, and treat the parametric-surrogate version as a future extension within the same family.
   • (c) Something else entirely.

3. Validation. Should the in-PR test assert specifically against the eq-(5) Lorenz NRMSE (after tuning to push below 0.1 at some agreed horizon — e.g. 4 t_λ), or do you want a tighter set of unit tests + Mackey-Glass + Lorenz in the docs as eye-tests only?

Will hold off on writing real code until you've had a chance to push back on any of this — easier to redo the design now than after I've committed 600 lines.

[MartinuzziFrancesco]

Thanks for running the experiments, they are quite interesting and I think necessary even for merging #450. To address you point number 1) I think having the two code snippets side to side would help in understanding if this is something that could be a gotcha for the user, and therefore if we should find a way to deal with it in a elegant manner. As far as 2), I think we can go with ContinuousESN thin wrapper of eq 5). And as far as validation I think we can go with the eye test of attractor reconstruction following the example of the readme. If you want a more qualitative test we can use the difference in fractal dimensions using FractalDimensions's Grassberger-Procaccia algorithm. I would say that this is a bit more solid wrt to vpt, and it should be enough for PR3

[Saswatsusmoy]

Thanks @MartinuzziFrancesco. Re-ran both paths from scratch for the side-by-side, and lining up the other two answers below.

Cold vs warmed predict — side-by-side

Both scripts are identical except for the warmup step. Reproducible numbers on Julia 1.12.5 with MersenneTwister(17), N = 300, train = 5000, predict = 1250, dt = 0.02, Tsit5 at reltol = 1e-6, ‖W_r‖₂ = 0.9, ‖W_in‖ ≈ 0.1.

lorenz_cold.jl — predict's u0 is whatever the user gave the ODEProblem (here zeros(N))

import Pkg; Pkg.activate("/tmp/lorenz_env"; io = devnull)
using Random, Statistics, LinearAlgebra, Printf
using ReservoirComputing, SciMLBase, DataInterpolations, OrdinaryDiffEqTsit5

const SEED, N_RES = 17, 300
const TRAIN_LEN, PREDICT_LEN, DT, SHIFT = 5000, 1250, 0.02, 300
const T_LYAPUNOV = 1 / 0.9056

lorenz!(du, u, p, t) = begin
    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, (SHIFT + TRAIN_LEN + PREDICT_LEN + 100) * DT), [10.0, 28.0, 8 / 3])
data = Array(solve(data_prob, Tsit5(); saveat = DT))
input_data  = data[:, SHIFT:(SHIFT + TRAIN_LEN - 1)]
target_data = data[:, (SHIFT + 1):(SHIFT + TRAIN_LEN)]
test_data   = data[:, (SHIFT + TRAIN_LEN):(SHIFT + TRAIN_LEN + PREDICT_LEN - 1)]

rng = MersenneTwister(SEED); Random.seed!(SEED)
Wr_raw = randn(rng, N_RES, N_RES) ./ sqrt(N_RES)
Wr  = (0.9 / opnorm(Wr_raw)) .* Wr_raw
Win = 0.1 .* randn(rng, N_RES, 3)
b   = 0.05 .* randn(rng, N_RES)

esn_rhs!(dx, x, p, t) = (dx .= .-x .+ tanh.(p.Wr * x .+ p.Win * p.input(t) .+ p.b); nothing)

function build_rc(tspan_len, u0)
    prob = ODEProblem(esn_rhs!, u0, (0.0, Float64(tspan_len)), (Wr = Wr, Win = Win, b = b))
    res = SciMLProblemReservoir(prob, TerminalStateSampling(),
        (0.0, Float64(tspan_len)), Tsit5(); reltol = 1.0e-6, abstol = 1.0e-8)
    return ReservoirComputer(res, LinearReadout(N_RES => 3))
end

rc_train = build_rc(TRAIN_LEN, zeros(N_RES))
ps, st = setup(rng, rc_train)
ps, st = train!(rc_train, input_data, target_data, ps, st)

# COLD: predict's u0 stays at the prob's original u0 (zeros).
rc_pred = build_rc(PREDICT_LEN, zeros(N_RES))
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_data[:, 1])

nrmse(yhat, y) = mean(sqrt(mean((yhat[i, :] .- y[i, :]) .^ 2)) / std(y[i, :]) for i in 1:size(y, 1))
for lt in (1.0, 2.0, 3.0, 4.0, 6.0, 8.0)
    h = min(round(Int, lt * T_LYAPUNOV / DT), PREDICT_LEN)
    @printf("  horizon %.1f t_λ  (%4d steps): NRMSE = %.4f\n", lt, h, nrmse(output[:, 1:h], test_data[:, 1:h]))
end

lorenz_warm.jl — same script, with the warmup inserted before build_rc(PREDICT_LEN, …)

# … everything above identical to lorenz_cold.jl through `train!` …

# WARMUP: re-run collectstates, grab terminal reservoir state.
trained_states, _ = collectstates(rc_train, input_data, ps, st)
warm_u0 = collect(trained_states[:, end])

# WARMED: predict's u0 = trained reservoir terminal state.
rc_pred = build_rc(PREDICT_LEN, warm_u0)
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_data[:, 1])

# … NRMSE loop identical to lorenz_cold.jl …

The semantic delta is exactly: collectstates once after train!, take trained_states[:, end], rebuild the predict reservoir with that as u0. Four added lines, one swapped line.

Results (autoregressive rollout NRMSE per Lukoševičius 2012 eq 1, mean over channels of RMSE / std):

Horizon	Cold (u0 = 0)	Warmed (u0 = trained terminal)
1 t_λ (55 steps)	1.4893	0.1568
2 t_λ (110 steps)	1.4960	0.1183
3 t_λ (166 steps)	1.7431	0.1118
4 t_λ (221 steps)	1.6959	0.2099
6 t_λ (331 steps)	1.6444	0.7294
8 t_λ (442 steps)	1.4708	0.9808

Train NRMSE (teacher-forced one-step on training data) was 0.0004 in both cases — the readout fits cleanly; the cold path's failure is purely the predict-time u0.

Two small caveats on what I ran:

• Eq (5) bundles bias as the leading column of W_in[1;u]; my script splits it into a separate +b term — algebraically identical, just different parameter layout.
• warm_u0 = trained_states[:, end] is the post-modifier state if any state_modifiers are wired up. None here, so it equals the raw reservoir terminal. If a real warmup API were to expose this, "warm from the raw reservoir terminal" and "warm from the modified terminal" would need to be a deliberate choice.

If you read this as a real footgun, the lightest-touch fix I can see is a warmup_data::Union{Nothing, AbstractMatrix} = nothing kwarg on the SciMLProblemReservoir paths of predict, where a non-nothing value triggers collectstates(rc, warmup_data, ps, st) internally and seeds u0 from the last column. No behaviour change for users who don't pass it. Happy to follow your lead on whether this lands in #450 or in PR3's wrapper.

On the naming

ContinuousESN works for me. I'll rename the stub on add-ctesn → new branch + new file + re-titled PR. The Anantharaman parametric-surrogate flavour can be its own type in its own PR if anyone needs it.

On validation via Grassberger-Procaccia

Will pull in FractalDimensions.jl (v1.9.6 as of writing). Plan is grassberger_proccacia_dim(StateSpaceSet(...)) on both the true Lorenz attractor and the warmed-rollout attractor, then assert the predicted dimension matches reference (Lorenz '63 correlation dimension is reported at ≈ 2.05 in the literature) to some agreed tolerance, alongside the attractor-reconstruction eye-test from the README. I won't quote actual fractal-dim numbers from my run until I've actually computed them — flagging the plan rather than the result.

[MartinuzziFrancesco]

this can go on its own page since it's a model

[MartinuzziFrancesco]

not needed if you ahve the example below

[Saswatsusmoy]

Both done in 7a7a84b5. Moved ContinuousESN into a new "Continuous-time Echo State Networks" section in api/models.md — agreed, it's a model not a layer. Dropped the standalone using block in the tutorial and folded the one-liner about which packages do what into the intro paragraph above the @example block.

[MartinuzziFrancesco]

too much here, the ContinuousESN should just hold the reservoir, state modifires, and readout. the construction is similar to the ESN, we just provide a convenient wrapper for the user bu we use the internal components. so the the esn equations are fed into the SciMLProblemReservoir, the dims into the ESNCell etcetc

[Saswatsusmoy]

Happy to restructure to the 3-field shape — want to make sure I land the composition right before I rewrite. Three things I'd like to confirm:

1. Where does the cell sit? I see two readings: (a) reservoir::SciMLProblemReservoir with ESNCell used only at construction time to mint prob.p, or (b) introduce a small ContinuousReservoir <: AbstractSciMLProblemReservoir that holds the ESNCell as a field and delegates initialparameters to it (so weights stay in ps.reservoir and the existing _build_solve_params flow keeps working). I lean (b) since (a) bakes the weights into prob.p at construction and they'd no longer live in ps. Which were you picturing?
2. leak_coefficient — agreed eq (5) doesn't carry α; α is the Euler step. Want me to drop the leak_coefficient kwarg entirely (user controls α implicitly via tspan / n_samples), or keep it as a thin convenience that adjusts tspan under the hood?
3. ESNCell init signatures — ESNCell's initializers take (rng, dims...) (no eltype), but the current ContinuousESN exposes T and passes (rng, T, dims...). To reuse ESNCell as-is I'd lose the T kwarg and inherit ESNCell's Float32 default. Acceptable, or should I keep T?

[MartinuzziFrancesco]

I see this raises a good point on adding the ContinuousESNCell as an additional layer. This way we can save those elements in the new cell, which should then also contain the esn equations. something like this maybe

@concrete struct ContinuousESNESNCell <: AbstractEchoStateNetworkCell #maybe change here
    activation
    in_dims <: IntegerType
    out_dims <: IntegerType
    init_bias
    init_reservoir
    init_input
    #init_feedback::F
    init_state
    use_bias <: StaticBool
    equations
end

also make sure to use IntegerType and StaticBool for Int and Bools

[Saswatsusmoy]

Pushed d9b5fda9 with the restructure. ContinuousESN is now 3-field (reservoir, states_modifiers, readout) exactly like ESN; ContinuousESNCell <: AbstractSciMLProblemReservoir so the existing continuous dispatch in RCODEReservoirExt fires on rc.reservoir. initialparameters returns the canonical (input_matrix, reservoir_matrix, [bias]) trio matching ESNCell. Default equations is eq (5) without α — α emerges purely from the Euler step Δt, per your correction.

One field-list deviation worth flagging — the cell carries three more fields than you sketched (tspan, args, kwargs). The reason: since the cell is the dispatch target for _collectstates / _predict, the ext needs cell.tspan to set up the saveat grid and cell.args / cell.kwargs to call solve. If those lived on ContinuousESN instead, the cell would still need a way to reach them at dispatch time, which would mean either threading them through every _collectstates call or duplicating dispatch on the outer model. Adding them to the cell felt like the lowest-friction option, but happy to move them onto ContinuousESN and reach for them via rc.tspan etc. if you'd rather keep the cell strictly ESNCell-shaped — just say the word.

Also renamed u_t → input_t per your other comment; equivalence test now uses Float-exact α ({1.0, 0.5}) so the 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).

[MartinuzziFrancesco]

try not to use single letter variables

[MartinuzziFrancesco]

Lukoševičius 2012 §3.2.6 eq (5) does not have alpha in it, it comes out as a consequence of discretization of the ode

[MartinuzziFrancesco]

I would also put this in a different subsection, similar to the models. Put it in a Continuous layers section

[MartinuzziFrancesco]

leaky-integrator

[MartinuzziFrancesco]

is this necessary now? I think the implementation we have is not the one provided by this paper

[MartinuzziFrancesco]

comments not really useful here

[MartinuzziFrancesco]

not needed

[MartinuzziFrancesco]

too low level information for a docstring

[MartinuzziFrancesco]

we still need a constructor here like for the normal ESNCell

[MartinuzziFrancesco]

§3.2.6, eq 5)

[MartinuzziFrancesco]

Structurally identical to ESN, but the reservoir field is the
ContinuousESNCell rather than a StatefulLayer(ESNCell)

[MartinuzziFrancesco]

there are too many implementation level details that the user may not be familiar with or not care about

[MartinuzziFrancesco]

not needed

[MartinuzziFrancesco]

I think we are almost there! There are a few points to clarify:

• the documentation, especially docstrings, has to be clear to human users. This means that details about the routing or the dispatch or how things work under the hood are not needed. When I use + I am not interested in how the dispatch was made but in what it does for my use case.
• Similarly, too many details on sources can be a bit much. It suffices to say which paper the model comes from using the citation.
• Comments are not needed, the code usually should be clear enough to roughly understand what is going on. I would remove most of the comments there are

Also, we do need a convenience constructor for the ContinuousESNCell.

Overall it's ok to use llms for coding, but make sure the code is up to standard. Forcing the SciML guide on a prompt usually helps in catching a lot of these issues, as well as having a local .md similar to this one

[Saswatsusmoy]

I will have to look into this seriously, I was trying and experimenting documentations and code styling changes through autonomous agents (mostly claude & hermes) will have to fine tune the agents (will force the SciML guide on a project level). For now I will fallback to the prev workflow (without LLM agents) for the PR

[Saswatsusmoy]

I kept _continuous_esn_rhs! unexported. It is documented and used as a default kwarg (equations = _continuous_esn_rhs!), so it shows up in signatures and @ref links, but the _ prefix is a Julia convention for internal and no other _-prefixed helper in the codebase is exported. Fine as is, or would you rather export it?

[Saswatsusmoy]

Coming back to the warmup discussion. The cold vs warm autoregressive predict gap is still there, and the user currently has to manually collectstates after training and rebuild the reservoir with the terminal state as u0. In the side-by-side scripts the gap was NRMSE ~1.5 (cold) vs ~0.11 (warmed).

Do you want me to wire in automatic warmup as part of this PR, or keep it as-is and defer to a separate PR? Two lightweight approaches on the table:

• Add a warmup_data kwarg to predict(rc, steps, ...) for the continuous paths, where a non-nothing value triggers collectstates internally and seeds u0 from the last column. No behaviour change for users who do not pass it.

• Keep the current manual workflow as the documented path and just note the cold/warm distinction in the tutorial.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review is ineligible. To be eligible to request a review, you need a paid Copilot license, or your organization must enable Copilot code review.

[MartinuzziFrancesco]

I think we can add the warmup as part of a separate PR, since I would like to investigate that a bit further. As far as this pr goes I think we can merge, so if you want up the version as usual when tests pass this should be ready to go