Large allocation problem while simulating agents model

Hello, This is a continuation of previous thread multithreading in agents. I have put an example here which overall is weird as a model but gives idea of memory allocation regions in my original model. Please ignore if you feel im collecting same data again and again :sweat_smile:. Its not the case in original model, all these collections are required. I just wanted to setup example in a way that there are allocations.

Original model is more complex but with this example I have pointed out locations where I feel issue reside.

using Agents, Random, StaticArrays

@agent struct SocialAgent(ContinuousAgent{2, Float64})
    mass::Float64
    count::Int64
    status::Symbol
    type::Symbol
end

function ball_model(; speed = 0.002)
    space2d = ContinuousSpace((1, 1); spacing = 0.02)
    model = StandardABM(SocialAgent, space2d; agent_step!, properties = Dict(:dt => 1.0, :xx => []),
                        rng = MersenneTwister(42))
    # And add some agents to the model
    for ind in 1:500
        pos = Tuple(rand(abmrng(model), 2))
        type = ind ≀ 0.5*500 ? :A : :B   
        status = type == :A ? :yes : :no 
        count =  0
        vel = sincos(2Ο€ * rand(abmrng(model))) .* speed
        add_agent!(pos, model, vel, 1.0, count, status, type)
    end
    return model
end

function interact!(a1, a2)
      ### Choosing Interaction based on cell type and status of Cell 

    if (a2.type == :A) && (a1.type == :B )

   
        a1.status = :blank
         a1.vel = @SVector[0.0,0.0]

    end

    if (a2.type == :B) && (a1.type == :A)

        a2.status = :blank
         a2.vel = @SVector[0.0,0.0]
    end

end

function rule!(agent, model)

    if agent.status == :yes
        agent.count += 1
        push!(model.xx, (agent.id, abmtime(model), agent.count, agent.status))

        if agent.count == 5

            push!(model.xx, (agent.id, abmtime(model), agent.count, agent.status))
            agent.count = 0

        end
    end

    if agent.status == :no
        agent.count += 1
        push!(model.xx, (agent.id, abmtime(model), agent.count, agent.status))

        if agent.count == 5

            push!(model.xx, (agent.id, abmtime(model), agent.count, agent.status))
            agent.count = 0
            agent.status == :blank
       
        end
    end

    if agent.status == :blank
        agent.count += 1
        push!(model.xx, (agent.id, abmtime(model), agent.count, agent.status))

        if agent.count == 5

            push!(model.xx, (agent.id, abmtime(model), agent.count, agent.status))
            agent.count = 0
            agent.status == :no
            
        end
    end
end

function agent_step!(agent, model) 

    rule!(agent, model)

    move_agent!(agent, model, model.dt)

end

function model_step!(model)
    for (a1, a2) in interacting_pairs(model, 0.012, :nearest)
        interact!(a1, a2)
        elastic_collision!(a1, a2, :mass)
    end
end


model = ball_model()


function runn(model)

    x = zeros(100)

    for i in 1:100

        model = ball_model()
        run!(model, 160)

        x[i] = sum(stack(model.xx, dims = 1 )[:,3])   

    end

    return x
end

@time runn(model)
julia> @time runn(model)
  4.937694 seconds (63.26 M allocations: 4.141 GiB, 4.96% gc time)

Thank you for help :slight_smile:

Note that

Tuple(rand(abmrng(model), 2))

is not stable.

julia> function f()
       return Tuple(rand(2))
       end
f (generic function with 1 method)

julia> function g()
       return (rand(), rand())
       end
g (generic function with 1 method)

julia> @code_warntype f()
MethodInstance for f()
  from f() @ Main REPL[6]:1
Arguments
  #self#::Core.Const(f)
Body::Tuple{Vararg{Float64}}
1 ─ %1 = Main.rand(2)::Vector{Float64}
β”‚   %2 = Main.Tuple(%1)::Tuple{Vararg{Float64}}
└──      return %2


julia> @code_warntype g()
MethodInstance for g()
  from g() @ Main REPL[7]:1
Arguments
  #self#::Core.Const(g)
Body::Tuple{Float64, Float64}
1 ─ %1 = Main.rand()::Float64
β”‚   %2 = Main.rand()::Float64
β”‚   %3 = Core.tuple(%1, %2)::Tuple{Float64, Float64}
└──      return %3

(The highlighting is not captured here, but note that the types in @code_warntype f() show a lot of red - the varargs are not good.)

