DifferentialEquations.jl and systems with heterogenous units

Hi everybody,

I am trying to solve the following system of ODE with DifferentialEquations.jl

function newton!(t, y, dy, μ)
    r = norm(y[1:3])
    dy[1:3] = y[4:6]
    dy[4:6] = -μ * y[1:3] / r^3
end

Without units this works perfectly:

function test()
    r0 = [1131.340, -2282.343, 6672.423]
    v0 = [-5.64305, 4.30333, 2.42879]
    Δt = 86400.0*365
    mu = 398600.4418
    rv0 = [r0; v0]
    prob = ODEProblem((t, y, dy) -> newton!(t, y, dy, mu, rv0, (0.0, Δt))
    solve(prob, Vern8())[end]
end

The unitful version on the other hand does not:

function test2()
    r0 = [1131.340, -2282.343, 6672.423]u"km"
    v0 = [-5.64305, 4.30333, 2.42879]u"km/s"
    Δt = 86400.0*365u"s"
    mu = 398600.4418u"km^3/s^2"
    rv0 = [r0; v0]

    prob = ODEProblem((t, y, dy) -> newton!(t, y, dy, mu), rv0, (0.0u"s", Δt))
    solve(prob, Vern8())[end]
end
TypeError: setindex!: in typeassert, expected Quantity{Float64, Dimensions:{}, Units:{}}, got Float64

The state vector rv0 with heterogenous units seems to be the problem.

I realise that I could avoid this problem by defining the system of ODE as separate functions in the non-mutating form but this is not an option because the above will be just a small part of the RHS in the complete system.

Is there a way I can have my cake and also eat it?

Yes, but it’s not too simple. The ODE functions accept AbstractArrays, so we can build one for this purpose. As an example, DEDataArrays is one such supported array:

https://github.com/JuliaDiffEq/DiffEqBase.jl/blob/master/src/data_array.jl

So let’s build one for your purpose. Let’s call it a HeterogenousArray. Since you have two different components, let’s make it have two parts:

type HeterogeneousArray{T,T2} <: AbstractVector
  r::T
  v::T2
end

Now we need to make it actually act like an array. To do this, we implement the array interface. See

http://docs.julialang.org/en/stable/manual/interfaces/#abstract-arrays

It should just need:

size(A::HeterogeneousArray)   = (length(A.x) + length(A.y),)
getindex( A::HeterogeniousArray,    i::Int) = (i <= length(A.x) ? A.x[i] : A.y[i-length(A.x)]))
setindex!(A::HeterogeniousArray, x, i::Int) = (i <= length(A.x) ? (A.x[i] = x) : A.y[i-length(A.x)] = x) )

Now you should be able to create this array:

    r0 = [1131.340, -2282.343, 6672.423]u"km"
    v0 = [-5.64305, 4.30333, 2.42879]u"km/s"
    rv0 = HeterogeniousArray(r0,v0)

and then it should be that r0 == u0[1:3] and v0 == u0[4:6]. This means you can also simplify your ODE;

function newton!(t, y, dy, μ)
    r = norm(y[1:3])
    dy.r0 .= dy.v0
    dy.v0 .= -μ .* y[1:3] / r^3
end

Then you should be able to solve the ODE with this array:

prob = ODEProblem((t, y, dy) -> newton!(t, y, dy, mu), rv0, (0.0u"s", Δt))

That should work, but I didn’t debug this code so there may be a bug in this. But that should at least get you very close.

Caveats

I should explain the caveats here. Unitful.jl treats units as type information. This means that the eltype of your HeterogeneousArray is not concrete, and so accessing the values is not type stable. The way your newton! function is written here will be type stable, but using the A[i] iteration will not be. This will slow the solver down a bit.

Essentially, the solver does:

for i in eachindex(A)
  tmp[i] = tmp[i] + some_constant*dy[i]
end

Maybe there is a way to use broadcast:

tmp .= tmp .+ some_constant*dy

and have it iterate fast on such heterogenous arrays? Maybe @stevengj would know. That would be a good trick to find out!

If we can find out how to do some kind of “fast type-stable broadcast” here, and switch the internal for loops to broadcast (I’m waiting on stack-allocated views, as until then there is a performance issue with small broadcasts), then this approach should be fine performance-wise.

1 Like

Excellent! Thanks, Chris!

I was planning to do something along this line anyway since the r and v are already part of a State type. I will try this out and post here again if run into trouble.

