Indexing StaticArrays

This is such a basic question that I almost feel bad asking it, but it would be nice to get some clarification and examples of how others are using StaticArrays.jl. I’ve found this thread, but I think my question is slightly more general.

using StaticArrays
x = SVector{5}([1.1, 1.2, 1.3, 1.4, 1.5])

julia> x[[2,5]]
2-element Vector{Float64}:
1.2
1.5

julia> x[2:3]
2-element Vector{Float64}:
1.2
1.3

First, it seems very strange to me that the return type is not SVector. I’m wondering what the intended baseline use case is where you’d want the type to change when indexing. Or maybe there’s not a way to make it work with multiple dispatch with StaticArrays.

After a bit of digging, I’ve found that I can do either of the following:

julia> x[SVector{2}([2,5])]
2-element SVector{2, Float64} with indices SOneTo(2):
1.2
1.5

julia> x[StaticArrays.SUnitRange(2,3)]
2-element SVector{2, Float64} with indices SOneTo(2):
1.2
1.3

Given that SUnitRange isn’t even exported, I’m guessing that I’m looking in the wrong place. This also seems to involve writing a lot of custom code when you want to use StaticArrays.

Ultimately, I think my questions are:

What is the “right” way to index StaticArrays?

How should we write general code that is dispatched to return a specified subset of an AbstractArray with the correct return type?

Some of my heavily optimized code (non-allocating):

function residual!(res, yd, y::MVector{S, SimFloat}, s::KPS4, time) where S
    T = S-2 # T: three times the number of particles excluding the origin
    segments = div(T,6) - KITE_PARTICLES
    
    # Reset the force vector to zero.
    for i in 1:segments+KITE_PARTICLES+1
        s.forces[i] .= SVector(0.0, 0, 0)
    end
    # extract the data for the winch simulation
    length,  v_reel_out  = y[end-1],  y[end]
    lengthd, v_reel_outd = yd[end-1], yd[end]
    # extract the data of the particles
    y_  = @view y[begin:end-2]
    yd_ = @view yd[begin:end-2]
    # unpack the vectors y and yd
    part  = reshape(SVector{T}(y_),  Size(3, div(T,6), 2))
    partd = reshape(SVector{T}(yd_), Size(3, div(T,6), 2))
    pos1, vel1 = part[:,:,1], part[:,:,2]
    pos = SVector{div(T,6)+1}(if i==1 SVector(0.0,0,0) else SVector(pos1[:,i-1]) end for i in 1:div(T,6)+1)
    vel = SVector{div(T,6)+1}(if i==1 SVector(0.0,0,0) else SVector(vel1[:,i-1]) end for i in 1:div(T,6)+1)
    posd1, veld1 = partd[:,:,1], partd[:,:,2]
    posd = SVector{div(T,6)+1}(if i==1 SVector(0.0,0,0) else SVector(posd1[:,i-1]) end for i in 1:div(T,6)+1)
    veld = SVector{div(T,6)+1}(if i==1 SVector(0.0,0,0) else SVector(veld1[:,i-1]) end for i in 1:div(T,6)+1)

Perhaps it gives you an idea…

The docs recommend indexing with a StaticVector, since the compiler is otherwise ignorant of output shape without the type-domain size information carried by StaticArrays. That could look like this:

julia> x[SA[2, 5]]
2-element SVector{2, Float64} with indices SOneTo(2):
 1.2
 1.5

julia> x[SA[2:3...]]
2-element SVector{2, Float64} with indices SOneTo(2):
 1.2
 1.3
4 Likes

What does this mean?

Without splatting (...), you’d have a one-element vector of eltype UnitRange:

julia> SA[2:3]
1-element SVector{1, UnitRange{Int64}} with indices SOneTo(1):
 2:3

julia> SA[2:3...]
2-element SVector{2, Int64} with indices SOneTo(2):
 2
 3

It might be nice to add a convenience indexing syntax to the @SArray interface, although I think it’d be a breaking change?

julia> @SArray x[2:3]
ERROR: ArgumentError: Static Array parameter T must be a type, got [1.1, 1.2, 1.3, 1.4, 1.5]
Stacktrace:

One thing you should avoid is doing the above. Here you create an ordinary vector, and then convert it to an SVector.

Remove the brackets, then you don’t need to include the length either:

SVector(1.1, 1.2, 1.3, 1.4, 1.5)

because the number and the type of inputs is known to the compiler:

julia> @btime SVector{5}([1.1, 1.2, 1.3, 1.4, 1.5]);
  38.889 ns (1 allocation: 96 bytes)

julia> @btime SVector(1.1, 1.2, 1.3, 1.4, 1.5);
  2.000 ns (0 allocations: 0 bytes)

As for indexing, how can the compiler know the size (and hence the type) of the output array from an the input argument? Can you answer this:

What should the output type of getindex(::SVector{5, Float64}, ::Vector{Int}) be? If you cannot answer it, how can the compiler?

On the other hand: What is the output type of getindex(::SVector{5, Float64}, ::Svector{3, Int})? That you can answer.

4 Likes

Great tips! I understand your point about Vector vs SVector for getindex, but what about UnitRange? Could you direct getindex(::SVector{5, Float64}, ::UnitRange) to return an SVector?

Perhaps not, since UnitRange is just a Tuple and the size isn’t in the type.

My motivation was the desire to avoid needing to write an additional method for anything in a codebase that did indexing.

Is there any syntax that works for StaticArrays that also happens to work for a normal array?

Exactly, you ask yourself what the type of getindex(::SVector{5, Float64}, ::UnitRange) is. There’s no answer that includes the size.

I suppose the design is intended to better maintain type stability? Because a getindex-like method function myget(x::SVector, i) x[SA[i...]] end is type-unstable (same input types can result in multiple different return types). Using a dynamically sized index i to make a SVector must require type instability somewhere, but you have better control over type instabilities in your own code (function barriers) than in StaticArrays.jl’s code.

1 Like