From 881018f43c829013428009b85c5b7f1c93a79014 Mon Sep 17 00:00:00 2001 From: Saswat Susmoy Date: Thu, 13 Aug 2026 21:35:13 +0530 Subject: [PATCH] fix(ext): persist continuous ODE state in st.carry Discrete models already keep the terminal reservoir state in st after collect/train. Continuous was restarting AR from zeros / prob.u0. Write the raw ODE terminal into st.reservoir.carry and seed the next solve from it. resetcarry! already understands that shape. --- ext/RCODEReservoirExt.jl | 46 +++++++++++++++------- test/Extensions/continuous_esn_tests.jl | 38 ++++++++++++++++++ test/Extensions/lsm_tests.jl | 26 +++++++++++- test/Extensions/ode_reservoir_ext_tests.jl | 27 +++++++++++++ 4 files changed, 121 insertions(+), 16 deletions(-) diff --git a/ext/RCODEReservoirExt.jl b/ext/RCODEReservoirExt.jl index 48e9e5457..436fff240 100644 --- a/ext/RCODEReservoirExt.jl +++ b/ext/RCODEReservoirExt.jl @@ -138,6 +138,14 @@ function __sample(::TerminalStateSampling, sol) return reduce(hcat, sol.u) end +# Copy so later `reinit!` / `solve` cannot mutate `st`. +@inline function __seed_u(st_res, default) + c = get(st_res, :carry, nothing) + return c === nothing ? default() : copy(first(c)) +end + +@inline __with_carry(st_res, u) = merge(st_res, (; carry = (copy(u),))) + function __apply_modifiers_continuous( modifiers::Tuple, states_matrix::AbstractMatrix, ps_mods, st_mods ) @@ -194,7 +202,12 @@ function __collectstates( input_interp = __make_input_fn(data, input_ts) solve_p = __build_solve_params(res.prob.p, ps.reservoir, input_interp) - prob_remade = remake(res.prob; tspan = res.tspan, p = solve_p) + prob_remade = remake( + res.prob; + tspan = res.tspan, + p = solve_p, + u0 = __seed_u(st.reservoir, () -> res.prob.u0), + ) sol = solve( prob_remade, res.args...; @@ -210,7 +223,7 @@ function __collectstates( ) newst = ( - reservoir = st.reservoir, + reservoir = __with_carry(st.reservoir, sol.u[end]), state_modifiers = st_mods, readout = st.readout, ) @@ -259,12 +272,7 @@ function __predict( window_starts = @view ts[1:(end - 1)] window_ends = @view ts[2:end] - # Preserve `u0`'s original type — `collect` would degrade `SVector` / - # `ComponentArray` / scalar states into a plain `Vector` and either - # error (no `collect(::Number)` method) or silently flatten the - # user's chosen representation. We only ever read `current_state`, - # never mutate it in place, so a direct reference is safe. - current_state = res.prob.u0 + current_state = __seed_u(st.reservoir, () -> res.prob.u0) current_input = initialdata st_mods = st.state_modifiers @@ -320,7 +328,7 @@ function __predict( end newst = ( - reservoir = st.reservoir, + reservoir = __with_carry(st.reservoir, current_state), state_modifiers = st_mods, readout = st_ro, ) @@ -414,7 +422,9 @@ function __collectstates( # `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) + u0 = __seed_u( + st.reservoir, () -> zeros(eltype(ps.reservoir.input_matrix), cell.out_dims) + ) jac_prototype = known(cell.use_jac_prototype) ? __reservoir_jac_prototype(cell.equations, ps.reservoir.reservoir_matrix) : nothing @@ -435,7 +445,7 @@ function __collectstates( ) newst = ( - reservoir = st.reservoir, + reservoir = __with_carry(st.reservoir, sol.u[end]), state_modifiers = st_mods, readout = st.readout, ) @@ -463,7 +473,9 @@ function __predict( 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_state = __seed_u( + st.reservoir, () -> zeros(eltype(ps.reservoir.input_matrix), cell.out_dims) + ) current_input = initialdata st_mods = st.state_modifiers @@ -516,7 +528,7 @@ function __predict( end newst = ( - reservoir = st.reservoir, + reservoir = __with_carry(st.reservoir, current_state), state_modifiers = st_mods, readout = st_ro, ) @@ -758,6 +770,8 @@ function __lsm_pack(cell::LSMCell, ps_res, input_fn, n_units::Int, ::Type{T}) wh end function __lsm_u0(cell::LSMCell, st_res, n_units::Int, ::Type{T}) where {T} + c = get(st_res, :carry, nothing) + c === nothing || return copy(first(c)), st_res rng = replicate(st_res.rng) u0 = zeros(T, 2 * n_units) copyto!(view(u0, 1:n_units), vec(cell.init_state(rng, n_units, 1))) @@ -835,7 +849,9 @@ function __collectstates( rc.state_modifiers, features, ps.state_modifiers, st.state_modifiers ) newst = ( - reservoir = merge(st_res_new, (encoder = st_enc_new,)), + reservoir = __with_carry( + merge(st_res_new, (encoder = st_enc_new,)), sol.u[end] + ), state_modifiers = st_mods, readout = st.readout, ) @@ -944,7 +960,7 @@ function __predict( end newst = ( - reservoir = st_res_new, + reservoir = __with_carry(st_res_new, current_state), state_modifiers = st_mods, readout = st_ro, ) diff --git a/test/Extensions/continuous_esn_tests.jl b/test/Extensions/continuous_esn_tests.jl index 834b03083..2d000829c 100644 --- a/test/Extensions/continuous_esn_tests.jl +++ b/test/Extensions/continuous_esn_tests.jl @@ -261,4 +261,42 @@ begin @test eltype(ps.reservoir.reservoir_matrix) == Float64 end + @testset "ContinuousESN: autoregressive predict uses carry" begin + rng = MersenneTwister(99) + dim, res_dim, T_steps, steps = 2, 16, 20, 5 + esn = ContinuousESN( + dim, res_dim, dim, (0.0, 2.0), Tsit5(); + reltol = 1.0e-8, abstol = 1.0e-10, + ) + esn_mod = ContinuousESN( + dim, res_dim, dim, (0.0, 2.0), Tsit5(); + reltol = 1.0e-8, abstol = 1.0e-10, + state_modifiers = (NLAT2(),), + ) + ps, st0 = setup(rng, esn) + data = randn(Float32, dim, T_steps) + init = data[:, end] + + states, st1 = collectstates(esn, data, ps, st0) + @test first(st1.reservoir.carry) ≈ states[:, end] + + cold, _ = predict(esn, steps, ps, st0; initialdata = init) + warm, st2 = predict(esn, steps, ps, st1; initialdata = init) + @test cold ≉ warm + @test first(st2.reservoir.carry) ≉ first(st1.reservoir.carry) + + st_clear = resetcarry!(MersenneTwister(0), esn, st1) + @test get(st_clear.reservoir, :carry, nothing) === nothing + cold2, _ = predict(esn, steps, ps, st_clear; initialdata = init) + @test cold2 ≈ cold + + s_cont, _ = collectstates(esn, data, ps, st1) + @test s_cont ≉ states + + ps_m, st_m = setup(MersenneTwister(0), esn_mod) + sm, stm = collectstates(esn_mod, data, ps_m, st_m) + @test first(stm.reservoir.carry) ≉ sm[:, end] + @test length(first(stm.reservoir.carry)) == res_dim + end + end diff --git a/test/Extensions/lsm_tests.jl b/test/Extensions/lsm_tests.jl index 08a96b842..e2cfa4dac 100644 --- a/test/Extensions/lsm_tests.jl +++ b/test/Extensions/lsm_tests.jl @@ -323,7 +323,7 @@ begin y[:, t] .= 0.6 .* u[:, t] .+ 0.4 .* reverse(u[:, t - 1]) end - states, st = collectstates(lsm, u, ps, st) + states, _ = collectstates(lsm, u, ps, st) @test mean(abs, states) > 0.05 nrmse0 = norm(zeros(n_out, T_steps) .- y) / norm(y .- mean(y; dims = 2)) @@ -375,6 +375,30 @@ begin @test all(isfinite, a1) end + @testset "AR predict uses carry" begin + n, T_steps, steps = 8, 12, 4 + lsm = _lsm_f64( + 1, n, 1, (0.0, 0.03), Tsit5(); + neuron = LIFNeuron(; tau_m = 0.02, tau_ref = 0.002, tau_syn = 0.008), + feature_map = MembraneVoltageFeature(), + reltol = 1.0e-7, abstol = 1.0e-9, dtmax = 5.0e-4, + ) + ps, st0 = setup(MersenneTwister(3), lsm) + data = ones(Float64, 1, T_steps) + states, st1 = collectstates(lsm, data, ps, st0) + u_end = first(st1.reservoir.carry) + @test length(u_end) == 2n + @test u_end[1:n] ≈ states[:, end] + + s_cont, _ = collectstates(lsm, data, ps, st1) + @test s_cont ≉ states + + init = [1.0] + cold, _ = predict(lsm, steps, ps, st0; initialdata = init) + warm, _ = predict(lsm, steps, ps, st1; initialdata = init) + @test cold ≉ warm + end + @testset "PoissonRateEncoder: event drive produces spikes" begin n_win = 40 T_end = 2.0 diff --git a/test/Extensions/ode_reservoir_ext_tests.jl b/test/Extensions/ode_reservoir_ext_tests.jl index 99b348799..ded07ee39 100644 --- a/test/Extensions/ode_reservoir_ext_tests.jl +++ b/test/Extensions/ode_reservoir_ext_tests.jl @@ -212,6 +212,33 @@ begin @test preds1 ≈ preds2 end + @testset "Autoregressive predict uses carry" begin + rng = MersenneTwister(37) + res_dim, dim, T_steps, steps = 8, 2, 16, 4 + tspan = (0.0, 2.0) + prob, _, _, _, _ = build_esn_problem(rng, dim, res_dim, tspan) + prob = remake(prob; u0 = 0.3 .* randn(rng, res_dim)) + res = SciMLProblemReservoir( + prob, TerminalStateSampling(), tspan, Tsit5(); + reltol = 1.0e-8, abstol = 1.0e-10, + ) + rc = ReservoirComputer(res, LinearReadout(res_dim => dim)) + ps, st0 = setup(MersenneTwister(0), rc) + data = randn(rng, dim, T_steps) + init = data[:, end] + + states, st1 = collectstates(rc, data, ps, st0) + @test first(st1.reservoir.carry) ≈ states[:, end] + + cold, _ = predict(rc, steps, ps, st0; initialdata = init) + warm, _ = predict(rc, steps, ps, st1; initialdata = init) + @test cold ≉ warm + + st_forced = merge(st0, (reservoir = (; carry = (copy(prob.u0),)),)) + from_u0, _ = predict(rc, steps, ps, st_forced; initialdata = init) + @test from_u0 ≈ cold + end + # --------------------------------------------------------------------------- # 6. State modifiers compose with the continuous path #