Try making the change so that it’s (rand(abmrng(model)), rand(abmrng(model))) and see what happens?

2 Likes

After change : pos = Tuple(rand(abmrng(model), 2)) to pos = (rand(abmrng(model)), rand(abmrng(model)))

julia> @time runn(model)
  4.949795 seconds (63.01 M allocations: 4.133 GiB, 4.44% gc time)

Profile for same…

        - #using Agents, Random, StaticArrays
        - 
        - @agent struct SocialAgent(ContinuousAgent{2, Float64})
        -     mass::Float64
        -     count::Int64
        -     status::Symbol
        -     type::Symbol
        - end
        - 
        0 function ball_model(; speed = 0.002)
        0     space2d = ContinuousSpace((1, 1); spacing = 0.02)
    41600     model = StandardABM(SocialAgent, space2d; agent_step!, properties = Dict(:dt => 1.0, :xx => []),
        -                         rng = MersenneTwister(42))
        -     # And add some agents to the model
        0     for ind in 1:500
        0         pos = (rand(abmrng(model)), rand(abmrng(model)))
        0         type = ind ≀ 0.5*500 ? :A : :B   
        0         status = type == :A ? :yes : :no 
        -         count =  0
        0         vel = sincos(2Ο€ * rand(abmrng(model))) .* speed
  7584000         add_agent!(pos, model, vel, 1.0, count, status, type)
        0     end
        0     return model
        - end
        - 
        - function interact!(a1, a2)
        -       ### Choosing Interaction based on cell type and status of Cell 
        - 
        -     if (a2.type == :A) && (a1.type == :B )
        - 
        -    
        -         a1.status = :blank
        -          a1.vel = @SVector[0.0,0.0]
        - 
        -     end
        - 
        -     if (a2.type == :B) && (a1.type == :A)
        - 
        -         a2.status = :blank
        -          a2.vel = @SVector[0.0,0.0]
        -     end
        - 
        - end
        - 
        - function rule!(agent, model)
        - 
        0     if agent.status == :yes
        0         agent.count += 1
201214400         push!(model.xx, (agent.id, abmtime(model), agent.count, agent.status))
        - 
        0         if agent.count == 5
        - 
146828800             push!(model.xx, (agent.id, abmtime(model), agent.count, agent.status))
        0             agent.count = 0
        - 
        -         end
        -     end
        - 
        0     if agent.status == :no
        0         agent.count += 1
245068800         push!(model.xx, (agent.id, abmtime(model), agent.count, agent.status))
        - 
        0         if agent.count == 5
        - 
 59353600             push!(model.xx, (agent.id, abmtime(model), agent.count, agent.status))
        0             agent.count = 0
        -             agent.status == :blank
        -        
        -         end
        -     end
        - 
        0     if agent.status == :blank
        0         agent.count += 1
        0         push!(model.xx, (agent.id, abmtime(model), agent.count, agent.status))
        - 
        0         if agent.count == 5
        - 
        0             push!(model.xx, (agent.id, abmtime(model), agent.count, agent.status))
        0             agent.count = 0
        0             agent.status == :no
        -             
        -         end
        -     end
        - end
        - 
        - 
        - function agent_step!(agent, model) 
        - 
        0     rule!(agent, model)
        - 
1408000000     move_agent!(agent, model, model.dt)
        - 
        - end
        - 
        - function model_step!(model)
        -     for (a1, a2) in interacting_pairs(model, 0.012, :nearest)
        -         interact!(a1, a2)
        -         elastic_collision!(a1, a2, :mass)
        -     end
        - end
        - 
        - 
        - model = ball_model()
        - 
        - 
        - function runn(model)
        - 
      896     x = zeros(100)
        - 
        0     for i in 1:100
        - 
        0         model = ball_model()
        0         run!(model, 160)
        - 
     1600         x[i] = sum(stack(model.xx, dims = 1 )[:,3])   
        - 
        0     end
        - 
        0     return x
        - end
        - 
        - #@time runn(model)

Changing :xx => [] to :xx => Tuple{Int,Int,Int,Symbol}[] helps a little bit also, but not all of it.

properties = Dict{Symbol, Union{Float64, Vector{Tuple{Int,Int,Int,Symbol}}}}(:dt => 1.0, :xx => Tuple{Int,Int,Int,Symbol}[])

is a bit better .

1 Like

Looking at

using Agents, Random, StaticArrays

@agent struct SocialAgent(ContinuousAgent{2, Float64})
    mass::Float64
    count::Int64
    status::Symbol
    type::Symbol
