Optional Struct-of-Arrays layout for agents

[Tortar]

While currently Agents.jl uses an Array-of-Structs (AoS) layout for Agents, it would be useful in some situations to have a Struct-of-Arrays (SoA), in particular when many operations are easily vectorizable, and the number of operations which are not are a small part of the total number of operations. I think it should be possible to use AoS or SoA just by building agents under the hood differently, but without changing any interface, or even any functions operating on agents inside the library. And we also have a precedent for this in #492. If @fbanning would like to share their opinion on this given that he worked on this (albeit a lot of time ago) it would be really appreciated. I think that SoA is bad when removal of agents are allowed, but even now we have a version of ABM models which don't allow for removals and have their uses. And actually not having removals is fine for my use case for SoA which is mainly to refactor BeforeIT.jl to be based on Agents.jl with an easier transition and without incurring any perf penalty.

[Datseris]

I think it should be possible to use AoS or SoA just by building agents under the hood differently, but without changing any interface, or even any functions operating on agents inside the library.

The only issue I see here is the user defined agent-stepping functions. Typically as a user you do something like a = model[id]; a.property = x. To make this non-breaking, you need to make model[id] return a special data structure, so that when doing a.property = x this actually updates the id-th entry in the property-th array in the struct-of-arrays. This doesn't sound hard at all actually, seems to me like a 10-20 lines of code helper data structure?

The same thing would happen when add_agent is called.

Worth doing a MWE that operates this way and extends the AgentBasedModel interface.

A downside I see: what happens with @agent? How does this (admidetly super duper nice) framework interplace with struct-of-arrays?

[Tortar]

mmmh yes actually the most problematic aspects I see of @agent is that it should work differently from how it works now (we can use an option in the macro for this) and most importantly it should define the same object as model[id] for creation, but everything seems doable

[Datseris]

I think a new macro with a new name would be more fitting here, that has similar enough functionality. @agent is already quite complex; arguably the most complex thing in Agents.jl. But this is something not really crucial for the start of this MWE.

[Datseris]

Actually, thinking about this more: we don't need a new macro. We just need a constructor of this new SoA-based ABM model that takes as input the same thing all other models take: either an agent type, or a union type, or a @multiagent output. It can extract the fields and just make a struct with fields the same name.

[Tortar]

that's okay but consider that model[id] should return the same thing as the creation of agents, and that means that we still need something which creates an agent differently not to have type instability, but given that we have a constructor for agents which accepts the model instance we cold make that constructor dispatch on the model type and be fine I think

[Datseris]

This is what I am arguing: agents are never actually created. Let's say the new model is called SoABM. Then

function Base.getindex(model::SoABM, id::Int)
    return SoAWrapper(model.container, id)
end

struct SoAWrapper{C}
    soa::C
    id::Int
end

function Base.getpropert(agent::SoAWrapper, name)
    return agent.soa[name][agent.id]
end

function Base.setproperty!(agent::SoAWrapper, name, x)
    agent.soa[name][agent.id] = x
    return agent
end

# Now this just works (TM)
function model_step!(model::SoABM)
    agent = model[25]
    agent.currency = 23
    return
end

so this has the same syntax Agents.jl has been using so far but actually utilizes this SoA approach.

We then can define a helper macro @agents_broadcast that takes the model, like

@agents_broacast model.property1 = model.property2 + 42 which does this broadcasting oepration across agent properties (not model properties), which you say it is advantageous and easily parallelizable in the SoA approach.

We should also consult #492 again, I don't remember it anymore besides the conclusion that it was not worth it. But maybe with this new approach I am proposing things are different.

[Tortar]

yeah, makes sense, and in the creation of the agent we could return the wrapper. Regarding #492 I think that removing agents will be a bit annoying, but with a "swap and pop" approach it should be feasible. At the same time, I think that probably it should be possible to have a layout which works well with a SoA approach, and yours seems fine, but having #492 is a bit of a red flag which I'm maybe wrong :-)

[Datseris]

The good thing is, now you really won't need to write more than 100 loc to have a fully runnable MWE that actually extends the ABM interface. So it doesn't cost much to try. Have the BeforeIT.jl team tested both versions or just happened to go with the SoA approach?

[Tortar]

No they didn't test both, but I think the SoA should be advantageous there because there are a lot of vectorized operations.

The cool aspect of all of this is that we can actually support a StructVector from StructArrays as a container and everything could automagically work

[Tortar]

e.g. this is sketch from the outside:

using Agents, Random, StructArrays, CommonSolve, LazyArrays

@agent struct Payer(NoSpaceAgent)
    alive::Bool
    age::Int
    wage::Float64
    wealth::Float64
end

struct StructArrayABM{
		S<:Agents.SpaceType, 
		A<:AbstractAgent,
		C,T,G,K,F,P,R} <: ABM{S}
    agents::C
    agent_step::G
    model_step::K
    space::S
    scheduler::F
    properties::P
    rng::R
    agents_types::T
    agents_first::Bool
    maxid::Base.RefValue{Int64}
    time::Base.RefValue{Int}
end
function StructArrayABM(
        ::Type{A}, space::S = nothing;
        agent_step!::G = dummystep,
    	model_step!::K = dummystep,
    	scheduler::F = Schedulers.fastest,
        properties::P = nothing,
        rng::R = Random.default_rng(),
        agents_first::Bool = true,
    ) where {A<:AbstractAgent,S<:Agents.SpaceType,G,K,F,P,R<:AbstractRNG}

	fields_with_types = zip(fieldnames(A), fieldtypes(A))
	tuple_init = NamedTuple(x => T[] for (x, T) in fields_with_types)
	agents = StructVector{A}(tuple_init)
	agents_types = Agents.union_types(A)
    T = typeof(agents_types)
	return StructArrayABM{S,A,typeof(agents),T,G,K,F,P,R}(agents, agent_step!, model_step!, space, scheduler,
               properties, rng, agents_types, agents_first, Ref(0), Ref(0))
