Plotting no longer works for MTK JumpSystems over ODEProblems

[isaacsas]

using ModelingToolkit, OrdinaryDiffEqTsit5, JumpProcesses, Plots
t = ModelingToolkit.t_nounits
@variables A(t) B(t) C(t)
@parameters k
vrj = VariableRateJump(k * (sin(t) + 1), [A ~ A + 1, C ~ C + 2])
js = complete(JumpSystem([vrj], t, [A, C], [k]; name = :js, observed = [B ~ C * A]))
oprob = ODEProblem(js, [A => 0, C => 0], (0.0, 10.0), [k => 1.0])
jprob = JumpProblem(js, oprob)
sol = solve(jprob, Tsit5())
plot(sol)

gives

julia> plot(sol)
ERROR: ArgumentError: Collection must contain exactly 1 element
Stacktrace:
  [1] _only
    @ ./iterators.jl:1561 [inlined]
  [2] only(s::Set{Any})
    @ Base ./set.jl:178
  [3] macro expansion
    @ ~/.julia/packages/SciMLBase/u1M8h/src/solutions/solution_interface.jl:242 [inlined]
  [4] apply_recipe(plotattributes::AbstractDict{Symbol, Any}, sol::SciMLBase.AbstractTimeseriesSolution)
    @ SciMLBase ~/.julia/packages/RecipesBase/BRe07/src/RecipesBase.jl:300
  [5] _process_userrecipes!(plt::Any, plotattributes::Any, args::Any)
    @ RecipesPipeline ~/.julia/packages/RecipesPipeline/BGM3l/src/user_recipe.jl:38
  [6] recipe_pipeline!(plt::Any, plotattributes::Any, args::Any)
    @ RecipesPipeline ~/.julia/packages/RecipesPipeline/BGM3l/src/RecipesPipeline.jl:72
  [7] _plot!(plt::Plots.Plot, plotattributes::Any, args::Any)
    @ Plots ~/.julia/packages/Plots/kLeqV/src/plot.jl:223
  [8] plot(args::Any; kw...)
    @ Plots ~/.julia/packages/Plots/kLeqV/src/plot.jl:102
  [9] plot(args::Any)
    @ Plots ~/.julia/packages/Plots/kLeqV/src/plot.jl:93
 [10] top-level scope
    @ REPL[19]:1

Environment:

(jl_D37Hzc) pkg> st
Status `/private/var/folders/8s/kss4zcyx10v_l22h64zdzs6m0000gn/T/jl_D37Hzc/Project.toml`
  [ccbc3e58] JumpProcesses v9.14.0
  [961ee093] ModelingToolkit v9.49.0
  [b1df2697] OrdinaryDiffEqTsit5 v1.1.0
  [0bca4576] SciMLBase v2.58.1

For a JumpProblem over a DiscreteProblem things still seem to work:

@variables A(t) B(t) 
@parameters k
crj = ConstantRateJump(k*A, [A ~ A - 1, B ~ B + 1])
js = complete(JumpSystem([crj], t, [A, B], [k]; name = :js))
dprob = DiscreteProblem(js, [A => 10, B => 0], (0.0, 10.0), [k => 1.0])
jprob = JumpProblem(js, dprob)
sol = solve(jprob)
plot(sol)

[ChrisRackauckas]

@AayushSabharwal are jumps captured as discrete variables now?

[AayushSabharwal]

I don't think so. There's nothing special we do for discrete variables in JumpSystem to my knowledge

[AayushSabharwal]

The problem here is that solutions of jump systems seem to do some black magic that SII isn't aware of:

infil> sol
retcode: Success
Interpolation: specialized 4th order "free" interpolation
t: 36-element Vector{Float64}:
  0.0
# ...
u: 36-element Vector{ExtendedJumpArray{Float64, 1, Vector{Float64}, Vector{Float64}}}:
 [0.0, 0.0]
# ...

Note how it appears sol.u[1] is a 2-length vector

julia> sol.u[1]
3-element ExtendedJumpArray{Float64, 1, Vector{Float64}, Vector{Float64}}:
  0.0
  0.0
 -0.2529329586712836

I'm not aware of what this is or what it's supposed to mean. As far as SII/the JumpSystem knows, the system has one unknown and the solution should have a timeseries of a single variable

[ChrisRackauckas]

Yeah they have pseudo variables that are used for the event handling in jump detection. They aren't real variables. I have thought for over a year now that we should probably just not hide them from the user and just use a ComponentArray by default or something. But SII validation needs to somehow allow this. Maybe we need to declare pseudonames for them in the MTK side?

[AayushSabharwal]

If the system specification can determine the number and names of pseudo variables, SII will be happy. Right now, plot(sol) lowers to plot(sol, idxs = 1:length(sol.u[1])) which is plot(sol, idxs = 1:3). SII says "I know 1 is a continuous variable" but throws up its hands at 2 and 3. Plotting code assumes that either a variable is continuous, or discrete. SII says neither. I can refactor the plotting code to consider "neither" as continuous, but not as discrete which is what they seem to be here.

The true solution would be for JumpSystem to declare all its pseudo variables, their types and whether they're discrete. Then, either the generated code or JumpProcesses needs to hook into SII to save discrete variables properly instead of artificially inserting them into the timeseries. I get the feeling, though, that jump processes considers these variables as piecewise continuous and not discrete in the true sense, which means we need [email protected] with its discrete indexing re-re-refactor.

[ChrisRackauckas]

I get the feeling, though, that jump processes considers these variables as piecewise continuous and not discrete in the true sense, which means we need [email protected] with its discrete indexing re-re-refactor.

Yes, that will be the case for the general solution.

But for a more immediate solution, maybe we just need an interface plottable_indices(u) = 1:length(u) that then gets overriden for ExtendedJumpArray?

[AayushSabharwal]

Do we not want to plot the first two?

[ChrisRackauckas]

You want to plot only the first two.

[AayushSabharwal]

Alright yeah then plottable_indices is the best immediate solution

[AayushSabharwal]

I'd need an MWE to tell you more. Just from the stacktrace, it could be either of the two reasons.

[isaacsas]

Note that propensities/intensities in JumpProcesses arising from VariableRateJumps, the "pseudo-variables" in the discussion above, can be continuous in time. So it wouldn't be appropriate to treat them as discrete (or even as piecewise-constant in general).