end

function ball_model(; speed = 0.002)
    space2d = ContinuousSpace((1, 1); spacing = 0.02)
    model = StandardABM(SocialAgent, space2d; agent_step!, properties = Dict{Symbol, Union{Float64, Vector{Tuple{Int,Int,Int,Symbol}}}}(:dt => 1.0, :xx => Tuple{Int,Int,Int,Symbol}[]),
                        rng = MersenneTwister(42))
    # And add some agents to the model
    for ind in 1:500
        pos = (rand(abmrng(model)), rand(abmrng(model)))
        type = ind ≀ 0.5*500 ? :A : :B   
        status = type == :A ? :yes : :no 
        count =  0
        vel = sincos(2Ο€ * rand(abmrng(model))) .* speed
        add_agent!(pos, model, vel, 1.0, count, status, type)
    end
    return model
end

function interact!(a1, a2)
      ### Choosing Interaction based on cell type and status of Cell 

    if (a2.type == :A) && (a1.type == :B )

   
        a1.status = :blank
         a1.vel = @SVector[0.0,0.0]

    end

    if (a2.type == :B) && (a1.type == :A)

        a2.status = :blank
         a2.vel = @SVector[0.0,0.0]
    end

end

function _rule!(agent, model, xx)
    if agent.status == :yes
        agent.count += 1
        push!(xx, (agent.id, abmtime(model), agent.count, agent.status))

        if agent.count == 5

            push!(xx, (agent.id, abmtime(model), agent.count, agent.status))
            agent.count = 0

        end
    end

    if agent.status == :no
        agent.count += 1
        push!(xx, (agent.id, abmtime(model), agent.count, agent.status))

        if agent.count == 5

            push!(xx, (agent.id, abmtime(model), agent.count, agent.status))
            agent.count = 0
            agent.status == :blank
       
        end
    end

    if agent.status == :blank
        agent.count += 1
        push!(xx, (agent.id, abmtime(model), agent.count, agent.status))
        
        if agent.count == 5
            push!(xx, (agent.id, abmtime(model), agent.count, agent.status))
            agent.count = 0
            agent.status == :no
        end
    end
end
function rule!(agent, model)
    return _rule!(agent, model, model.xx::Vector{Tuple{Int,Int,Int,Symbol}})
end

function agent_step!(agent, model) 

    rule!(agent, model)

    move_agent!(agent, model, model.dt::Float64)

end

function model_step!(model)
    for (a1, a2) in interacting_pairs(model, 0.012, :nearest)
        interact!(a1, a2)
        elastic_collision!(a1, a2, :mass)
    end
end


model = ball_model()

function __build_x!(x, model, xx, i)
    x[i] = sum(view(stack(xx, dims = 1 ), :,3))
    return x
end
function _build_x!(x, model, i)
    xx = model.xx
    return __build_x!(x, model, xx, i)
end
function runn(model)

    x = zeros(100)

    for i in 1:100

        model = ball_model()
        run!(model, 160)

        #_build_x!(x, model, i)

    end

    return x
end

@time runn(model)

(note that I commented out the part where you set x), I get

0.823700 seconds (554.80 k allocations: 585.380 MiB, 8.13% gc time)

while with the x part uncommented and put behind a barrier,

1.606706 seconds (35.81 M allocations: 2.241 GiB, 17.70% gc time)

So, that is what you should focus on.

1 Like

I just realised now what you were doing with stack. I would just loop over and sum. Since getindex would be a union in this case since the Tuple is not all of the same type, wrapping the xx entries in a struct (or a named tuple) is appropriate.

using Agents, Random, StaticArrays

@agent struct SocialAgent(ContinuousAgent{2, Float64})
    mass::Float64
    count::Int64
    status::Symbol
    type::Symbol
end

struct XX 
    id::Int
    time::Int
    count::Int
    status::Symbol
end

function ball_model(; speed = 0.002)
    space2d = ContinuousSpace((1, 1); spacing = 0.02)
    model = StandardABM(SocialAgent, space2d; agent_step!, properties = Dict{Symbol, Union{Float64, Vector{XX}}}(:dt => 1.0, :xx => XX[]),
                        rng = MersenneTwister(42))
    # And add some agents to the model
    for ind in 1:500
        pos = (rand(abmrng(model)), rand(abmrng(model)))
        type = ind ≀ 0.5*500 ? :A : :B   
        status = type == :A ? :yes : :no 
        count =  0
        vel = sincos(2Ο€ * rand(abmrng(model))) .* speed
        add_agent!(pos, model, vel, 1.0, count, status, type)
    end
    return model
