using HybridArrays.jl

[SteffenPL]

I'm trying to simulate interacting particle systems. The state vector has dimension '2N' or '3N'.
It seems beneficial to use an Array of StaticArrays for the 2D/3D points.

HybridArrays.jl would be super nice here, since it avoids a Vector of SVectors
and is still very fast. (timings are in the code below; notice that access is exactly as for a 2xN Array!).

Issues:

• HybridArrays does not work with all methods (for example not with Rosenbrock23(), Rosenbrock23(autodiff=false)).
• Missing plot recipe (should be easy to implement)

I can try to work on the integration if wanted. Of course, hints are very welcome :D

I think HybridArrays could be an alternative to an Array of StaticArrays and without problems like #47 and SciML/LabelledArrays.jl#68.

using HybridArrays, StaticArrays
using DifferentialEquations, BenchmarkTools, Random

N = 50
Dim = 2

# just some random dynamics to show speed benefit of static arrays
function f_ode(du, u, p, t)
    N = size(u,2)
    du[:] .= 0.
    for i = 1:N
        for j = 1:i-1
            v = u[:,i] - u[:,j]
            dist_squared = sum( x -> x^2, v)
            du[:,i] -= v * dist_squared
            du[:,j] += v * dist_squared
        end
    end
end

# same as f_ode but for vector of svectors
function f_ode_vec_svec(du, u, p, t)
    N = size(u,1)
    du[:] .-= du[:]  #(sorry, I don't know how to zero a vector of svectors fast.)
    for i = 1:N
        for j = 1:i-1
            v = u[i] - u[j]
            dist_squared = sum( x -> x^2, v)
            du[i] -= v * dist_squared
            du[j] += v * dist_squared
        end
    end
end


Random.seed!(1)
z0 = rand(Dim,N)
z0_hybrid = HybridArray{Tuple{Dim,StaticArrays.Dynamic()}}(z0)

Random.seed!(1)
z0_vec_svec = rand(SVector{Dim,Float64},N)


prob          = ODEProblem(f_ode,          z0,          (0., 10.))
prob_hybrid   = ODEProblem(f_ode,          z0_hybrid,   (0., 10.))
prob_vec_svec = ODEProblem(f_ode_vec_svec, z0_vec_svec, (0., 10.))

sol          = @btime solve(prob,          Euler(), dt=0.01) seconds=2
# 682.401 ms (33153208 allocations: 1.65 GiB)
sol_hybrid   = @btime solve(prob_hybrid,   Euler(), dt=0.01) seconds=2
# 7.337 ms (14070 allocations: 3.04 MiB)
sol_vec_svec = @btime solve(prob_vec_svec, Euler(), dt=0.01) seconds=2
# 5.019 ms (20074 allocations: 4.78 MiB)


sol2          = solve(prob,          Rosenbrock23())
# works
sol2_hybrid   = solve(prob_hybrid,   Rosenbrock23())
# MethodError: copyto!(::DiffEqBase.DiffEqBC{Array{Float64,2}}, ::Base.Broadcast.Broadcasted{HybridArrays.HybridArrayStyle{2},Tuple{Base.OneTo{Int64},Base.OneTo{Int64}},typeof(muladd),Tuple{Base.Broadcast.Broadcasted{HybridArrays.HybridArrayStyle{2},Nothing,typeof(DiffEqBase.ODE_DEFAULT_NORM),Tuple{HybridArray{Tuple{2,StaticArrays.Dynamic()},Float64,2,2,Array{Float64,2}},Float64}},Float64,Float64}})
# is ambiguous.
sol2_vec_svec = solve(prob_vec_svec, Rosenbrock23())
# ArgumentError: Cannot create a dual over scalar type SArray{Tuple{2},Float64,1,2}.
#  If the type behaves as a scalar, define FowardDiff.can_dual.

using Plots
plot( sol[1,1,:], sol[2,1,:])
plot( sol_hybrid[1,1,:], sol_hybrid[2,1,:])
plot( getindex.(sol_vec_svec[1,:],1) , getindex.(sol_vec_svec[1,:],2))
# plots look all the same

plot( sol )
# works
plot( sol_hybrid )
# MethodError: no method matching u_n(...)  
# Note: Here a ::CartesianIndex{2}  does not match the expected ::Int64
plot( sol_vec_svec )
# MethodError: no method matching one(::Type{SArray{Tuple{2},Float64,1,2}})

I can provide full error messages if needed.

[mateuszbaran]

Hi!

Thanks for using HybridArrays 🙂 . I can see at least two problems here:

1. copyto!(dest::DiffEqBC, bc::$B) from diffeqfastbc.jl needs to be specialized for HybridArrayStyle somewhere. I guess https://github.com/SciML/DiffEqBase.jl/blob/master/src/diffeqfastbc.jl would be the right place, just executed on demand using Requires.jl. I think the same copyto! as here: https://github.com/SciML/DiffEqBase.jl/blob/d6769e8aeac3eb29e54d2e557ceacc4f4fa27891/src/diffeqfastbc.jl#L46 , just for HybridArrayStyle, would be enough.
2. HybridArrays needs to overload some additional functions. I'm not sure which ones, definitely zero to make the right cache, but if someone more experienced with DiffEq ecosystem could make a list (or link to one) it would be helpful. At the moment HybridArrays doesn't overload that much, mostly broadcasting.

I can help fixing this on the HybridArrays side but I'd like to hear from DiffEq people before I start working on it.

And, just in case, HybridArrays will most likely drop SSubArray in the near-ish future since allocation improvements in Julia 1.5 should make this trick obsolete.

[mateuszbaran]

