Create SVector with one non-zero value

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])

I tried the trick described here in the docs:

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

You can probably use the function ntuple

2 Likes

It worked, many thanks!

For the record, the correct function is:

# 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.

You can also try FillArrays.jl OneElement type, which does exactly what your fc(i, ::ord) function does:

using FilArrays
OneElement(i,order) # Vector of lenght order, where everything is zero except for index i, which is 1

That might be beneficial since it also knows about the sparse structure and won’t try to multiply the zero elements with something else :slight_smile:

2 Likes

Actually, it wouldn’t work in my case because I need the SVector type, but I will keep that in mind for other applications!

You may convert a OneElement to a SVector:

julia> using FillArrays, StaticArrays

julia> f(i) = SVector{3}(OneElement(i, 3))
f (generic function with 1 method)

julia> f(2)
3-element SVector{3, Int64} with indices SOneTo(3):
 0
 1
 0

julia> @btime f($(Ref(2))[])
  2.954 ns (0 allocations: 0 bytes)
3-element SVector{3, Int64} with indices SOneTo(3):
 0
 1
 0
1 Like

I believe the canonical solution is to use StaticArrays.setindex (without !), like this:

fc(i::Int) = setindex(zeros(SVector{10, Float64}), 1, i)

No other libraries necessary.

If you want the size to be configurable, use the regular Val provided by Base instead of a custom Order:

fc(::Val{N}, i::Int) where {N} = setindex(zeros(SVector{N, Float64}), 1, i)

Example:

julia> fc(Val(5), 2)
5-element SVector{5, Float64} with indices SOneTo(5):
 0.0
 1.0
 0.0
 0.0
 0.0
2 Likes

This seems like “cleaner” Julia code indeed.

I used Order instead of Val in my original code to make the code more readable; I will just write Order = Val to keep it this way.

You can do that, but Val is the standard way, and is easily recognizable for other users. So personally, I would stick with Val.

1 Like

Oh, BTW, if you absolutely prefer using Order, you should write

const Order = Val

to make the binding constant, otherwise you may experience performance issues.

1 Like

Thank you for the answer!! I will think about what’s best for my use case.

If it suits you, defining your own StaticArrays type is dead simple:

using StaticArrays

struct OneHotSVector{N,T} <: StaticVector{N,T}
	val::T
	ind::Int
end

OneHotSVector{N}(val::T, ind) where {N,T} = OneHotSVector{N,T}(val, ind) # autodetect the eltype

Base.@propagate_inbounds function Base.getindex(v::OneHotSVector, i::Int)
	@boundscheck checkbounds(v, i)
	return i==v.ind ? v.val : zero(v.val)
end

julia> OneHotSVector{3}(1.0, 2) * OneHotSVector{2}(1.0, 1)'
3Ă—2 SMatrix{3, 2, Float64, 6} with indices SOneTo(3)Ă—SOneTo(2):
 0.0  0.0
 1.0  0.0
 0.0  0.0
4 Likes