I am trying to define a function fc(i) returning a fixed size SVector with zeros everywhere, and a 1 at position i. Is it possible to do this with no allocations, and type stability? My first two tries either were type unstable or allocated some memory.
using StaticArrays
# No allocation, type unstable
fc_unstable(i) = SVector{10, Float64}(
zeros(SVector{i}),
vcat(1, zeros(SVector{9-i}))
))
# Type stable, allocates 144 bytes
fc_allocating(i) = SVector{10, Float64}([x == i ? 1 : 0 for x in 1:10])
struct Order{N} end
Order(N) = Order{N}()
fc_order(::Order{i}) where {i} = SVector(vcat(
zeros(SVector{i}),
vcat(1, zeros(SVector{9-i}))
))
This, however, only works if in fc_order(Order(n)), n is known at compilation. When used in a loop in a function, it also provoked type instability (below, x is of type Any):
julia> function loop_fc(::Order{n}) where n
for i in 1:n
x = fc_order(Order(i))
end
end
loop_fc (generic function with 1 method)
julia> @code_warntype loop_fc(Order(3))
MethodInstance for loop_fc(::Order{3})
from loop_fc(::Order{n}) where n @ Main REPL[65]:1
Static Parameters
n = 3
Arguments
#self#::Core.Const(loop_fc)
_::Core.Const(Order{3}())
Locals
@_3::Union{Nothing, Tuple{Int64, Int64}}
i::Int64
x::Any
Body::Nothing
1 ─ %1 = (1:$(Expr(:static_parameter, 1)))::Core.Const(1:3)
│ (@_3 = Base.iterate(%1))
│ %3 = (@_3::Core.Const((1, 1)) === nothing)::Core.Const(false)
│ %4 = Base.not_int(%3)::Core.Const(true)
└── goto #4 if not %4
2 ┄ %6 = @_3::Tuple{Int64, Int64}
│ (i = Core.getfield(%6, 1))
│ %8 = Core.getfield(%6, 2)::Int64
│ %9 = Main.Order(i)::Order
│ (x = Main.fc_order(%9))
│ (@_3 = Base.iterate(%1, %8))
│ %12 = (@_3 === nothing)::Bool
│ %13 = Base.not_int(%12)::Bool
└── goto #4 if not %13
3 ─ goto #2
4 ┄ return nothing
# Type stable, no allocation
fc_tuple(i) = SVector( ntuple( x->x == i ? 1 : 0, Val(10) ) )
# With size as parameter
fc(i, ::Order{ord}) where {ord} = SVector(
ntuple(x->x == i ? 1 : 0, Val(ord))
)
EDIT: attention! ntuple(f, n) forgets the size information if n > 10. You have to write ntuple(f, Val(n)) to keep the size information. I edited the functions accordingly.