# How to solve an ODEProblem with u0 in the format of vector{vector} and be able to use it for Zygote differentiation

Hello, I am currently facing an ode problem. The format of u0 in this problem is vector{vector}, for example [[1,2], [1,2,3], [1]], the particularity of this u0 is that the length of each element is different, and I cannot easily convert it into a matrix as u0. I first directly use vector{vector} as u0, the code is as follows:

``````
function foo_1(u, p, t)
p1, p2 = p[1], p[2]
[u[1] .* p1, u[2] .* p2]
end

prob = ODEProblem(foo_1, [[2.0, 3.0], [2.1, 3.2, 4.2]], (0, 5), [1.01, 0.99])
sol = solve(prob, Tsit5(), saveat=1.0)
``````

The results of running this code suggest that I use RecursiveArrayTools.jl or ComponentArrays.jl as the format of u0, so I refactored the code as follows:

``````# based on RecursiveArrayTools
function foo_2(u, p, t)
p1, p2 = p[1], p[2]
VectorOfArray([u[1] .* p1, u[2] ./ p2])
end

prob = ODEProblem(foo_2, VectorOfArray([[2.0, 3.0], [2.1, 3.2, 4.2]]), (0, 5), [1.01, 0.99])
sol = solve(prob, Tsit5(), saveat=1.0)

# based on ComponentArrays
function foo_3(u, p, t)
p1, p2 = p[1], p[2]
ComponentVector(a=u.a .* p1, b=u.b ./ p2)
end

prob = ODEProblem(f2, ComponentVector(a=[2.0, 3.0], b=[2.1, 3.2, 4.2]), (0, 5), [1.01, 0.99])
sol = solve(prob, Tsit5(), saveat=1.0)
``````

Both methods can solve the ode problem (the results seem to be different), and then I try to use zygote to find the gradients calculated by these two methods:

``````
function g2(p)
function foo_2(u, p, t)
p1, p2 = p[1], p[2]
VectorOfArray([u[1] .* p1, u[2] ./ p2])
end

prob = ODEProblem(foo_2, VectorOfArray([[2.0, 3.0], [2.1, 3.2, 4.2]]), (0, 5), p)
sol = solve(prob, Tsit5(), saveat=1.0)
sum(sol)
end

``````

The way you use RecursiveArrayTools will cause an error:

``````ERROR: Can only convert non-ragged VectorOfArray to Array
Stacktrace:
[1] error(s::String)
@ Base .\error.jl:35
[2] convert(::Type{Array}, VA::VectorOfArray{Float64, 2, Vector{Vector{Float64}}})
@ RecursiveArrayTools D:\Julia\Julia-1.10.4\packages\packages\RecursiveArrayTools\xGKIm\src\vector_of_array.jl:760
[3] vec
@ D:\Julia\Julia-1.10.4\packages\packages\RecursiveArrayTools\xGKIm\src\vector_of_array.jl:756 [inlined]
[4] (::SciMLSensitivity.var"#330#339"{…})()
@ SciMLSensitivity D:\Julia\Julia-1.10.4\packages\packages\SciMLSensitivity\PstNN\src\concrete_solve.jl:953
[5] unthunk
@ D:\Julia\Julia-1.10.4\packages\packages\ChainRulesCore\I1EbV\src\tangent_types\thunks.jl:204 [inlined]
[6] wrap_chainrules_output
@ D:\Julia\Julia-1.10.4\packages\packages\Zygote\nsBv0\src\compiler\chainrules.jl:110 [inlined]
[7] map
@ .\tuple.jl:293 [inlined]
[8] map (repeats 3 times)
@ .\tuple.jl:294 [inlined]
[9] wrap_chainrules_output
@ D:\Julia\Julia-1.10.4\packages\packages\Zygote\nsBv0\src\compiler\chainrules.jl:111 [inlined]
[10] ZBack
@ D:\Julia\Julia-1.10.4\packages\packages\Zygote\nsBv0\src\compiler\chainrules.jl:211 [inlined]
[11] kw_zpullback
@ D:\Julia\Julia-1.10.4\packages\packages\Zygote\nsBv0\src\compiler\chainrules.jl:237 [inlined]
[12] #291
@ D:\Julia\Julia-1.10.4\packages\packages\Zygote\nsBv0\src\lib\lib.jl:206 [inlined]
[13] (::Zygote.var"#2169#back#293"{…})(Δ::ODESolution{…})
[14] #solve#51
@ D:\Julia\Julia-1.10.4\packages\packages\DiffEqBase\V6SCE\src\solve.jl:1003 [inlined]
[15] (::Zygote.Pullback{…})(Δ::ODESolution{…})
@ Zygote D:\Julia\Julia-1.10.4\packages\packages\Zygote\nsBv0\src\compiler\interface2.jl:0
[16] #291
@ D:\Julia\Julia-1.10.4\packages\packages\Zygote\nsBv0\src\lib\lib.jl:206 [inlined]
[17] (::Zygote.var"#2169#back#293"{…})(Δ::ODESolution{…})
[18] solve
@ D:\Julia\Julia-1.10.4\packages\packages\DiffEqBase\V6SCE\src\solve.jl:993 [inlined]
[19] (::Zygote.Pullback{…})(Δ::ODESolution{…})
@ Zygote D:\Julia\Julia-1.10.4\packages\packages\Zygote\nsBv0\src\compiler\interface2.jl:0
[20] g2
@ e:\JlCode\LumpedHydro\test\tmp_test.jl:42 [inlined]
[21] (::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Float64)
@ Zygote D:\Julia\Julia-1.10.4\packages\packages\Zygote\nsBv0\src\compiler\interface2.jl:0
[22] (::Zygote.var"#75#76"{Zygote.Pullback{Tuple{…}, Tuple{…}}})(Δ::Float64)
@ Zygote D:\Julia\Julia-1.10.4\packages\packages\Zygote\nsBv0\src\compiler\interface.jl:91
@ Zygote D:\Julia\Julia-1.10.4\packages\packages\Zygote\nsBv0\src\compiler\interface.jl:148
[24] top-level scope
@ e:\JlCode\LumpedHydro\test\tmp_test.jl:46
Some type information was truncated. Use `show(err)` to see complete types.
``````

Then using ComponentArrays seems to work:

``````function g3(p)
function foo_3(u, p, t)
p1, p2 = p[1], p[2]
ComponentVector(a=u.a .* p1, b=u.b ./ p2)
end

prob = ODEProblem(foo_3, ComponentVector(a=[2.0, 3.0], b=[2.1, 3.2, 4.2]), (0, 5), p)
sol = solve(prob, Tsit5(), saveat=1.0)
sum(sol)
end

The result is: `([5436.3132778556665, -10543.639885459668],)`