That plot issue is really intriguing. On one hand, HybridArrays doesn't expose linear indexing where it could. On the other hand, plotting recipe system doesn't work with arrays that are not linearly indexable. On yet another hand, I've noticed that for some reason eachindex on static arrays returns Base.OneTo instead of SOneTo.

[mateuszbaran]

And, by the way, I could implement ArrayInteface.jl methods for HybridArray, seems easy enough. I'm just not sure where should I put the code.

[ChrisRackauckas]

It's best to put the code in the array library since less Requires will mean less compile times. Right now Requires is just used in order to jump start the adoption of an extended trait ecosystem: if the standard arrays don't have the trait, then there's no point in using it.

[mateuszbaran]

OK, I've implemented ismutable, can_setindex, parent_type and restructure in HybridArrays. That should be enough I guess.

[mateuszbaran]

FYI, with the changes from JuliaArrays/HybridArrays.jl#19 HybridArray is almost as fast as Vector of SVectors in this example (4.968 ms vs 4.492 ms) and plotting also works.

[SteffenPL]

Wow, thanks for the implementation! Yes, I can confirm that the speed is very good and plotting works.

Implicit methods still fail, since ldiv! is not defined (I tested ImplicitEuler, Rosenbrock23).
But this also doesn't work with SVectors and MVectors, so it might be more complicated. (I'm not sure if DiffEq has special treatment for those.)

PS: For my application, implicit methods are not important. Your changes already solve the other difficulties I had 👍

using LinearAlgebra

A = qr([1 2.2 4; 3.1 0.2 3; 4 1 2]);
X = [1; 2.5; 3];
Yd = HybridArray{Tuple{StaticArrays.Dynamic()}}(X)
Ys = HybridArray{Tuple{3}}(X)
Ysa = SVector{3,Float64}(X)
Ym = MVector{3,Float64}(X)

ldiv!(A,X) # works

ldiv!(A,Yd)
# ERROR: MethodError: no method matching ldiv!(::LinearAlgebra.QRCompactWY{Float64,Array{Float64,2}}, ::HybridArray{Tuple{StaticArrays.Dynamic()},Float64,1,1,Array{Float64,1}})
# Closest candidates are:
#   ldiv!(::IterativeSolvers.Identity, ::Any) at C:\Users\Steffen Plunder\.julia\packages\IterativeSolvers\3g7hG\src\common.jl:31
#   ldiv!(::Number, ::AbstractArray) at D:\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v1.4\LinearAlgebra\src\generic.jl:252
#   ldiv!(::Adjoint{T,#s48} where #s48<:(BandedMatrices.BandedLU{T,#s41} where #s41<:BandedMatrices.BandedMatrix), ::Union{AbstractArray{T,1}, AbstractArray{T,2}}) where T<:Real at C:\Users\Steffen Plunder\.julia\packages\BandedMatrices\vcLoH\src\banded\linalg.jl:54

ldiv!(A,Ys)
# ERROR: MethodError: no method matching ldiv!(::LinearAlgebra.QRCompactWY{Float64,Array{Float64,2}}, ::HybridArray{Tuple{3},Float64,1,1,Array{Float64,1}})
# ...

ldiv!(A,Ysa)
# ERROR: MethodError: no method matching ldiv!(::LinearAlgebra.QRCompactWY{Float64,Array{Float64,2}}, ::SArray{Tuple{3},Float64,1,3})
# ...

ldiv!(A,Ym)
# ERROR: MethodError: no method matching ldiv!(::LinearAlgebra.QRCompactWY{Float64,Array{Float64,2}}, ::SArray{Tuple{3},Float64,1,3})
# ...

[mateuszbaran]

I've implemented vec for HybridArray and now implicit solvers also work. They are also much faster on HybridArray:

julia> sol   = @btime solve(prob,   ImplicitEuler(), dt=0.01); seconds=2
  760.786 ms (5048795 allocations: 534.85 MiB)
2

julia> sol_hybrid   = @btime solve(prob_hybrid,   ImplicitEuler(), dt=0.01); seconds=2
  27.109 ms (10635 allocations: 976.16 KiB)
2

julia> sol_hybrid   = @btime solve(prob_hybrid,   Rosenbrock23(), dt=0.01); seconds=2
  38.912 ms (5940 allocations: 622.34 KiB)
2

julia> sol   = @btime solve(prob,   Rosenbrock23(), dt=0.01); seconds=2
  884.613 ms (3313352 allocations: 723.02 MiB)
2

[ChrisRackauckas]

Sounds great. I think this is done?

[ChrisRackauckas]

How are hybrid arrays so much faster there? That's awesome.

[SteffenPL]

I guess for NBody or particle simulations this is very attractive ( https://github.com/SciML/DiffEqPhysics.jl )
maybe we could document somewhere that HybridArrays are working?

(A tutorial/example somewhere?)

[SteffenPL]

Is there a way to test that an array type is fully compatible with DiffEq?

[mateuszbaran]

Sounds great. I think this is done?

More or less. I guess we could also check if it's possible to wrap CuArray in HybridArray and get even better performance on GPU but I don't have an NVidia GPU.

How are hybrid arrays so much faster there? That's awesome.

HybridArrays uses the same tricks for indexing and broadcasting as StaticArrays, just adapted to arrays where only a subset of dimensions have statically known sizes. As long as indexing or broadcasting only touches parts with statically known dimensions, the generated code is more or less on par with StaticArrays.

The rest is just the magic of Julia 🙂 .

By the way, I've tried solving that problem using MArray. Compilation of ImplicitEuler solve used all my memory and I had to kill the process...