end

struct AgentWrapperSoA{C}
    soa::C
    id::Int
end

function Base.getindex(model::StructArrayABM, id::Int)
    return AgentWrapperSoA(Agents.agent_container(model), id)
end

function Base.getproperty(agent::AgentWrapperSoA, name::Symbol)
    return getproperty(getfield(agent, :soa), name)[getfield(agent, :id)]
end

function Base.setproperty!(agent::AgentWrapperSoA, name::Symbol, x)
    getproperty(getfield(agent, :soa), name)[getfield(agent, :id)] = x
    return agent
end

Agents.nextid(model::StructArrayABM) = nagents(model) + 1
Agents.hasid(model::StructArrayABM, id::Int) = id ≤ nagents(model)

function Agents.add_agent_to_model!(agent::AbstractAgent, model::StructArrayABM)
    agent.id != nagents(model) + 1 && error(lazy"Cannot add agent of ID $(agent.id) in a vector ABM of $(nagents(model)) agents. Expected ID == $(nagents(model)+1).")
    tuple_agent = NamedTuple(x => getfield(agent, x) for x in fieldnames(typeof(agent)))
    push!(Agents.agent_container(model), tuple_agent)
    return
end

Agents.agenttype(::StructArrayABM{S,A}) where {S,A} = A

function CommonSolve.step!(model::StructArrayABM, n::Union{Real, Function} = 1)
    agent_step! = Agents.agent_step_field(model)
    model_step! = Agents.model_step_field(model)
    t = getfield(model, :time)
    step_ahead!(model, agent_step!, model_step!, n, t)
    return model
end

function step_ahead!(model::StructArrayABM, agent_step!, model_step!, n, t)
    agents_first = getfield(model, :agents_first)
    t0 = t[]
    while Agents.until(t[], t0, n, model)
        !agents_first && model_step!(model)
        for id in schedule(model)
            # ensure we don't act on agent that doesn't exist
            Agents.agent_not_removed(id, model) || continue
            agent_step!(model[id], model)
        end
        agents_first && model_step!(model)
        t[] += 1
    end
end
function step_ahead!(model::StructArrayABM, agent_step!::typeof(dummystep), model_step!, n, t)
    t0 = t[]
    while Agents.until(t[], t0, n, model)
        model_step!(model)
        t[] += 1
    end
    return model
end

and we can see that it works better than our current containers with this simple example:

function taxsystem(ModelType, model_step!; nagents = 100000)
    if ModelType == StandardABM
        model = ModelType(Payer; model_step!, container=Vector, rng = Xoshiro(42))
    else
        model = ModelType(Payer; model_step!, rng = Xoshiro(42))
    end
    for _ in 1:nagents
        add_agent!(model, true, rand(abmrng(model), 20:30), rand(abmrng(model), 500:5000), 0)
    end
    return model
end

function model_step_1!(model)
    for agent in allagents(model)
        agent.age += 1
        agent.wealth += agent.wage - (agent.wage < 1000) * agent.wage * 0.2
    end
end
function model_step_2!(model)
    for id in allids(model)
        agent = model[id]
        agent.age += 1
        agent.wealth += agent.wage - (agent.wage < 1000) * agent.wage * 0.2
    end
end
function model_step_3!(model)
    c = Agents.agent_container(model)
    @~ c.age .+= 1
    @~ c.wealth .+= c.wage .- (c.wage .< 1000) .* c.wage .* 0.2
end

# compile
model1 = taxsystem(StandardABM, model_step_1!);
model2 = taxsystem(StructArrayABM, model_step_2!);
model3 = taxsystem(StructArrayABM, model_step_3!);
@time step!(model1, 100);
@time step!(model2, 100);
@time step!(model3, 100);

# benchmark
model1 = taxsystem(StandardABM, model_step_1!);
model2 = taxsystem(StructArrayABM, model_step_2!);
model3 = taxsystem(StructArrayABM, model_step_3!);
@time step!(model1, 100);
@time step!(model2, 100);
@time step!(model3, 100);

since

julia> @time step!(model1, 100);
  0.010449 seconds

julia> @time step!(model2, 100);
  0.007805 seconds

julia> @time step!(model3, 100);
  0.006923 seconds

but it is clear enough we can support all of this just by specifying container=StructVector in the model construction

[Datseris]

Well that looks great and we are pretty much ready for the PR? One thing I'd like to change is doing

    c = Agents.agent_container(model)
    @~ c.age .+= 1

I'd prefer if the agent_container remains an internal function. Could we define a couple of special methods specially so that there is broadcasting across agents for the SoA model?

One way would be to re-use the allagents function, which would return as special iterator of AgentWrapperSoA, and create a macro that does the same broadcasting as @~:

c = allagents(model)
@aa c.age += 1

with @aa meaning "across agents". We can then document that the whole purpose of SoAABM is to provide this broadcasted overloading of allagents. Thoughts?

[Tortar]

Yes, indeed the vectorized case (as well as other things) are just sketches. My idea is that to allow more flexibility as possible we shouldn't have a special macro for the @~, instead we should allow the user to use that macro or other things (e.g ThreadX.Map) for the vectorization. The way to go is my opinion to specialize the getproperty of the allagents(model) special iterator to return the underlying vectors of the property, so that everything will work without more interventions