StaticArrays and Statsbase

Viewed 57

I want to use StaticArray with StatsBase. Consider the following function

function update_weights_1(N, M)
    weights_vector_to_update      = ones(N) / N

    wvector = Weights(weights_vector_to_update, 1)

    res = [0.0]
    for m in 1:M
        sample!(M_vector, wvector, res)
    end
end

function update_weights_2(N, M)
    weights_vector_to_update      = ones(N) / N

    res = [0.0]
    for m in 1:M
        sample!(M_vector, Weights(weights_vector_to_update, 1), res)
    end
end

update_weights_1 requires substantially less memory allocation than update_weights_2 because Weights(weights_vector_to_update, 1) needs memory allocation. However, suppose I have a list of small vectors, say z,

z = [ones(3) / 3 for i in 1:10000]

and this function

function update_weights_3(z,M)

    N   = size(z[1],1)
    M_vector = 1:N
    for i in 1:size(z,1)
        rand!(z[i])
        res = [0.0]
        for m in 1:M
            sample!(M_vector, Weights(z[i]), res)
        end
    end
end

update_weights_3(z,1000) allocates a lot of memory. I know that using StaticArrays for z can significantly speed up the code and reduce memory allocation. However, following the procedure in this post, whenever I wrap Weights around a StaticArray, it creates memory.

Would you know how to apply StaticArray in this case? Essentially I have a collection of small arrays that I would like to transform into Weights.

1 Answers

Weights is a mutable type, which can cause unnecessary heap allocations (sometimes they are stack allocated... I don't fully understand when this optimization happens). You can define your own immutable weights type, though:

struct StaticWeights{S<:Real, T<:Real, N, V<: StaticVector{N, T}} <: AbstractWeights{S, T, V}
    values::V
    sum::S
end
StaticWeights(values) = StaticWeights(values, sum(values))

Used in your example:

function update_weights_3(z,M)
    N   = size(z[1],1)
    M_vector = 1:N
    for i in 1:size(z,1)
        rand!(z[i])
        res = [0.0]
        for m in 1:M
            sample!(M_vector, StaticWeights(z[i]), res)
        end
    end
end

With this change I don't see any allocations in the inner loop.

Related