Some solvers don't work with ArrayPartition + ExtendedJumpArray

[kaandocal]

prob = ODEProblem((du, u, p, t) -> du .= 0, ArrayPartition([0.]), (0, 10.))
jump = VariableRateJump((u, p, t) -> 1, int -> (int.u += 1))
jprob = JumpProblem(prob, Direct(), jump)
solve(jprob, KenCarp4())          # can replace this by TRBDF2

yields

BoundsError: attempt to access 1-element ArrayPartition{Float64, Tuple{Vector{Float64}}} at index [2:2]

Stacktrace:
  [1] throw_boundserror(A::ArrayPartition{Float64, Tuple{Vector{Float64}}}, I::Tuple{UnitRange{Int64}})
    @ Base ./abstractarray.jl:691
  [2] checkbounds
    @ ./abstractarray.jl:656 [inlined]
  [3] _setindex!
    @ ./multidimensional.jl:893 [inlined]
  [4] setindex!
    @ ./abstractarray.jl:1315 [inlined]
  [5] _typed_vcat!(a::ArrayPartition{Float64, Tuple{Vector{Float64}}}, V::Tuple{ArrayPartition{Float64, Tuple{Vector{Float64}}}, Vector{Float64}})
    @ Base ./abstractarray.jl:1551
  [6] _typed_vcat
    @ ./abstractarray.jl:1543 [inlined]
  [7] typed_vcat
    @ ./abstractarray.jl:1619 [inlined]
  [8] vcat
    @ ./abstractarray.jl:1535 [inlined]
  [9] zeromatrix(A::ExtendedJumpArray{Float64, 1, ArrayPartition{Float64, Tuple{Vector{Float64}}}, Vector{Float64}})
    @ DiffEqJump ~/.julia/packages/DiffEqJump/x05Qi/src/extended_jump_array.jl:36

The issue seems to be that vec applied to an ArrayPartition returns a tuple, not a vector. Running the above with Vern7 gives