end

function interact!(a1, a2)
      ### Choosing Interaction based on cell type and status of Cell 

    if (a2.type == :A) && (a1.type == :B )

   
        a1.status = :blank
         a1.vel = @SVector[0.0,0.0]

    end

    if (a2.type == :B) && (a1.type == :A)

        a2.status = :blank
         a2.vel = @SVector[0.0,0.0]
    end

end

function _rule!(agent, model, xx)
    if agent.status == :yes
        agent.count += 1
        push!(xx, XX(agent.id, abmtime(model), agent.count, agent.status))

        if agent.count == 5

            push!(xx, XX(agent.id, abmtime(model), agent.count, agent.status))
            agent.count = 0

        end
    end

    if agent.status == :no
        agent.count += 1
        push!(xx, XX(agent.id, abmtime(model), agent.count, agent.status))

        if agent.count == 5

            push!(xx, XX(agent.id, abmtime(model), agent.count, agent.status))
            agent.count = 0
            agent.status == :blank
       
        end
    end

    if agent.status == :blank
        agent.count += 1
        push!(xx, XX(agent.id, abmtime(model), agent.count, agent.status))
        
        if agent.count == 5
            push!(xx, XX(agent.id, abmtime(model), agent.count, agent.status))
            agent.count = 0
            agent.status == :no
        end
    end
end
function rule!(agent, model)
    return _rule!(agent, model, model.xx::Vector{XX})
end

function agent_step!(agent, model) 

    rule!(agent, model)

    move_agent!(agent, model, model.dt::Float64)

end

function model_step!(model)
    for (a1, a2) in interacting_pairs(model, 0.012, :nearest)
        interact!(a1, a2)
        elastic_collision!(a1, a2, :mass)
    end
end


model = ball_model()

function __build_x!(x, model, xx, i)
    x[i] = 0.0
    for j in eachindex(xx)
        x[i] += xx[j].count
    end
    return x
end
function _build_x!(x, model, i)
    xx = model.xx
    return __build_x!(x, model, xx, i)
end
function runn(model)

    x = zeros(100)

    for i in 1:100

        model = ball_model()
        run!(model, 160)

        _build_x!(x, model, i)

    end

    return x
end

@time runn(model)
0.809722 seconds (554.80 k allocations: 585.380 MiB, 6.08% gc time)

Still some work that could be done but it’s a good reduction.

(Your stack was creating a Matrix{Any}. That is extremely slow.) Profiling I think shows that most of the remaining allocations are from push!, though I don’t know what type of allocations are typical with Agents.jl’s internals.

1 Like

I will try to implement this :slight_smile: on my model. Thank you:)

Instead of

 properties = Dict{Symbol, Union{Float64, Vector{XX}}}(:dt => 1.0, :xx => XX[])

you could use

properties = (dt = 1.0, xx = XX[])

this gives a little improvement.

1 Like

At the moment I have it like this

mutable struct prop
    x::Float64
    y::Float64
    z::Float64
    w::Float64
end

and then in model initiation function I give it values


function model_init(
    x ::Float64 = 3.1,
    y ::Float64 = 0.1,
    z ::Float64 = 2.0,
    w ::Float64= 3.0)
    
    .
    .
    .
    properties = prop(x, y, z, w)
    .
    .
end
1 Like

I think that is also good if you need to change those model properties during the simulation :slight_smile:

1 Like

In my actual model profile I see a large allocation for model property model.count but i dont understand why its allocating so much. I can’t share whole script but can you guys probably point why it could be the case? As its just increment. Is it because of model.count =0 term beacuse it will be read again and again so allocation?

        - function model_step_mimic!(model)

        - 
 37119792     model.count += 1                        
        - 
        0     if model.count == nagents(model)
        - 
        0         x = round(abmtime(model), digits = 1)
        - 
   307200         if  x in model.ss[2:end, 1]
        - 
    37120             model.pkl = model.ss[Int64(x)+1, 4] * n
        -         end
        - 
   371200         model.count = 0
        -      end
        - end


I pass this function inside this one which runs for all agents for every single time-point:

   - function main_func!(agent, model)
        -     model_step_mimic!(model)
        -     some_other_func!(agent, model)
   -  end

The only thing I could think of is that your model.count is not properly annotated, maybe it is considered Any? you could look at the output of @code_warntype

1 Like