Problem
When using SolverStrangThreads (or SolverStrangSerial), prob.f.sys is Nothing after building the ODEProblem. This makes it impossible to programmatically determine which state index corresponds to which variable.
This is critical for two use cases:
- Restart from saved output: mapping each species' saved concentration to the correct state index
- Custom initial conditions: any code that sets per-species IC needs the species-to-index mapping
Root cause
In solver_strategy_strang.jl, the ODEProblem is constructed by passing nonstiff_op (a raw operator) directly:
ODEProblem(nonstiff_op, view(u0, :), (start, finish), p; ...)
Since nonstiff_op is not an ODEFunction, ODEProblem wraps it in a new ODEFunction with sys=nothing. The sys_mtk system is already available in scope but not attached.
Why unknowns(convert(System, model)) is not a workaround
The pre-simplification ordering from unknowns(convert(System, model)) does not match the actual state ordering in prob.u0. For example, with a 7-species surrogate chemistry model:
unknowns(convert(System, model)): [c7, c6, c5, c4, c3, c2, c1] # pre-simplification
Actual prob.u0 ordering: [c1, c2, c3, c4, c5, c6, c7] # after mtkcompile
Using unknowns(convert(System, model)) to assign initial conditions results in species getting each other's data (e.g., O3 gets NO2's IC and vice versa).
Proposed fix (1 line)
Wrap nonstiff_op in ODEFunction with sys=sys_mtk before passing to ODEProblem:
# Before:
ODEProblem(nonstiff_op, view(u0, :), (start, finish), p; ...)
# After:
nonstiff_fn = ODEFunction(nonstiff_op; sys = sys_mtk)
ODEProblem(nonstiff_fn, view(u0, :), (start, finish), p; ...)
This has zero runtime cost -- ODEProblem already wraps raw operators in ODEFunction internally. The fix simply makes it explicit and attaches the symbolic system metadata.
Usage example: restart from NetCDF checkpoint
With the fix, users can do:
prob = ODEProblem(model, SolverStrangThreads(Rodas5P(), 300f0))
vars = unknowns(prob.f.sys) # correct ordering matching prob.u0
ds = NCDataset("previous_run_output.nc")
n_species = length(vars)
init_cond = zeros(n_species, nx, ny, nz)
for (state_idx, var) in enumerate(vars)
nc_key = replace(string(var), "(t)" => "")
init_cond[state_idx, :, :, :] .= ds[nc_key][:, :, :, end]
end
prob = remake(prob, u0 = vec(init_cond))
sol = solve(prob, SSPRK22(); dt=300f0)
Testing
I have tested this fix locally:
- Full test suite passes (263 pass, 0 new failures)
unknowns(prob.f.sys) matches fingerprint-detected ordering from prob.u0
- Restart workflow with
remake(prob, u0=...) with IC built from prob.f.sys ordering preserves correct species-to-index mapping
I have a branch with the fix ready at xk-y/EarthSciMLBase.jl:expose-sys-in-strang-solver and can open a PR if this approach looks good.
This issue was written with the assistance of Claude Code.
Problem
When using
SolverStrangThreads(orSolverStrangSerial),prob.f.sysisNothingafter building theODEProblem. This makes it impossible to programmatically determine which state index corresponds to which variable.This is critical for two use cases:
Root cause
In
solver_strategy_strang.jl, theODEProblemis constructed by passingnonstiff_op(a raw operator) directly:Since
nonstiff_opis not anODEFunction,ODEProblemwraps it in a newODEFunctionwithsys=nothing. Thesys_mtksystem is already available in scope but not attached.Why
unknowns(convert(System, model))is not a workaroundThe pre-simplification ordering from
unknowns(convert(System, model))does not match the actual state ordering inprob.u0. For example, with a 7-species surrogate chemistry model:Using
unknowns(convert(System, model))to assign initial conditions results in species getting each other's data (e.g., O3 gets NO2's IC and vice versa).Proposed fix (1 line)
Wrap
nonstiff_opinODEFunctionwithsys=sys_mtkbefore passing toODEProblem:This has zero runtime cost --
ODEProblemalready wraps raw operators inODEFunctioninternally. The fix simply makes it explicit and attaches the symbolic system metadata.Usage example: restart from NetCDF checkpoint
With the fix, users can do:
Testing
I have tested this fix locally:
unknowns(prob.f.sys)matches fingerprint-detected ordering fromprob.u0remake(prob, u0=...)with IC built fromprob.f.sysordering preserves correct species-to-index mappingI have a branch with the fix ready at
xk-y/EarthSciMLBase.jl:expose-sys-in-strang-solverand can open a PR if this approach looks good.This issue was written with the assistance of Claude Code.