Allocation-free weighted samples

General Usage Performance
, ,
8

Assuming the type (including length for an SVector) of the weights does not change, you could exploit the mutability of Weights to update it in-place:

using Accessors, StaticArrays, StatsBase, BenchmarkTools
StatsBase.weights(w::Weights) = w
function plusone(wghts, res)
    for i in 1:100
        # Update in any way compatible with typeof(wghts).
        # Here this is Weights{Float64, Float64, SVector{16, Float64}}.
        # (The parameters are the types of the sum, the entries, and the weights vector.)
        wghts.values = @SVector rand(length(wghts.values))  
        wghts.sum = sum(wghts.values) 
        idx = wsample(wghts)
        @reset res[idx] += 1
    end
   return res
end

wts = @SVector rand(16)
wghts = Weights(wts)  # In the example wts is not used directly, but fixes the type
res = @SVector fill(0, 16)
@btime plusone($wghts, $res);
    # 5.250 μs (0 allocations: 0 bytes)

To avoid forgetting to update wghts.sum, you could also use a function

function update!(w::Weights{S, T, V}, new_wts::V) where {S, T, V}
    w.values = new_wts
    w.sum = sum(w.values)
end
2 Likes