MethodError: no method matching (::DiffEqJump.var"#jump_f#123"{ODEProblem{ArrayPartition{Float64, Tuple{Vector{Float64}}}, Tuple{Float64, Float64}, true, SciMLBase.NullParameters, ODEFunction{true, var"#308#309", UniformScaling{Bool}, (...), typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, JumpSet{Tuple{VariableRateJump{var"#310#312", var"#311#313", Nothing, Float64, Int64}}, Tuple{}, Nothing, Nothing}})(::Vector{Float64}, ::ExtendedJumpArray{Float64, 1, ArrayPartition{Float64, Tuple{Vector{Float64}}}, Vector{Float64}}, ::SciMLBase.NullParameters, ::Float64)
Closest candidates are:
  (::DiffEqJump.var"#jump_f#123")(::ExtendedJumpArray, ::ExtendedJumpArray, ::Any, ::Any) at ~/.julia/packages/DiffEqJump/x05Qi/src/problem.jl:125

Stacktrace:
  [1] (::ODEFunction{true, DiffEqJump.var"#jump_f#123"{ODEProblem{ArrayPartition{Float64, Tuple{Vector{Float64}}}, Tuple{Float64, Float64}, true, SciMLBase.NullParameters, ODEFunction{true, var"#308#309", UniformScaling{Bool}, (...), typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, JumpSet{Tuple{VariableRateJump{var"#310#312", var"#311#313", Nothing, Float64, Int64}}, Tuple{}, Nothing, Nothing}}, UniformScaling{Bool}, (...), typeof(SciMLBase.DEFAULT_OBSERVED), Nothing})(::Vector{Float64}, ::Vararg{Any})
    @ SciMLBase ~/.julia/packages/SciMLBase/pr0Dt/src/scimlfunctions.jl:345
  [2] ode_determine_initdt(u0::ExtendedJumpArray{...}, t::Float64, tdir::Float64, dtmax::Float64, abstol::Float64, reltol::Float64, internalnorm::typeof(DiffEqBase.ODE_DEFAULT_NORM), prob::ODEProblem{...}, ...)
    @ OrdinaryDiffEq ~/.julia/packages/OrdinaryDiffEq/Op0Oq/src/initdt.jl:49

Here the problem is in initdt, providing dt = ... to the solve call bypasses the issue. Sorry for the long error messages; the problem seems to be that ArrayPartition does not follow some standard array interface conventions - any advice for convenient alternatives when solving problems with multiple components?

[ChrisRackauckas]

Does a ComponentArray work for this?

[kaandocal]

It does, indeed! Thanks for the hint.

[isaacsas]

@kaandocal I'm going to close this, but feel free to reopen if you have further issues/questions!

[kaandocal]

I'd argue that this is an example where two rather important Julia packages don't work together - I understand that this might not be solved now (I could have a look at the ArrayPartition interface), but there might be an argument in favour of keeping this open for now...

[isaacsas]

Sure, I've opened this back up. My impression had been that this is more of an ArrayPartition issue from your comments than anything specific in DiffEqJump, but if there are concrete dispatches that could be added here to fix this please do let us know (or make a PR as you said).

[kaandocal]

Thanks! It's about the way EJA and APs interact as stiff solvers work for pure ODEProblems - I can open an issue at RecursiveArrayTools.jl or we can shift this one there. In the second case (Vern7) I'm not sure what exactly happens, but DiffEqJump ends up calling a function with the wrong argument types, which may be down to either of the packages. I'd need to delve into the details to figure out what the culprit is...

[ChrisRackauckas]

Yeah in theory it can work. I'm not sure why vec would return a tuple. I gave a workaround to get you going since I'll be travelling a lot and so this one might snooze for a bit. But it most likely is just some simple AbstractArray interface break.

[kaandocal]

Thanks, much appreciated! This is not an urgent issue given that ComponentArrays do the job, but it might stump future users, which is why I would prefer keeping this open for now...

[paolo-mgi]

Please bear with me, my programming skills are very limited. It seems, however, to me that both
ArrayPartition and ComponentArrays have an issue with ExtendedJumpArray.
Namely the example code below ends in both cases giving the error message

type ExtendedJumpArray has no field v ( if I use ComponentArrays) or x (if I use ArrayPartition)

I would prefer the latter option to work because I am dealing with a system of jump equations where one real component
depends though a real quantity on the remaining complex components.

using DifferentialEquations, LinearAlgebra, Plots
#using RecursiveArrayTools.
psi0 = [1.0 + 0.0im, 1.0 + 0.0im]./sqrt(2)
#u0= ArrayPartition(psi0, [1.0])
tf =15.0
using ComponentArrays
u0= ComponentArray(v=psi0, m=1.0)

p = (;  
	process_operators = ([ 1.0 0.0 ; 0.0 1.0], [ 0.0 1.0  ; 1.0  0.0]), 
         temp_psi = similar(psi0)
	)

function psi!(du, u, p, t)
   Heff = @. - im * p.process_operators[1]
 #  mul!(du.x[1] ,Heff , u.x[1])
 # du.x[2] .= u.x[2] 
    mul!(du.v ,Heff , u.v)
    du.m = u.m 
end

function affect1!(integrator)
 #  vv = copyto!(integrator.p.temp_psi, integrator.u.x[1]) 
 # mart =  integrator.u.x[2] `
  vv = copyto!(integrator.p.temp_psi, integrator.u.v) 
  mart =  integrator.u.m # copy this to avoid aliasing later 
   jump_prob = real(dot(vv, integrator.p.process_operators[2], vv)) 
   normv = sqrt(jump_prob ) 
  #mul!(integrator.u.x[1], integrator.p.process_operators[2], vv, 1/normv,false) 
   #integrator.u.x[2] = mart +1.0 
  mul!(integrator.u.v, integrator.p.process_operators[2], vv, 1/normv,false)
   integrator.u.m = mart +1.0
end
var_jump1 = VariableRateJump(rate1!, affect1!)
num_traj=1  
sse_prob = ODEProblem(psi!, u0, (0.0, tf),p)
jump_prob = JumpProblem(sse_prob, Direct() , var_jump1)
ensemble_sse = EnsembleProblem(jump_prob)
jump_sol_sse = DifferentialEquations.solve(ensemble_sse, Tsit5(), EnsembleSerial(), trajectories=num_traj; save_everystep=true, dense=false,abstol=1e-8,reltol=1e-8,maxiters=Int(1e6));

Many thanks for your attention.

[isaacsas]

Your actual ComponentArray or ArrayPartition is stored as u.u and du.u would be the analogous output location. So I think you'd want something like du.u.x[1] etc. Does that help?

See also the ExtendedJumpArray doc string: https://docs.sciml.ai/JumpProcesses/stable/api/#JumpProcesses.ExtendedJumpArray

[paolo-mgi]

Many thanks for your comment!

If I replace affect! as

function affect1!(integrator)
#  vv = copyto!(integrator.p.temp_psi, integrator.u.u.x[1]) # copy this to avoid aliasing later
#  mart =  integrator.u.u.x[2] # copy this to avoid aliasing later
  vv = copyto!(integrator.p.temp_psi, integrator.u.u.v) # copy this to avoid aliasing later
  mart =  integrator.u.u.m # copy this to avoid aliasing later
  jump_prob = real(dot(vv, integrator.p.process_operators[2], vv))
  normv = sqrt(jump_prob )
#  mul!(integrator.u.u.x[1], integrator.p.process_operators[2], vv, 1/normv,false)
#  integrator.u.u.x[2] = mart +1.0
  mul!(integrator.u.u.v, integrator.p.process_operators[2], vv, 1/normv,false)
  integrator.u.u.m = mart +1.0
end

The code is executed but no jump occurs. So the/an issue remains. If I use instead ArrayPartition I get the error message.

MethodError: no method matching +(::Vector{Float64}, ::Float64)
For element-wise addition, use broadcasting with dot syntax: array .+ scalar

Many thanks!

[isaacsas]

Did you update your psi! function too?

function psi!(du, u, p, t)
   Heff = @. - im * p.process_operators[1]
   mul!(du.u.v, Heff, u.u.v)
   du.u.m = u.u.m 
end

[isaacsas]

Could you give a simple, short MWE here without all the commented code and intermediary operations? It would make testing / debugging easier.

[paolo-mgi]

Sorry I did not further update. Indeed after further debugging I realised that my code correctly works with

using ComponentArrays
u0= ComponentArray(v=psi0, m=1.0)

and with u.u as you suggested. The problem only remains with

using RecursiveArrayTools
u0= ArrayPartition(psi0, [1.0])

so the situation is exactly the one described at the beginning of this thread. If you are interested I can provide a MWE

[isaacsas]

Awesome, happy to hear ComponentArrays work! No need for a MWE with ArrayPartitions, The issue there is missing dispatches, and I think the general recommendation in SciML is to use ComponentArray over ArrayPartitions because the former is more feature complete.

[paolo-mgi]

Excellent. Many thanks for your help!