Dealing with type stability when nothing can be returned

Specific Domains Modelling & Simulations
, ,
1

I have a relatively complex Agent Based Model (ABM) that I am profiling. One function in particular (get_shark_pos) is dominating the runtime according to the flame graph. I am surprised because the function is very simple and not dealing with large data structures which are processed in many other parts of the model.

When I run @code_warntype on get_shark_pos I get the following yellow warnings:
::Union{Nothing, Vector{Tuple{Int64, Int64, Int64}}}
::Union{Nothing, Tuple{Fishes, Agents.IdIteratorState}}

I think this is because the function nearby_agents (used by get_shark_pos) can return different things depending on whether there are other agents of species :shark in visible range or not.

I thought creating an empty container (shark_pos = Tuple{Int, Int, Int}[]) to push to would help resolve this issue but it has not.

My question:

  • What is a better way to set up my get_shark_pos function to deal with the type stability warnings?

Below is a MWE.

using Agents

space = GridSpace((100, 100, 100))

@agent struct Fishes(GridAgent{3}) 
    species::Symbol
end

# agent stepping function
function fish_step!(fish::Fishes, model::StandardABM)
  shark_pos = get_shark_pos(fish, model)
  # do other stuff with shark pos
end

# define model
fish_model = StandardABM(
    Fishes,
    space;
    agent_step! = fish_step!
)

# add sharks
for n in 1:1e4
    add_agent!(fish_model, :shark)
end

# add dolphins
for n in 1:1e4
    add_agent!(fish_model, :dolphin)
end

# get the position of all :shark within range = 1
function get_shark_pos(fish::Fishes, model::StandardABM)
    near_all = nearby_agents(fish, model, 1)
    shark_pos = Tuple{Int, Int, Int}[]
    for x in near_all
        @inbounds if x.species === :shark
            push!(shark_pos, x.pos)
        end
    end
    return shark_pos
end 

# check type stability
@code_warntype get_shark_pos(fish_model[1], fish_model)
2

I don’t think these unions have a bad performance impact. These short unions are handled well by the compiler (3 or fewer is handled by very efficient code). Indeed, such unions with Nothing are created by every for loop.

What you can do is to preallocate the shark_pos array:

    shark_pos = Tuple{Int, Int, Int}[]
    sizehint!(shark_pos, length(near_all))
    for x in near_all
    ... # as before

This may allocate more memory than needed, but will avoid reallocation inside the loop. Whether it’s better depends on the number of entries which are sharks.

Alternatively you can rewrite get_shark_pos as:

function get_shark_pos(fish::Fishes, model::StandardABM)
    near_all = nearby_agents(fish, model, 1)
    return [x.pos for x in near_all if x.species === :shark]
end

I’m not sure if it will be faster.

Another possibility is to reuse a preallocated array, if that is doable in your program. I.e. keep a Tuple{Int,Int,Int}[] hanging around, and use a get_shark_pos! to fill it:

function get_shark_pos!(shark_pos, fish::Fishes, model::StandardABM)
    near_all = nearby_agents(fish, model, 1)
    empty!(shark_pos)
    for x in near_all
        (x.species === :shark) && push!(shark_pos, x.pos)
    end
    return shark_pos
end

If you always call the get_shark_pos! with the same shark_pos vector, there will be allocations only when the vector has to increase to something larger than it has been before (because empty! only sets the current size, not the allocated size). Of course, if you then need to save some positions the shark_pos must be copied, because it will be overwritten at every call.