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).

[AayushSabharwal]

Yeah, so those should be turned into states instead of being anonymous variables and we should get better discrete/piecewise continuous support and use that here

[isaacsas]

How can JumpProcesses turn them into states? They are only added at the level of the JumpProcesses.JumpProblem internals, and not in MTK.

edit: I mean turn them into states in a way that works with symbolic indexing?

[AayushSabharwal]

The JumpSystem contains enough information to compute how many of them are required, right? So it should just be able to add them as artificial states in, say, a custom structural_simplify in the right order. That way symbolic indexing will work for propensities/intensities.

[isaacsas]

Maybe we can have a kwarg that is passed to a JumpProblem that MTK generates, u0_includes_propensities and then add them as symbolic states in JumpProblem(::JumpSystem) as you suggest. I'd have to look throughout JumpProcesses to see if that could mess anything else up / be handled in a non-breaking way.

For now can we do what Chris suggested above to get plot(sol) working? It seems solution indexing is otherwise ok I think. Given these are pseudo-states it would be nice to have a general way to indicate they shouldn't be plotted by default anyways (even for non-symbolic JumpProblems, where they do now get plotted in the default recipe last I checked).

[AayushSabharwal]

Maybe we can have a kwarg that is passed to a JumpProblem that MTK generates, u0_includes_propensities and then add them as symbolic states in JumpProblem(::JumpSystem) as you suggest. I'd have to look throughout JumpProcesses to see if that could mess anything else up / be handled in a non-breaking way.

Why would we need to do that? Just because they're unknowns in the system doesn't mean they have to be part of u0.

For now can we do what Chris suggested above to get plot(sol) working?

Already merged to SciMLBase#master

[ChrisRackauckas]

Just because they're unknowns in the system doesn't mean they have to be part of u0.

Yes, they don't need to be. A better implementation would use https://docs.sciml.ai/DiffEqCallbacks/stable/integrating/. A slightly more complicated discretecallback could be used to do the integral approximation with a direct rootfinding approach on it. But doing the full thing here would be a good chunk of work.

[isaacsas]

OK, that is great to hear. I have no idea how one can get MTK to generate a system/Problem where the final generated u0 and the number of unknowns in the finalized system are not required to be the same length, but if it can be done that sounds good to me. If you have an example you can point me to I'll take a look.

[isaacsas]

Already merged to SciMLBase#master

Awesome!

[ChrisRackauckas]

#3193 is the issue for it.

If we modeled the variable rate jumps symbolically and MTK did this optimization, we'd actually be quite a bit faster in many cases. But we'd have to move the rootfinding to the callback as well. But, because the interpolation of the ODE solution is a polynomial, that rootfinding would be a polynomial rootfind on the integral of the interpolating polynomial, and so we could use HomotopyContinuation to guarantee that we find all roots w.r.t. t. That would make it much more robust and quite a bit faster.

It's a good chunk of work but that is probably what we really need to make the hybrid modeling shine.

The nice part of doing it explicitly through MTK though is that this optimization can apply to many systems. A lot of systems people write down can simplify out part of it, so you'd get some good speedups in lots of places.

ExtendedJumpArray only had its weird behavior because SII didn't exist. Without SII and symbolic indexing, changing someone's u0 behind the scenes as part of jump "magic" can screw up the user's ability to understand the array, so we faked it as "an array that didn't change". In the context of MTK, we don't really need to do that though because there's already interfaces for handling the fact that the state can change and the user needs to respect that.

So with all of the SII things, and then also the inline SCC nonlinear solving stuff (which is an even bigger change to most Catalyst integration, especially with homotopy integration), we'll really need to reconsider the variable rate jump implementation since I would not be surprised to see it get 100x faster through connecting with the other optimizations going on.

[isaacsas]

With the SciMLBase fix and SciML/JumpProcesses.jl#461 this seems to work now. Thanks @AayushSabharwal for the quick fix.