Improve speed of Nearby Look-ups in GridSpaces (#816)

---------

Co-authored-by: George Datseris <[email protected]>

CHANGELOG.md

@@ -1,5 +1,6 @@
 # main
 
+- Nearby look-ups with `nearby_positions`, `nearby_ids` and derivatives are now incrementally faster than before more the radius increases.
 - The `randomwalk!` function is now supported for any number of dimensions in ContinuousSpace when used to create isotropic/uniform random walks. For all type of `AbstractGridSpace`, the `randomwalk!` function supports a new keyword `force_motion`, which is false by default. See the docs to be informed on the effect of setting this keyword. Besides, in the continuous space default case random walks are up to 2 times faster than before.
 - The `ByProperty` scheduler can now accept any type of (ordered) properties, while before it was restricted to only floats. The `ByID` scheduler of an `UnremovableABM` is now as fast as the `Fastest` scheduler since in this case they are actually equivalent.
 

src/spaces/grid_general.jl

@@ -134,18 +134,28 @@ function nearby_positions(
     ) where {D}
     stored_ids = space.stored_ids
     nindices = get_indices_f(space, r)
-    positions_iterator = (n .+ pos for n in nindices)
-    return Base.Iterators.filter(
-        pos -> checkbounds(Bool, stored_ids, pos...), positions_iterator
-    )
+    space_size = size(stored_ids)
+    # check if we are far from the wall to skip bounds checks
+    if all(i -> r < pos[i] <= space_size[i] - r, 1:D)
+        return (n .+ pos for n in nindices)
+    else
+        return (n .+ pos for n in nindices if checkbounds(Bool, stored_ids, (n .+ pos)...))
+    end
 end
 function nearby_positions(
         pos::ValidPos, space::AbstractGridSpace{D,true}, r = 1,
         get_indices_f = offsets_within_radius_no_0 # NOT PUBLIC API! For `ContinuousSpace`.
     ) where {D}
+    stored_ids = space.stored_ids
     nindices = get_indices_f(space, r)
     space_size = size(space)
-    return (mod1.(n .+ pos, space_size) for n in nindices)
+    # check if we are far from the wall to skip bounds checks
+    if all(i -> r < pos[i] <= space_size[i] - r, 1:D)
+        return (n .+ pos for n in nindices)
+    else
+        return (checkbounds(Bool, stored_ids, (n .+ pos)...) ? 
+                n .+ pos : mod1.(n .+ pos, space_size) for n in nindices)
+    end
 end
 
 function random_nearby_position(pos::ValidPos, model::ABM{<:AbstractGridSpace{D,false}}, r=1; kwargs...) where {D}

src/spaces/grid_multi.jl

@@ -109,7 +109,9 @@ nearby_ids(pos::NTuple, model::ABM{<:GridSpace}, r::Real = 1) = nearby_ids(pos,
 function nearby_ids(pos::NTuple{D, Int}, space::GridSpace{D,P}, r::Real = 1) where {D,P}
     nindices = offsets_within_radius(space, r)
     stored_ids = space.stored_ids
-    return GridSpaceIdIterator{P}(stored_ids, nindices, pos)
+    space_size = size(stored_ids)
+    nocheck = all(i -> r < pos[i] <= space_size[i] - r, 1:D)
+    return GridSpaceIdIterator{P}(stored_ids, nindices, pos, space_size, nocheck)
 end
 
 # Iterator struct. State is `(pos_i, inner_i)` with `pos_i` the index to the nearby indices
@@ -120,40 +122,43 @@ struct GridSpaceIdIterator{P,D}
     origin::NTuple{D,Int}             # origin position nearby is measured from
     L::Int                            # length of `indices`
     space_size::NTuple{D,Int}         # size of `stored_ids`
+    nocheck::Bool                     # skip bound checks if we are far from the edges
 end
-function GridSpaceIdIterator{P}(stored_ids, indices, origin::NTuple{D,Int}) where {P,D}
+function GridSpaceIdIterator{P}(stored_ids, indices, origin::NTuple{D,Int}, space_size, nocheck) where {P,D}
     L = length(indices)
     @assert L > 0
-    space_size = size(stored_ids)
-    return GridSpaceIdIterator{P,D}(stored_ids, indices, origin, L, space_size)
+    return GridSpaceIdIterator{P,D}(stored_ids, indices, origin, L, space_size, nocheck)
 end
 Base.eltype(::Type{<:GridSpaceIdIterator}) = Int # It returns IDs
 Base.IteratorSize(::Type{<:GridSpaceIdIterator}) = Base.SizeUnknown()
 
 # Instructs how to combine two positions. Just to avoid code duplication for periodic
-combine_positions(pos, origin, ::GridSpaceIdIterator{false}) = pos .+ origin
-function combine_positions(pos, origin, iter::GridSpaceIdIterator{true})
+combine_positions(pos, origin, ::GridSpaceIdIterator{false}, ::Bool) = pos .+ origin
+function combine_positions(pos, origin, iter::GridSpaceIdIterator{true}, nocheck)
     # the mod function is not needed for many positions and it's expensive compared
     # with checking for bounds so it is better to apply it only as a fallback.
-    off_pos = pos .+ origin
-    checkbounds(Bool, iter.stored_ids, off_pos...) && return off_pos
-    return mod1.(off_pos, iter.space_size)
+    npos = pos .+ origin
+    if nocheck || checkbounds(Bool, iter.stored_ids, npos...)
+        return npos
+    else
+        return mod1.(npos, iter.space_size)
+    end
 end
 
 # Initialize iteration
 function Base.iterate(iter::GridSpaceIdIterator)
     @inbounds begin
-        stored_ids, indices, L, origin = getproperty.(
-        Ref(iter), (:stored_ids, :indices, :L, :origin))
+    stored_ids, indices, L, origin, nocheck = getproperty.(
+        Ref(iter), (:stored_ids, :indices, :L, :origin, :nocheck))
     pos_i = 1
-    pos_index = combine_positions(indices[pos_i], origin, iter)
+    pos_index = combine_positions(indices[pos_i], origin, iter, nocheck)
     # First, check if the position index is valid (bounds checking)
     # AND whether the position is empty. If not, proceed to next position index.
-    while invalid_access(pos_index, iter)
+    while invalid_access(pos_index, iter, nocheck)
         pos_i += 1
         # Stop iteration if `pos_index` exceeded the amount of positions
         pos_i > L && return nothing
-        pos_index = combine_positions(indices[pos_i], origin, iter)
+        pos_index = combine_positions(indices[pos_i], origin, iter, nocheck)
     end
     # We have a valid position index and a non-empty position
     ids_in_pos = stored_ids[pos_index...]
@@ -163,11 +168,15 @@ function Base.iterate(iter::GridSpaceIdIterator)
 end
 
 # Must return `true` if the access is invalid
-function invalid_access(pos_index, iter::GridSpaceIdIterator{false})
-    valid_bounds = checkbounds(Bool, iter.stored_ids, pos_index...)
-    return !valid_bounds || @inbounds isempty(iter.stored_ids[pos_index...])
+function invalid_access(pos_index, iter::GridSpaceIdIterator{false}, nocheck)
+    if nocheck
+        return @inbounds isempty(iter.stored_ids[pos_index...])
+    else
+        valid_bounds = checkbounds(Bool, iter.stored_ids, pos_index...)
+        return !valid_bounds || @inbounds isempty(iter.stored_ids[pos_index...])
+    end
 end
-function invalid_access(pos_index, iter::GridSpaceIdIterator{true})
+function invalid_access(pos_index, iter::GridSpaceIdIterator{true}, ::Bool)
     @inbounds isempty(iter.stored_ids[pos_index...])
 end
 
@@ -176,8 +185,8 @@ end
 # known knowledge of `pos_i` being a valid position index.
 function Base.iterate(iter::GridSpaceIdIterator, state)
     @inbounds begin
-    stored_ids, indices, L, origin = getproperty.(
-        Ref(iter), (:stored_ids, :indices, :L, :origin))
+    stored_ids, indices, L, origin, nocheck = getproperty.(
+        Ref(iter), (:stored_ids, :indices, :L, :origin, :nocheck))
     pos_i, inner_i, ids_in_pos = state
     X = length(ids_in_pos)
     if inner_i > X
@@ -186,12 +195,12 @@ function Base.iterate(iter::GridSpaceIdIterator, state)
         # Stop iteration if `pos_index` exceeded the amount of positions
         pos_i > L && return nothing
         inner_i = 1
-        pos_index = combine_positions(indices[pos_i], origin, iter)
+        pos_index = combine_positions(indices[pos_i], origin, iter, nocheck)
         # Of course, we need to check if we have valid index
-        while invalid_access(pos_index, iter)
+        while invalid_access(pos_index, iter, nocheck)
             pos_i += 1
             pos_i > L && return nothing
-            pos_index = combine_positions(indices[pos_i], origin, iter)
+            pos_index = combine_positions(indices[pos_i], origin, iter, nocheck)
         end
         ids_in_pos = stored_ids[pos_index...]
     end

src/spaces/grid_single.jl

@@ -86,25 +86,35 @@ function nearby_ids(pos::NTuple{D, Int}, model::ABM{<:GridSpaceSingle{D,true}},
     nindices = get_offset_indices(model, r)
     stored_ids = model.space.stored_ids
     space_size = size(stored_ids)
-    array_accesses_iterator = (stored_ids[(mod1.(pos .+ β, space_size))...] for β in nindices)
-    # Notice that not all positions are valid; some are empty! Need to filter:
-    valid_pos_iterator = Base.Iterators.filter(x -> x ≠ 0, array_accesses_iterator)
-    return valid_pos_iterator
+    position_iterator = (pos .+ β for β in nindices)
+    # check if we are far from the wall to skip bounds checks
+    if all(i -> r < pos[i] <= space_size[i] - r, 1:D)
+        ids_iterator = (stored_ids[p...] for p in position_iterator 
+                        if stored_ids[p...] != 0)
+    else
+        ids_iterator = (checkbounds(Bool, stored_ids, p...) ? 
+                        stored_ids[p...] : stored_ids[mod1.(p, space_size)...]
+                        for p in position_iterator if stored_ids[mod1.(p, space_size)...] != 0)
+    end
+    return ids_iterator
 end
 
 function nearby_ids(pos::NTuple{D, Int}, model::ABM{<:GridSpaceSingle{D,false}}, r = 1,
         get_offset_indices = offsets_within_radius # internal, see last function
     ) where {D}
     nindices = get_offset_indices(model, r)
     stored_ids = model.space.stored_ids
-    positions_iterator = (pos .+ β for β in nindices)
-    # Here we combine in one filtering step both valid accesses to the space array
-    # but also that the accessed location is not empty (i.e., id is not 0)
-    array_accesses_iterator = Base.Iterators.filter(
-        pos -> checkbounds(Bool, stored_ids, pos...) && stored_ids[pos...] ≠ 0,
-        positions_iterator
-    )
-    return (stored_ids[pos...] for pos in array_accesses_iterator)
+    space_size = size(stored_ids)
+    position_iterator = (pos .+ β for β in nindices)
+    # check if we are far from the wall to skip bounds checks
+    if all(i -> r < pos[i] <= space_size[i] - r, 1:D)
+        ids_iterator = (stored_ids[p...] for p in position_iterator 
+                        if stored_ids[p...] != 0)
+    else
+        ids_iterator = (stored_ids[p...] for p in position_iterator 
+                        if checkbounds(Bool, stored_ids, p...) && stored_ids[p...] != 0)
+    end
+    return ids_iterator
 end
 
 # Contrary to `GridSpace`, we also extend here `nearby_ids(a::Agent)`.