# Allocation-free weighted samples

**URL:** <https://discourse.julialang.org/t/allocation-free-weighted-samples/119265>\
**Category:** Performance\
**Tags:** memory-allocation, staticarrays, sampling\
**Created:** [September 11, 2024, 4:55am UTC](https://discourse.julialang.org/t/allocation-free-weighted-samples/119265 "2024-09-11T04:55:38Z")\
**Posts on this page:** 13\
**Page:** 1

<div class="post-metadata">

**Author:** ![Tetrakai](https://avatars.discourse-cdn.com/v4/letter/t/4da419/32.png) [@Tetrakai](https://discourse.julialang.org/u/Tetrakai)\
**Post date:** [September 11, 2024, 4:55am UTC](https://discourse.julialang.org/t/allocation-free-weighted-samples/119265/1 "2024-09-11T04:55:38Z")

</div>

Given an `SVector` of weights, is there a way to take a single weighted sample without any allocations? Ie, these work great but require an allocation:

```julia
x = [1]
wts = @SVector rand(16)
@btime wsample(1:16, $wts)
@btime wsample!(1:16, $wts, $x)[1]

```

```julia
julia> @btime wsample(1:16, $wts)
  52.173 ns (1 allocation: 144 bytes)
13

julia> @btime wsample!(1:16, $wts, $x)[1]
  62.325 ns (1 allocation: 144 bytes)
3

```

---

<div class="post-metadata">

**Author:** ![hz-xiaxz](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/hz-xiaxz/32/209585_2.png) [@hz-xiaxz](https://discourse.julialang.org/u/hz-xiaxz)\
**Post date:** [September 11, 2024, 6:20am UTC](https://discourse.julialang.org/t/allocation-free-weighted-samples/119265/2 "2024-09-11T06:20:01Z")

</div>

In my simple test, using staticarray increases allocation? I don’t know why.

```julia
julia> @btime wsample(1:16, $wts)
  27.587 ns (1 allocation: 144 bytes)
9
julia> wt = rand(16)
julia> @btime wsample(1:16, $wt)
  29.242 ns (1 allocation: 32 bytes)
6

```

---

<div class="post-metadata">

**Author:** ![Benny](https://avatars.discourse-cdn.com/v4/letter/b/49beb7/32.png) [@Benny](https://discourse.julialang.org/u/Benny)\
**Post date:** [September 11, 2024, 6:40am UTC](https://discourse.julialang.org/t/allocation-free-weighted-samples/119265/3 "2024-09-11T06:40:54Z")

</div>

I think it’s the instantiation of a mutable `Weights` instance

```julia
julia> @btime StatsBase.Weights($wts);
  14.358 ns (1 allocation: 144 bytes)

julia> ismutabletype(Weights) # usually uses 8 bytes for pointer
true

julia> sizeof(typeof( StatsBase.Weights(wts))) # 8+136=144
136

```

A bit of piracy can elide the instantiation if already provided a preallocated instance, but I don’t know if this is sound

```julia
julia> StatsBase.weights(w::Weights) = w

julia> @btime wsample(1:16, $(Weights(wts)))
  37.277 ns (0 allocations: 0 bytes)
6

```

---

<div class="post-metadata">

**Author:** ![Tetrakai](https://avatars.discourse-cdn.com/v4/letter/t/4da419/32.png) [@Tetrakai](https://discourse.julialang.org/u/Tetrakai)\
**Post date:** [September 11, 2024, 2:59pm UTC](https://discourse.julialang.org/t/allocation-free-weighted-samples/119265/4 "2024-09-11T14:59:25Z")

</div>

Thanks, can you explain how this works?

**Edit:**  
I quickly tried it in my function but it had no effect on the allocations. Likely I’m doing it wrong. Eg, this still allocates once per iteration:

```julia
using Accessors, StaticArrays, StatsBase, BenchmarkTools
StatsBase.weights(w::Weights) = w
function plusone(wts, res)
    for i in 1:100
        idx = wsample(Weights(wts))
        @reset res[idx] += 1
    end
        return res
    end

wts = @SVector rand(16)
res = @SVector fill(0, 16)
@btime plusone($wts, $res)

```

---

<div class="post-metadata">

**Author:** ![eldee](https://avatars.discourse-cdn.com/v4/letter/e/b5a626/32.png) [@eldee](https://discourse.julialang.org/u/eldee)\
**Post date:** [September 11, 2024, 5:45pm UTC](https://discourse.julialang.org/t/allocation-free-weighted-samples/119265/5 "2024-09-11T17:45:34Z")

</div>

If you look into the source code, you’ll see that `wsample(w)` calls `sample(default_rng(), weights(w))`, where `weights(w)` creates an instance of the mutable `Weights`. As @Benny pointed out, this allocates.

> **Relevant part of sampling.jl: lines 1044-1056**
>
> ```julia
> """
> wsample([rng], [a], w)
> 
> Select a weighted random sample of size 1 from `a` with probabilities proportional
> to the weights given in `w`. If `a` is not present, select a random weight from `w`.
> 
> Optionally specify a random number generator `rng` as the first argument
> (defaults to `Random.$(VERSION < v"1.3" ? "GLOBAL_RNG" : "default_rng()")`).
> """
> wsample(rng::AbstractRNG, w::AbstractVector{<:Real}) = sample(rng, weights(w))
> wsample(w::AbstractVector{<:Real}) = wsample(default_rng(), w)
> wsample(rng::AbstractRNG, a::AbstractArray, w::AbstractVector{<:Real}) = sample(rng, a, weights(w))
> wsample(a::AbstractArray, w::AbstractVector{<:Real}) = wsample(default_rng(), a, w)
> 
> ```

> **Relevant part of weights.jl: lines 4-23; 69; 82-89**
>
> ```julia
> """
> @weights name
> 
> Generates a new generic weight type with specified `name`, which subtypes `AbstractWeights`
> and stores the `values` (`V<:AbstractVector{<:Real}`) and `sum` (`S<:Real`).
> """
> macro weights(name)
> return quote
> mutable struct $name{S<:Real, T<:Real, V<:AbstractVector{T}} <: AbstractWeights{S, T, V}
> values::V
> sum::S
> function $(esc(name)){S, T, V}(values, sum) where {S<:Real, T<:Real, V<:AbstractVector{T}}
> isfinite(sum) || throw(ArgumentError("weights cannot contain Inf or NaN values"))
> return new{S, T, V}(values, sum)
> end
> end
> $(esc(name))(values::AbstractVector{T}, sum::S) where {S<:Real, T<:Real} = $(esc(name)){S, T, typeof(values)}(values, sum)
> $(esc(name))(values::AbstractVector{<:Real}) = $(esc(name))(values, sum(values))
> end
> end
> 
> @weights Weights
> 
> """
> weights(vs::AbstractArray{<:Real})
> 
> Construct a `Weights` vector from array `vs`.
> See the documentation for [`Weights`](@ref) for more details.
> """
> weights(vs::AbstractArray{<:Real}) = Weights(vec(vs))
> weights(vs::AbstractVector{<:Real}) = Weights(vs)
> 
> ```

By using `StatsBase.weights(w::Weights) = w` and directly supplying a `Weights` to `wsample`, you skip this instantiation. Note that creating our `Weights` instance still allocates. But `wsample(w)` does not if `w is Weights` .

So to improve your code snippet, you just need to move the `Weights(wts)` outside of the loop:

```julia
using Accessors, StaticArrays, StatsBase, BenchmarkTools
StatsBase.weights(w::Weights) = w
function plusone(wghts, res)
    for i in 1:100
        idx = wsample(wghts)
        @reset res[idx] += 1
    end
        return res
    end

wts = @SVector rand(16)
wghts = Weights(wts)
res = @SVector fill(0, 16)
@btime plusone($wghts, $res);
    # 2.400 μs (0 allocations: 0 bytes)

```

The reason why @Benny had no allocations in his `@btime wsample(1:16, $(Weights(wts)))` is because of the interpolation using `$`. For example,

```julia-repl
julia> @btime wsample(1:16, $(Weights(wts)));
  18.737 ns (0 allocations: 0 bytes)

julia> @btime wsample(1:16, Weights($wts));
  30.452 ns (1 allocation: 144 bytes)

julia> @btime $(rand(10^6));
  2.200 ns (0 allocations: 0 bytes)

```

---

<div class="post-metadata">

**Author:** ![eldee](https://avatars.discourse-cdn.com/v4/letter/e/b5a626/32.png) [@eldee](https://discourse.julialang.org/u/eldee)\
**Post date:** [September 11, 2024, 6:18pm UTC](https://discourse.julialang.org/t/allocation-free-weighted-samples/119265/6 "2024-09-11T18:18:33Z")

</div>

In `Weights(w)` we store `w` in the `values` field of the newly created `Weights` instance. The difference in allocation size is then the difference between the size of an `SVector` and a `Vector`. In the former case, we store the entries (128B), in the latter only a pointer (8B).  
I’m not completely sure about the details, but I assume the rest of the difference has to do with memory alignment.

---

<div class="post-metadata">

**Author:** ![Tetrakai](https://avatars.discourse-cdn.com/v4/letter/t/4da419/32.png) [@Tetrakai](https://discourse.julialang.org/u/Tetrakai)\
**Post date:** [September 11, 2024, 6:52pm UTC](https://discourse.julialang.org/t/allocation-free-weighted-samples/119265/7 "2024-09-11T18:52:28Z")

</div>

Thanks, the weights must be updated every iteration though. If i understand, using this method at some point I need to call `Weights` to use `wsample` within the loop. Ie, I can’t just pre-allocate a vector to store the weights then update that.

---

<div class="post-metadata">

**Author:** ![eldee](https://avatars.discourse-cdn.com/v4/letter/e/b5a626/32.png) [@eldee](https://discourse.julialang.org/u/eldee)\
**Post date:** [September 11, 2024, 7:18pm UTC](https://discourse.julialang.org/t/allocation-free-weighted-samples/119265/8 "2024-09-11T19:18:10Z")

</div>

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:

```julia
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

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

```

---

<div class="post-metadata">

**Author:** ![Tortar](https://avatars.discourse-cdn.com/v4/letter/t/6bbea6/32.png) [@Tortar](https://discourse.julialang.org/u/Tortar)\
**Post date:** [September 11, 2024, 7:27pm UTC](https://discourse.julialang.org/t/allocation-free-weighted-samples/119265/9 "2024-09-11T19:27:50Z")

</div>

Maybe an approach with reservoir sampling could fit your use case:

```julia
julia> using StreamSampling, BenchmarkTools

julia> function update_sample!(s, iter)
           for x in iter 
               update!(s, x, wts(x))
           end
           return s
       end;

julia> wts(x) = x;

julia> sample = ReservoirSample(Int, algAExpJ);

julia> iter = 1:16;

julia> @btime empty!(update_sample!($sample, $iter));
  62.166 ns (0 allocations: 0 bytes)

```

here you use a function for the weights (which can also point to an array if needed)

---

<div class="post-metadata">

**Author:** ![Tetrakai](https://avatars.discourse-cdn.com/v4/letter/t/4da419/32.png) [@Tetrakai](https://discourse.julialang.org/u/Tetrakai)\
**Post date:** [September 11, 2024, 10:18pm UTC](https://discourse.julialang.org/t/allocation-free-weighted-samples/119265/10 "2024-09-11T22:18:31Z")

</div>

I see, I was misusing `@reset`:

```julia
StatsBase.weights(w::Weights) = w
function plusone(wghts, res)
    for i in 1:100
        @reset wghts.values = @SVector rand(16)
        @reset wghts.sum = sum(wghts.values)
        idx = wsample(wghts)
        @reset res[idx] += 1
    end
    return res
end

wts = @SVector rand(16);
wghts = Weights(wts);
res = @SVector fill(0, 16);
@btime plusone($wghts, $res);

```

Thanks, seems this will work perfectly.

---

<div class="post-metadata">

**Author:** ![Tetrakai](https://avatars.discourse-cdn.com/v4/letter/t/4da419/32.png) [@Tetrakai](https://discourse.julialang.org/u/Tetrakai)\
**Post date:** [September 11, 2024, 10:30pm UTC](https://discourse.julialang.org/t/allocation-free-weighted-samples/119265/11 "2024-09-11T22:30:43Z")

</div>

I assume this would be a better solution given no type piracy though. I’ll try it out.

---

<div class="post-metadata">

**Author:** ![Benny](https://avatars.discourse-cdn.com/v4/letter/b/49beb7/32.png) [@Benny](https://discourse.julialang.org/u/Benny)\
**Post date:** [September 11, 2024, 11:23pm UTC](https://discourse.julialang.org/t/allocation-free-weighted-samples/119265/12 "2024-09-11T23:23:44Z")

</div>

I’m not actually sure why Weights is mutable. I know StatsBase started early in Julia’s history so it may have made a Weights wrapper assuming mutable AbstractArrays with variable sums stored in a mutable `sum` field, but support for immutable arrays seems justified. I can’t imagine it but if there’s an internal reason it needs separate mutable instances, then that type piracy can get dangerous.

---

<div class="post-metadata">

**Author:** ![abraunst](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/abraunst/32/6880_2.png) [@abraunst](https://discourse.julialang.org/u/abraunst)\
**Post date:** [September 12, 2024, 1:50pm UTC](https://discourse.julialang.org/t/allocation-free-weighted-samples/119265/13 "2024-09-12T13:50:12Z")

</div>

I suppose it depends on the use case, but you could also write a simple non-allocating sampler along these lines

```julia
function mysample(w)
    r = rand() * sum(w)
    for (i, x) in pairs(w)
         r -= x
         r < 0 && return i
    end
end

```

Note that if you only update a small number of the weights between samples, there may be faster alternatives.