Improve performance of EventQueueABM by better union splitting (#1065)

* Improve performance of EventQueueABM by better union splitting * Update Agents.jl * Update Project.toml * Update model_event_queue.jl
JuliaDynamics · Aug 11, 2024 · 50f5dcc · 50f5dcc
1 parent c5db69c
commit 50f5dcc
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 9 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "Agents"
 uuid = "46ada45e-f475-11e8-01d0-f70cc89e6671"
 authors = ["George Datseris", "Tim DuBois", "Aayush Sabharwal", "Ali Vahdati", "Adriano Meligrana"]
-version = "6.1.1"
+version = "6.1.2"
 
 [deps]
 CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"

diff --git a/src/Agents.jl b/src/Agents.jl
@@ -17,7 +17,7 @@ using DynamicSumTypes
 export variant, variantof
 import ProgressMeter
 using Random
-using StaticArrays: SVector
+using StaticArrays: SVector, SizedVector
 export SVector
 import LinearAlgebra
 import StreamSampling

diff --git a/src/core/model_event_queue.jl b/src/core/model_event_queue.jl
@@ -141,15 +141,15 @@ can be evolved in time.
   event for an agent after an event affected said agent has been triggered.
 """
 function EventQueueABM(
-        ::Type{A}, events::E,
+        ::Type{A}, events,
         space::S = nothing;
         container::Type = Dict,
         properties::P = nothing,
         rng::R = Random.default_rng(),
         warn = true,
         autogenerate_on_add = true,
         autogenerate_after_action = true,
-    ) where {A<:AbstractAgent,S<:SpaceType,E,P,R<:AbstractRNG}
+    ) where {A<:AbstractAgent,S<:SpaceType,P,R<:AbstractRNG}
     if warn
         @warn "This model type is still experimental which means that it is subject to breaking changes in the
         future. Also, while all the core functionalities have been implemented, this model type
@@ -158,9 +158,10 @@ function EventQueueABM(
     end
     C = construct_agent_container(container, A)
     agents = C()
-    I = events isa Tuple ? Int : keytype(events) # `Tuple` doesn't define `keytype`...
+    events = SizedVector{length(events), Union{typeof.(events)...}}(events...)
+
     # the queue stores pairs of (agent ID, event index) mapping them to their trigger time
-    queue = BinaryHeap(Base.By(last), Pair{Tuple{I, Int}, Float64}[])
+    queue = BinaryHeap(Base.By(last), Pair{Tuple{Int, Int}, Float64}[])
 
     agent_types = is_sumtype(A) ? values(allvariants(A)) : union_types(A)
     type_func = is_sumtype(A) ? variantof : typeof
@@ -195,8 +196,8 @@ function EventQueueABM(
     # because we use the index of `type_to_idx` to access them.
 
     # construct the type
-    TF,ET,PT,FPT,TI,Q = typeof.((
-        type_func, idx_events_each_type, propensities_each_type,
+    E,TF,ET,PT,FPT,TI,Q = typeof.((
+        events, type_func, idx_events_each_type, propensities_each_type,
         idx_func_propensities_each_type, type_to_idx, queue
     ))
     return EventQueueABM{S,A,C,P,E,R,TF,ET,PT,FPT,TI,Q}(
@@ -249,7 +250,7 @@ function add_event!(agent, model) # TODO: Study type stability of this function
     end
     # Then, select an event based on propensities
     event_idx = events_type[sample_propensity(abmrng(model), propensities_type)]    # The time to the event is generated from the selected event
-    selected_event = abmevents(model)[event_idx]
+    selected_event = events[event_idx]
     selected_prop = propensities_type[event_idx]
     t = selected_event.timing(agent, model, selected_prop)
     # we then propagate to the direct function