diff --git a/Project.toml b/Project.toml index 60d935b765..3ece79b168 100644 --- a/Project.toml +++ b/Project.toml @@ -12,6 +12,7 @@ Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6" Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6" +IteratorSampling = "ef79a3d2-ae9f-5cd2-ab61-e13847810a6e" JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" LazyArtifacts = "4af54fe1-eca0-43a8-85a7-787d91b784e3" LightOSM = "d1922b25-af4e-4ba3-84af-fe9bea896051" @@ -46,6 +47,7 @@ DataStructures = "0.18" Distributions = "0.25" Downloads = "1" Graphs = "1.4" +IteratorSampling = "0.2" JLD2 = "0.4" LazyArtifacts = "1.3.0" LightOSM = "0.2.0" diff --git a/src/Agents.jl b/src/Agents.jl index 189b5c5f4c..e181243bbd 100644 --- a/src/Agents.jl +++ b/src/Agents.jl @@ -7,15 +7,16 @@ module Agents read(path, String) end Agents -using Distributed +using DataFrames using DataStructures +using Distributed using Graphs -using DataFrames +using IteratorSampling using MacroTools +import ProgressMeter using Random using StaticArrays: SVector export SVector -import ProgressMeter import LinearAlgebra # Core structures of Agents.jl diff --git a/src/core/model_free_extensions.jl b/src/core/model_free_extensions.jl index 11cef3fc2a..0f922601e0 100644 --- a/src/core/model_free_extensions.jl +++ b/src/core/model_free_extensions.jl @@ -82,11 +82,15 @@ end function fallback_random_agent(model, condition, alloc) if alloc iter_ids = allids(model) - return sampling_with_condition_agents_single(iter_ids, condition, model) + id = sampling_with_condition_single(iter_ids, condition, model, id -> model[id]) + isnothing(id) && return nothing + return model[id] else iter_agents = allagents(model) iter_filtered = Iterators.filter(agent -> condition(agent), iter_agents) - return resorvoir_sampling_single(iter_filtered, model) + agent = IteratorSampling.itsample(abmrng(model), iter_filtered) + isnothing(agent) && return nothing + return agent end end diff --git a/src/core/space_interaction_API.jl b/src/core/space_interaction_API.jl index cea0a325b6..d1e2044a10 100644 --- a/src/core/space_interaction_API.jl +++ b/src/core/space_interaction_API.jl @@ -304,8 +304,8 @@ nearby_agents(a, model, r = 1; kwargs...) = """ random_nearby_id(agent, model::ABM, r = 1, f = nothing, alloc = false; kwargs...) → id -Return the `id` of a random agent near the position of the given `agent` using an optimized -algorithm from [Reservoir sampling](https://en.wikipedia.org/wiki/Reservoir_sampling#An_optimal_algorithm). +Return the `id` of a random agent near the position of the given `agent`. + Return `nothing` if no agents are nearby. The value of the argument `r` and possible keywords operate identically to [`nearby_ids`](@ref). @@ -317,18 +317,21 @@ is expensive since in this case the allocating version can be more performant. For discrete spaces, use [`random_id_in_position`](@ref) instead to return a random id at a given position. + +This function, as all the other methods which sample from lazy iterators, uses an optimized +algorithm which doesn't require to collect all elements to just sample one of them. """ function random_nearby_id(a, model, r = 1, f = nothing, alloc = false; kwargs...) iter = nearby_ids(a, model, r; kwargs...) if isnothing(f) - return resorvoir_sampling_single(iter, model) + return IteratorSampling.itsample(abmrng(model), iter) else if alloc return sampling_with_condition_single(iter, f, model) else iter_filtered = Iterators.filter(id -> f(id), iter) - return resorvoir_sampling_single(iter_filtered, model) - end + return IteratorSampling.itsample(abmrng(model), iter_filtered) + end end end @@ -348,21 +351,19 @@ For discrete spaces, use [`random_agent_in_position`](@ref) instead to return a position. """ function random_nearby_agent(a, model, r = 1, f = nothing, alloc = false; kwargs...) + iter_ids = nearby_ids(a, model, r; kwargs...) if isnothing(f) - id = random_nearby_id(a, model, r; kwargs...) - isnothing(id) && return nothing - return model[id] + id = IteratorSampling.itsample(abmrng(model), iter_ids) else - iter_ids = nearby_ids(a, model, r; kwargs...) if alloc - return sampling_with_condition_agents_single(iter_ids, f, model) + id = sampling_with_condition_single(iter_ids, f, model, id -> model[id]) else iter_filtered = Iterators.filter(id -> f(model[id]), iter_ids) - id = resorvoir_sampling_single(iter_filtered, model) - isnothing(id) && return nothing - return model[id] + id = IteratorSampling.itsample(abmrng(model), iter_filtered) end end + isnothing(id) && return nothing + return model[id] end """ @@ -380,69 +381,13 @@ is expensive since in this case the allocating version can be more performant. function random_nearby_position(pos, model, r=1, f = nothing, alloc = false; kwargs...) iter = nearby_positions(pos, model, r; kwargs...) if isnothing(f) - return resorvoir_sampling_single(iter, model) + return IteratorSampling.itsample(abmrng(model), iter) else if alloc return sampling_with_condition_single(iter, f, model) else iter_filtered = Iterators.filter(pos -> f(pos), iter) - return resorvoir_sampling_single(iter_filtered, model) + return IteratorSampling.itsample(abmrng(model), iter_filtered) end end -end - -####################################################################################### -# %% sampling functions -####################################################################################### - -function sampling_with_condition_single(iter, condition, model) - population = collect(iter) - n = length(population) - rng = abmrng(model) - @inbounds while n != 0 - index_id = rand(rng, 1:n) - el = population[index_id] - condition(el) && return el - population[index_id], population[n] = population[n], population[index_id] - n -= 1 - end - return nothing -end - -# almost a copy of sampling_with_condition_single, but it's better to call this one -# when selecting an agent since collecting ids is less costly than collecting agents -function sampling_with_condition_agents_single(iter, condition, model) - population = collect(iter) - n = length(population) - rng = abmrng(model) - @inbounds while n != 0 - index_id = rand(rng, 1:n) - el = population[index_id] - condition(model[el]) && return model[el] - population[index_id], population[n] = population[n], population[index_id] - n -= 1 - end - return nothing -end - -# Reservoir sampling function (https://en.wikipedia.org/wiki/Reservoir_sampling) -function resorvoir_sampling_single(iter, model) - res = iterate(iter) - isnothing(res) && return nothing - rng = abmrng(model) - el, state = res - w = rand(rng) - while true - skip_counter = ceil(Int, randexp(rng)/log(1-w)) - while skip_counter != 0 - skip_res = iterate(iter, state) - isnothing(skip_res) && return el - state = skip_res[2] - skip_counter += 1 - end - res = iterate(iter, state) - isnothing(res) && return el - el, state = res - w *= rand(rng) - end -end +end \ No newline at end of file diff --git a/src/simulations/sample.jl b/src/simulations/sample.jl index e28537eff4..7aa6f2eb75 100644 --- a/src/simulations/sample.jl +++ b/src/simulations/sample.jl @@ -121,3 +121,21 @@ end function choose_arg(x, kwargs_nt, agent) return deepcopy(getfield(hasproperty(kwargs_nt, x) ? kwargs_nt : agent, x)) end + +####################################################################################### +# %% sampling functions +####################################################################################### + +function sampling_with_condition_single(iter, condition, model, transform=identity) + population = collect(iter) + n = length(population) + rng = abmrng(model) + @inbounds while n != 0 + index_id = rand(rng, 1:n) + el = population[index_id] + condition(transform(el)) && return el + population[index_id], population[n] = population[n], population[index_id] + n -= 1 + end + return nothing +end diff --git a/src/spaces/discrete.jl b/src/spaces/discrete.jl index 5a9181faa7..fb5f786e60 100644 --- a/src/spaces/discrete.jl +++ b/src/spaces/discrete.jl @@ -109,7 +109,7 @@ function random_empty(model::ABM{<:DiscreteSpace}, cutoff = 0.998) end else empty = empty_positions(model) - return resorvoir_sampling_single(empty, model) + return IteratorSampling.itsample(abmrng(model), empty) end end @@ -145,13 +145,15 @@ function random_id_in_position(pos, model) isempty(ids) && return nothing return rand(abmrng(model), ids) end -function random_id_in_position(pos, model, f, alloc = false) +function random_id_in_position(pos, model, f, alloc = false, transform = identity) iter_ids = ids_in_position(pos, model) if alloc - return sampling_with_condition_single(iter_ids, f, model) + return sampling_with_condition_single(iter_ids, f, model, transform) else - iter_filtered = Iterators.filter(id -> f(id), iter_ids) - return resorvoir_sampling_single(iter_filtered, model) + iter_filtered = Iterators.filter(id -> f(transform(id)), iter_ids) + id = IteratorSampling.itsample(abmrng(model), iter_filtered) + isnothing(id) && return nothing + return id end end @@ -172,15 +174,9 @@ function random_agent_in_position(pos, model) return model[id] end function random_agent_in_position(pos, model, f, alloc = false) - iter_ids = ids_in_position(pos, model) - if alloc - return sampling_with_condition_agents_single(iter_ids, f, model) - else - iter_filtered = Iterators.filter(id -> f(model[id]), iter_ids) - id = resorvoir_sampling_single(iter_filtered, model) - isnothing(id) && return nothing - return model[id] - end + id = random_id_in_position(pos, model, f, alloc, id -> model[id]) + isnothing(id) && return nothing + return model[id] end ####################################################################################### diff --git a/src/spaces/walk.jl b/src/spaces/walk.jl index fd67224f28..2616000d9a 100644 --- a/src/spaces/walk.jl +++ b/src/spaces/walk.jl @@ -205,8 +205,8 @@ function random_empty_pos_in_offsets(offsets, agent, model) n_attempts -= 1 end targets = Iterators.map(β -> normalize_position(agent.pos .+ β, model), offsets) - check_empty = pos -> isempty(pos, model) - return sampling_with_condition_single(targets, check_empty, model) + empty_targets = Iterators.filter(pos -> isempty(pos, model), targets) + return IteratorSampling.itsample(abmrng(model), empty_targets) end """