Getting variable rate jumps to work

[kaandocal]

Is your feature request related to a problem? Please describe.

It would be nice to have VariableRateJumps (tied to an ODE system) work with ModelingToolkit out of the box. This would make it much easier to code up hybrid systems, and improve event handling capabilities!

Describe the solution you’d like

A simple and clear way to add VRJs to systems, e.g.

@variables x(t) y(y) 
@parameters p

...

# jump has propensity x + y^2 * u
jump = MT.VariableRateJump(x + y^2 * u, [ u ~ Pre(u) + 1 ])

...

# add jump to ODEProblem or JumpProblem

The interface for VariableRateJumps could be similar to that of continuous callbacks, but instead of a condition like x ~ 0 we supply a symbolic expression that computes the rate of a jump.

Describe alternatives you’ve considered

I currently have a dummy variable that integrates the jump intensity and use a callback based on that. This is a bit messy, and not very efficient since MTK by default adds interpolation points to continuous callbacks, which is unnecessary for VRJs with nonnegative rate.

Additional context

If I create a JumpProblem from an ODE problem and a VRJ, I obtain an error because ExtendedJumpArrays do not mesh with MTKs way of handling variables:

ERROR: ArgumentError: Expected `newu0` to be of same length as unknowns (6). Got ExtendedJumpArray{Float64, 1, Vector{Float64}, Vector{Float64}} of length 7
Stacktrace:
  [1] late_binding_update_u0_p(prob::ODEProblem{…}, sys::ODESystem, u0::ExtendedJumpArray{…}, p::Missing, t0::Float64, newu0::ExtendedJumpArray{…}, newp::MTKParameters{…})
    @ ModelingToolkit ~/.julia/packages/ModelingToolkit/aiAqm/src/systems/nonlinear/initializesystem.jl:680
  [2] late_binding_update_u0_p
    @ ~/.julia/packages/SciMLBase/MEOvx/src/remake.jl:1333 [inlined]
  [3] #updated_u0_p#906
    @ ~/.julia/packages/SciMLBase/MEOvx/src/remake.jl:1310 [inlined]
  [4] updated_u0_p
    @ ~/.julia/packages/SciMLBase/MEOvx/src/remake.jl:1286 [inlined]
  [5] #remake#868
    @ ~/.julia/packages/SciMLBase/MEOvx/src/remake.jl:228 [inlined]
  [6] remake
    @ ~/.julia/packages/SciMLBase/MEOvx/src/remake.jl:214 [inlined]
  [7] extend_problem(prob::ODEProblem{…}, jumps::Tuple{…}; rng::TaskLocalRNG)
    @ JumpProcesses ~/.julia/packages/JumpProcesses/orveP/src/variable_rate.jl:103
  [8] extend_problem
    @ ~/.julia/packages/JumpProcesses/orveP/src/variable_rate.jl:80 [inlined]
  [9] configure_jump_problem(prob::ODEProblem{…}, vr_aggregator::VR_FRM, jumps::JumpSet{…}, cvrjs::Tuple{…}; rng::TaskLocalRNG)
    @ JumpProcesses ~/.julia/packages/JumpProcesses/orveP/src/variable_rate.jl:63
 [10] configure_jump_problem
    @ ~/.julia/packages/JumpProcesses/orveP/src/variable_rate.jl:61 [inlined]
 [11] #JumpProblem#156
    @ ~/.julia/packages/JumpProcesses/orveP/src/problem.jl:287 [inlined]
 [12] JumpProblem
    @ ~/.julia/packages/JumpProcesses/orveP/src/problem.jl:226 [inlined]
 [13] #JumpProblem#149
    @ ~/.julia/packages/JumpProcesses/orveP/src/problem.jl:197 [inlined]
 [14] JumpProblem(prob::ODEProblem{…}, aggregator::Direct, jumps::VariableRateJump{…})
    @ JumpProcesses ~/.julia/packages/JumpProcesses/orveP/src/problem.jl:195