[ANN] LightSumTypes.jl v4

It was cool to implement this package and it works quite well on Julia <=1.10, it really patches Union-splitting there.

However, I’m seeing an incredible performance improvement in Union-splitting in 1.11rc1 for some reason, it seems to almost always work automatically which is really cool, e.g. all the benchmarks related to it here Union splitting vs C++ are okay now, no dynamic dispatch anywhere.

By extensive testing I found some generators having dynamic dispatch, but not more than that. There are some insights from this package which can maybe lead to some more improvements, e.g. actually in some cases the storage required is much less than what Julia currently achieves with multiples types, but Union-splitting works quite well now! So if you can update to Julia 1.11, even if I think this approach has some more guarantees, I would just first try to work with a Union :slight_smile:

The battle is indeed much harsher there:

Code
# The following simple model has a variable number of agent types,
# but there is no removing or creating of additional agents.
# It creates a model that has the same number of agents and does
# overall the same number of operations, but these operations
# are split in a varying number of agents. It shows how much of a
# performance hit is to have many different agent types.

using Agents, DynamicSumTypes, Random, BenchmarkTools

tn = 100
agent_types_s = [Symbol(:Agent, i) for i in 1:tn]


for (i, t) in enumerate(agent_types_s)
    eval(:(@agent struct $t(GridAgent{2})
               money::Any
           end))
end

for i in 1:tn
    eval(:(@sumtype $(Symbol(:AgentAll, i))($(agent_types_s[1:i]...)) <: AbstractAgent))
end

const agent_types = Tuple(eval.(agent_types_s))
const agent_all_t = NamedTuple(Symbol(:AgentAll, i) => eval(Symbol(:AgentAll, i)) for i in 1:tn)

function initialize_model_1(;n_agents=600,dims=(5,5))
    space = GridSpace(dims)
    model = StandardABM(Agent1, space; agent_step!,
                        scheduler=Schedulers.Randomly(),
                        rng = Xoshiro(42), warn=false)
    id = 0
    for id in 1:n_agents
        add_agent!(Agent1, model, 10)
    end
    return model
end

function initialize_model_sum(;n_agents=600, n_types=1, dims=(5,5))
    agents_used = agent_types[1:n_types]
    agent_all = agent_all_t[Symbol(:AgentAll, n_types)]
    space = GridSpace(dims)
    model = StandardABM(agent_all, space; agent_step!,
                        scheduler=Schedulers.Randomly(), warn=false,
                        rng = Xoshiro(42))
    agents_per_type = div(n_agents, n_types)
    for A in agents_used
        add_agents!(A, model, agents_per_type, agent_all)
    end
    return model
end

function initialize_model_n(;n_agents=600, n_types=1, dims=(5,5))
    agents_used = agent_types[1:n_types]
    space = GridSpace(dims)
    model = StandardABM(Union{agents_used...}, space; agent_step!,
                        scheduler=Schedulers.Randomly(), warn=false,
                        rng = Xoshiro(42))
    agents_per_type = div(n_agents, n_types)
    for A in agents_used
        add_agents!(A, model, agents_per_type)
    end
    return model
end

function add_agents!(A, model, n)
    for _ in 1:n
        a = A(model, random_position(model), 10 + rand(abmrng(model), 1:10))
        add_agent!(a, model)
    end
    return nothing
end
function add_agents!(A, model, n, W)
    for _ in 1:n
        a = W(A(model, random_position(model), 10 + rand(abmrng(model), 1:10)))
        add_agent!(a, model)
    end
    return nothing
end

function agent_step!(agent, model)
    move!(agent, model)
    agents = agents_in_position(agent.pos, model)
    for a in agents
        exchange!(agent, a)
    end
    return nothing
end

function move!(agent, model)
    cell = random_nearby_position(agent, model)
    move_agent!(agent, cell, model)
    return nothing
end

function exchange!(agent, other_agent)
    v1 = agent.money
    v2 = other_agent.money
    agent.money = v2
    other_agent.money = v1
    return nothing
end

function run_simulation_1(n_steps)
    model = initialize_model_1()
    Agents.step!(model, n_steps)
end

function run_simulation_sum(n_steps; n_types)
    model = initialize_model_sum(; n_types=n_types)
    Agents.step!(model, n_steps)
end

function run_simulation_n(n_steps; n_types)
    model = initialize_model_n(; n_types=n_types)
    Agents.step!(model, n_steps)
end

# %% Run the simulation, do performance estimate, first with 1, then with many
n_steps = 50
n_types = [2,3,4,5,10,20,30,40,50,60,70,80,90,100]

time_1 = @belapsed run_simulation_1($n_steps)
times_n = Float64[]
times_multi_s = Float64[]
for n in n_types
    println(n)
    t = @belapsed run_simulation_n($n_steps; n_types=$n)
    push!(times_n, t/time_1)
    t_sum = @belapsed run_simulation_sum($n_steps; n_types=$n)
    print(t/time_1, " ", t_sum/time_1)
    push!(times_multi_s, t_sum/time_1)
end

println("relative time of model with 1 type: 1.0")
for (n, t1, t2) in zip(n_types, times_n, times_multi_s)
    println("relative time of model with $n types: $t1")
    println("relative time of model with $n @sumtype: $t2")
end

using CairoMakie
fig, ax = CairoMakie.scatterlines(n_types, times_n; label = "Union");
scatterlines!(ax, n_types, times_multi_s; label = "@sumtype")
ax.xlabel = "# types"
ax.ylabel = "time relative to 1 type"
ax.title = "Union types vs @sumtype"
axislegend(ax; position = :lt)
ax.yticks = 0:1:ceil(Int, maximum(times_n))
ax.xticks = n_types
fig
6 Likes