Skip to content

prob.f.sys is Nothing in Strang solver - cannot query variable ordering #142

Description

@xk-y

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:

  1. Restart from saved output: mapping each species' saved concentration to the correct state index
  2. 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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions