Skip to content

fix(ext): persist continuous ODE state in st.carry #499

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
MartinuzziFrancesco merged 1 commit into from
Aug 17, 2026
+121 −16
Merged

fix(ext): persist continuous ODE state in st.carry #499

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 31 additions & 15 deletions ext/RCODEReservoirExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
Expand Down Expand Up @@ -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...;
Expand All @@ -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,
)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
)
Expand Down Expand Up @@ -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
Expand All @@ -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,
)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
)
Expand Down Expand Up @@ -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)))
Expand Down Expand Up @@ -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]
),
Comment on lines +852 to +854

@chatgpt-codex-connector chatgpt-codex-connector Bot Aug 13, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Persist LSM refractory state with the carry

When an LSM solve is resumed from this new carry, only sol.u[end] is saved, but the spike callback also mutates p.ref_until to enforce tau_ref; __lsm_pack recreates ref_until = typemin for the next solve. If a neuron spikes near the end of collectstates/AR and the returned st is used for the next call, that neuron is no longer refractory even though an uninterrupted solve would keep du[unit] = 0 until t_spike + tau_ref, so warm-started LSM dynamics and features can be wrong exactly in the spiking/refractory case this carry is meant to preserve.

Useful? React with 👍 / 👎.

@Saswatsusmoy Saswatsusmoy Aug 13, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ref_until is callback state on p, not part of the ODE u this PR is threading. also the next collect/AR restarts tspan, so copying the old absolute times would keep units refractory for the whole next interval. leftover tau_ref across calls is a fair lsm follow-up (same bucket as the exp filter). not changing this PR for it.

state_modifiers = st_mods,
readout = st.readout,
)
Expand Down Expand Up @@ -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,
)
Expand Down
38 changes: 38 additions & 0 deletions test/Extensions/continuous_esn_tests.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
26 changes: 25 additions & 1 deletion test/Extensions/lsm_tests.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down Expand Up @@ -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
Expand Down
27 changes: 27 additions & 0 deletions test/Extensions/ode_reservoir_ext_tests.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
#
Expand Down
Loading