Reducing the allocations of my function

Here’s a version that fills a preallocated array:

using StaticArrays
using LinearAlgebra: norm

function randv!(vreturn::AbstractMatrix, v, vstep)
    Base.require_one_based_indexing(vreturn)
    (num, n) = size(vreturn)
    n == 3 || throw(ArgumentError("vreturn must have 3 columns"))
    vlim = round(v/vstep)
    vrange = -vlim:vlim
    vcount = 0
    vtest = 0.01 * v
    while vcount < num
        vrand = SVector{3,Float64}(rand(vrange) for _ in 1:3)
        if abs(norm(vrand) * vstep - v) < vtest
            vcount += 1
            vreturn[vcount, :] .= vrand * vstep
        end
    end
    return vreturn
end

It doesn’t allocate and is reasonably fast:

julia> using BenchmarkTools

julia> vreturn = zeros(10, 3);

julia> @btime randv!(vreturn, 1.0, 0.01)
  4.980 μs (0 allocations: 0 bytes)
10×3 Matrix{Float64}:
  0.85   0.51   0.02
 -0.73  -0.18  -0.67
  0.57   0.12   0.81
  0.95   0.12  -0.3
  0.44  -0.87   0.25
  0.71   0.67  -0.24
 -0.26  -0.56   0.78
 -0.68  -0.15  -0.72
 -0.83  -0.55  -0.15
  0.85   0.29   0.43

Edit: Modified my original post to add a bang to the function name and put the modified argument first, plus a couple other minor edits.

4 Likes