I you don’t care too much about performance, and easier way to do this would be to just not have a strictly typed array.

    rv0 = Vector{Number}()
    r0 = [1131.340, -2282.343, 6672.423]u"km"
    v0 = [-5.64305, 4.30333, 2.42879]u"km/s"
    append!(rv0,r0); append!(rv0,v0)

But I noticed that

     Δt = 86400.0*365u"s"
    mu = 398600.4418u"km^3/s^2"
    prob = ODEProblem((t, y, dy) -> newton!(t, y, dy, mu), rv0, (0.0u"s", Δt))
    solve(prob, Vern8())[end]

this errors.

ERROR: DimensionError: 1//1000000 and 1.13134 km are not dimensionally compatible.
 in .+(::Rational{Int64}, ::Unitful.Quantity{Float64,Unitful.Dimensions{(Unitful.Dimension{:Length}(1//1),)},Unitful.FreeUnits{(Unitful.Unit{:Meter,Unitful.Dimensions{(Unitful.Dimension{:Length}(1//1),)}}(3,1//1),),Unitful.Dimensions{(Unitful.Dimension{:Length}(1//1),)}}}) at .\operators.jl:152
 in collect(::Base.Generator{Array{Unitful.Quantity{Float64,D,U},1},Base.##143#145{Rational{Int64}}}) at .\array.jl:307
 in .+(::Rational{Int64}, ::Array{Unitful.Quantity{Float64,D,U},1}) at .\arraymath.jl:90
 in +(::Rational{Int64}, ::Array{Unitful.Quantity{Float64,D,U},1}) at .\arraymath.jl:130
 in ode_determine_initdt(::Array{Number,1}, ::Unitful.Quantity{Float64,Unitful.Dimensions{(Unitful.Dimension{:Time}(1//1),)},Unitful.FreeUnits{(Unitful.Unit{:Second,Unitful.Dimensions{(Unitful.Dimension{:Time}(1//1),)}}(0,1//1),),Unitful.Dimensions{(Unitful.Dimension{:Time}(1//1),)}}}, ::Float64, ::Unitful.Quantity{Float64,Unitful.Dimensions{(Unitful.Dimension{:Time}(1//1),)},Unitful.FreeUnits{(Unitful.Unit{:Second,Unitful.Dimensions{(Unitful.Dimension{:Time}(1//1),)}}(0,1//1),),Unitful.Dimensions{(Unitful.Dimension{:Time}(1//1),)}}}, ::Rational{Int64}, ::Float64, ::OrdinaryDiffEq.#ODE_DEFAULT_NORM, ::DiffEqBase.ODEProblem{Array{Number,1},Quantity{Float64, Dimensions:{�}, Units:{s}},true,##3#4}, ::Int64) at C:\Users\Chris\.julia\v0.5\OrdinaryDiffEq\src\initdt.jl:5

This makes me think that the previous solution won’t work either. This isn’t going to work because it doesn’t know what to type some of the internal constants, and it usually types them using eltype(u0). In this case, eltype(u0) is not concrete, so it dies. To handle this case, the solvers will need to generalize how it does the typing.

Would you mind opening an issue? It’ll take some work to make heterogenous arrays like this work. And I think that with a proper broadcast (making them an iterator in the right), they can even work well. But it’ll take some work.

Done: https://github.com/JuliaDiffEq/DifferentialEquations.jl/issues/147

The example in the issue no longer works, I was able to adapt it to work as follows

using Unitful, RecursiveArrayTools, DiffEqBase, OrdinaryDiffEq, LinearAlgebra

function newton(u, p, t)
    mu = 398600.4418u"km^3/s^2"
    r = norm(u.x[1])
    ArrayPartition(u.x[2], -mu .* u.x[1] / r^3)
end

r0 = [1131.340, -2282.343, 6672.423]u"km"
v0 = [-5.64305, 4.30333, 2.42879]u"km/s"
Δt = 86400.0*365u"s"
rv0 = ArrayPartition(r0,v0)

prob = ODEProblem(newton, rv0, (0.0u"s", Δt))
sol = solve(prob, Vern8(),dt=100u"s",adaptive=false)

It’s probably inefficient, but I wasn’t able to get it to use the function signature that modifies du.

This also works now with ComponentArrays:

using ComponentArrays
using OrdinaryDiffEq
using LinearAlgebra
using Unitful

function newton(du, u, p, t)
    mu = 398600.4418u"km^3/s^2"
    r = norm(u.r)
    du.r = u.v
    du.v = -mu .* u.r / r^3
end

r0 = [1131.340, -2282.343, 6672.423]u"km"
v0 = [-5.64305, 4.30333, 2.42879]u"km/s"
Δt = 86400.0*365u"s"
rv0 = ComponentArray(r=r0, v=v0)

prob = ODEProblem(newton, rv0, (0.0u"s", Δt))
sol = solve(prob, Vern8(), dt=100u"s", adaptive=false)
3 Likes

This is the recommended solution.

2 Likes

I think we have an open issue on this. Can you post in there and we’ll close it?

1 Like

Is this the one?

yes

It should be pointed out that this relies on a bit of type twiddling underneath the hood. The ComponentArray type in this case will not have a concrete type for any of the elements.

The issue that I came across that led me here was that the algorithms in OrdinaryDiffEq rely on zero(rate_prototype). Surprisingly, zero(rv0) returns a ComponentVector with no units, because it resolves to similar.(rv0) .= 0. (To be clear, it doesn’t return elements of type Float64, but rather elements with a Unitful.Quantity type with NoUnits)

So this should work, but will likely suffer from a speed penalty. Just pointing this out in case others get stuck in the loop I was in.

Huh, that shouldn’t be the case. I’ll fix that.

Alright, it’s fixed now.

julia> rv0.r
3-element view(::Array{Quantity{Float64,D,U} where U where D,1}, 1:3) with eltype Quantity{Float64,D,U} where U where D:
   1131.34 km
 -2282.343 km
  6672.423 km

julia> rv0.v
3-element view(::Array{Quantity{Float64,D,U} where U where D,1}, 4:6) with eltype Quantity{Float64,D,U} where U where D:
 -5.64305 km s^-1
  4.30333 km s^-1
  2.42879 km s^-1

Now, there is still going to be some speed penalty due to the fact that the ComponentArray is wrapping a heterogeneous array. @ChrisRackauckas, this might be a use case where ArrayPartitions are the better choice.

ap_rv0 = ArrayPartition(r0, v0)

@btime $rv0 + $rv0       # 6.360 μs (46 allocations: 1.78 KiB)
@btime $ap_rv0 + $ap_rv0 # 67.111 ns (4 allocations: 272 bytes)

But weirdly enough, ComponentArrays outperform their underlying array…

julia> a_rv0 = getdata(rv0)
6-element Array{Quantity{Float64,D,U} where U where D,1}:
       1131.34 km
     -2282.343 km
      6672.423 km
 -5.64305 km s^-1
  4.30333 km s^-1
  2.42879 km s^-1

julia> @btime $a_rv0 + $a_rv0;
  14.499 μs (71 allocations: 3.81 KiB)
1 Like

Probably the broadcast overload + constant prop is helping

Nice!

1 Like

I think the underlying issue here is how to play nice with DifferentialEquations.jl underlying optimisations in preparing cache arrays, etc…

I think this is a discussion for a different thread, but I still want to mention what I’ve tried for this specific situation. The possibilities I’ve fiound are: (I’ll use Q{D} as a shorthand for Unitful.Quantity{Float64,D,U} which is not concrete)

  1. Plain Vector with non-concrete type Q{D} to represent quantities with different units. This breaks because zero(Q{D}) is (rightly) not defined.

  2. ComponentArray with non-concrete parameter Q{D}, but each component itself is of a definite type. Works fine and produces correct results, but is slow because the derivative function receives an object with non-concrete type.

  3. Hand-rolling my own struct:

struct Object{Tr,Tv}
r::Tr 
v::Tv 
end

I don’t think this works, because it is really difficult to make the internal DifferentialEquations.jl variables the correct type for units (e.g. handling the types returned by recursive_unitless_bottom_eltype, etc) in general. I think this approach would be fine if units were not an issue. (But then I would recommend ComponentArrays instead)

  1. Mimicking a vector with:
struct Object{Tr,Tv} <: AbstractVector
r::Tr 
v::Tv 
end

This doesn’t work as it is invalid inheritance. The AbstractVector needs an element type and it is impossible to put a concrete type there.

My conclusion is that currently DifferentialEquations.jl requires a single element type, but it can happily work with a non-concrete type, it will just be slow.

Perhaps the frustrating bit here is that it isn’t possible to override the internal DifferentialEquations.jl choices for types of derivatives. Would there be anything wrong if I were to attempt a pull request that extracts out those lines of code to functions like rate_type(u,t)? This might also reduce the size of functions like __init.

that would be helpful.

More correctly, Julia in general requires strict typing on containers to be fast. I don’t think that holding units in the type domain should be the final answer, and instead it should be isbits but (value,symbol) pairs IMO. But that’s a bigger story.

1 Like