new internal API for sampling (#875)

* new internal API for sampling

* Update space_interaction_API.jl

* dev

* dev2

* dev3

* Update sample.jl

* Update sample.jl

* Update walk.jl

* Update discrete.jl

* Update space_interaction_API.jl

* Update model_abstract.jl

* Update api.md

* Update model_abstract.jl

* Update sample.jl

* Update space_interaction_API.jl

* Update sample.jl

* update code

* update code 2

* update code 3

* Update Project.toml

* Update Project.toml

* Update Agents.jl

* fix code

* Update Agents.jl

* Update Agents.jl

* printing for debugging

* remove change for AgentsExampleZoo

* Update runtests.jl

* remove prints

* code review changes

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"

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

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
 

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

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

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
 
 #######################################################################################

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
 
 """