How To Avoiding Allocations in Static Structs

Hi, I’m doing an MD simulation and have a problem getting good performance with static structs. The idea is to have a struct that represents a particle. I checked that using static structs is the fastest (nice) way of doing it, even if I have to override the particle when editing a field. I have the following implementation inspired by AtomsBase.jl.

using  StaticArrays

struct Particle{D}
    r::SVector{D, Float64} 
    v::SVector{D, Float64}
    a::SVector{D, Float64}
    m::Float64
end

This works great. But I have a lot of questions about the workflow. For example, I can make a function that edits a field, for example

function Set_r(P::Particle{D}, r::SVector{D, Float64})::Particle{D} where D
    Particle(r,P.v,P.a,P.m,P.rad)
end

@btime Set_r(P, SVector(0.0,0.0,0.0))
3.370 ns (0 allocations: 0 bytes)

This works similar with all the vectors, but when I do it for the mass

function Set_m(P::Particle{D}, m::Float64)::Particle{D} where D
    Particle(P.r,P.v,P.a,m,P.rad)
end

@btime Set_m(P,0.0)
20.917 ns (1 allocation: 96 bytes)

This, for me, is very weird. Why a scalar would have 1 allocation and a vector of 3 scalars won’t? Also, if I wore to add fields to the struct, suddenly the set methods will become very cumbersome to work with. Is there another way to improve this workflow?

At last, the idea is to have a vector of structs to run the simulation. Therefore I’ll be overriding the particles of the vector each time I have to update the position (for example). Is there a way to avoid allocations while doing this? What is the best way of doing this sort of thing?

@btime particles[1] = Set_r(P,SVector(0.0,0.0,0.0))
@btime particles[2] = Set_m(P,0.0)
  20.628 ns (1 allocation: 96 bytes)
  27.820 ns (1 allocation: 96 bytes)

I can’t reproduce it:

julia> using StaticArrays

julia> struct Particle{D}
           r::SVector{D, Float64}
           v::SVector{D, Float64}
           a::SVector{D, Float64}
           m::Float64
       end

julia> function Set_r(P::Particle{D}, r::SVector{D, Float64})::Particle{D} where D
           Particle(r,P.v,P.a,P.m)
       end;

julia> function Set_m(P::Particle{D}, m::Float64)::Particle{D} where D
           Particle(P.r,P.v,P.a,m)
       end;

julia> const P = Particle(rand(SVector{3}), rand(SVector{3}), rand(SVector{3}), rand());

julia> @btime Set_r(P, SVector(0.0,0.0,0.0));
  1.207 ns (0 allocations: 0 bytes)

julia> @btime Set_m(P, 0.0);
  1.207 ns (0 allocations: 0 bytes)

Perhaps you were using a non-const global P?

1 Like

I can’t reproduce it either, but in any case to properly benchmark you need to interpolate the variables, i. e., use @btime Set_m($P,0.0), with the $. I would bet that solves your problem.

For example, I don´t get the allocations, but the benchmark is quite different:

julia> @btime particles[1] = Set_r(P,SVector(0.0,0.0,0.0));
  15.816 ns (0 allocations: 0 bytes)

julia> @btime $particles[1] = Set_r($P,SVector(0.0,0.0,0.0))
  3.898 ns (0 allocations: 0 bytes)

(as a side note, if you are doing particle simulations, take a look at CellListMap :-), might be useful)

Thanks, @fredrikekre, and @lmiq. It’s weird that the benchmark is so different but in fact, I was using a non-const P. I think the reason is that I was using a Jupyter notebook to test things on the fly. But somehow it calms me a bit to know that code works like I thought it should.

I made some research about array management and may help someone in the future with the same problem. I stumbled across StructArrays.jl. I found that makes management very easy. Also, a time step would look like this

aos = [P for i = 1:400]
soa = StructArray(aos);

r = soa.r + soa.v*0.1
@btime map((x,y) -> x.r = y, LazyRows(soa), r)
1.699 μs (4 allocations: 9.67 KiB)

which is 75 times faster in my notebook than doing it in a for loop and only does 8 allocations. Impressive!

@lmiq I feel honored that you answered my question, I love your package, it’s great. I was planning to add CellListMap to the simulation.

1 Like

Regarding the workflow you could also take a look at Accessors.jl

3 Likes