Skip to content

Improve performance of nearby_positions for GridSpaces #1160

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
May 10, 2025
Merged

Improve performance of nearby_positions for GridSpaces #1160

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
660ede0
Improve performance of nearby_positions for GridSpaces
Tortar May 9, 2025
fbc6de1
Update grid_general.jl
Tortar May 9, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
79 changes: 61 additions & 18 deletions src/spaces/grid_general.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -182,34 +182,31 @@ function nearby_positions(pos::GridPos{D}, model::ABM{<:AbstractGridSpace{D}}, a
return nearby_positions(pos, abmspace(model), args...)
end
function nearby_positions(
pos::GridPos{D}, space::AbstractGridSpace{D,false}, r=1,
pos::GridPos{D}, space::AbstractGridSpace{D,true}, r=1,
get_indices_f=offsets_within_radius_no_0 # NOT PUBLIC API! For `ContinuousSpace`.
) where {D}
nindices = get_indices_f(space, r)
space_size = spacesize(space)
# 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
stored_ids = space.stored_ids
return (n .+ pos for n in nindices if checkbounds(Bool, stored_ids, (n .+ pos)...))
end
stored_ids = space.stored_ids
do_bounds = !all(i -> r < pos[i] <= space_size[i] - r, 1:D)
return PeriodicNearbyPositions{D, eltype(nindices), typeof(stored_ids)}(
pos, nindices, stored_ids, space_size, do_bounds
)
end
function nearby_positions(
pos::GridPos{D}, space::AbstractGridSpace{D,true}, r=1,
pos::GridPos{D}, space::AbstractGridSpace{D,false}, r=1,
get_indices_f=offsets_within_radius_no_0 # NOT PUBLIC API! For `ContinuousSpace`.
) where {D}
nindices = get_indices_f(space, r)
space_size = spacesize(space)
# 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
stored_ids = space.stored_ids
return (checkbounds(Bool, stored_ids, (n .+ pos)...) ?
n .+ pos : mod1.(n .+ pos, space_size) for n in nindices)
end
stored_ids = space.stored_ids
do_bounds = !all(i -> r < pos[i] <= space_size[i] - r, 1:D)
return NonPeriodicNearbyPositions{D, eltype(nindices), typeof(stored_ids)}(
pos, nindices, stored_ids, space_size, do_bounds
)
end

# TODO: create custom iterator also for this case
function nearby_positions(
pos::GridPos{D}, space::AbstractGridSpace{D,P}, r=1,
get_indices_f=offsets_within_radius_no_0 # NOT PUBLIC API! For `ContinuousSpace`.
Expand All @@ -230,7 +227,53 @@ function nearby_positions(
end
end

function random_nearby_position(pos::Any, model::ABM{<:AbstractGridSpace{D,false}}, r=1; kwargs...) where {D}
struct PeriodicNearbyPositions{D,Q,I}
pos::GridPos{D}
nindices::Vector{Q}
stored_ids::I
space_size::NTuple{D, Int}
do_bounds::Bool
end
@inline function Base.iterate(iter::PeriodicNearbyPositions, state=1)
state > length(iter.nindices) && return nothing
@inbounds offset = iter.nindices[state]
nearby_pos = iter.pos .+ offset
# check if we are far from the wall to skip bounds checks
if iter.do_bounds && !checkbounds(Bool, iter.stored_ids, nearby_pos...)
nearby_pos = mod1.(nearby_pos, iter.space_size)
end
return (nearby_pos, state + 1)
end
Base.IteratorSize(::Type{<:PeriodicNearbyPositions}) = Base.HasLength()
Base.length(iter::PeriodicNearbyPositions) = length(iter.nindices)
Base.eltype(::Type{PeriodicNearbyPositions{D}}) where {D} = GridPos{D}

struct NonPeriodicNearbyPositions{D,Q,I}
pos::GridPos{D}
nindices::Vector{Q}
stored_ids::I
space_size::NTuple{D, Int}
do_bounds::Bool
end
@inline function Base.iterate(iter::NonPeriodicNearbyPositions, state=1)
state > length(iter.nindices) && return nothing
@inbounds offset = iter.nindices[state]
nearby_pos = iter.pos .+ offset
# check if we are far from the wall to skip bounds checks
if iter.do_bounds
while !checkbounds(Bool, iter.stored_ids, nearby_pos...)
state == length(iter.nindices) && return nothing
state += 1
@inbounds offset = iter.nindices[state]
nearby_pos = iter.pos .+ offset
end
end
return (nearby_pos, state + 1)
end
Base.IteratorSize(::Type{<:NonPeriodicNearbyPositions}) = Base.SizeUnknown()
Base.eltype(::Type{NonPeriodicNearbyPositions{D}}) where {D} = GridPos{D}

function random_nearby_position(pos::GridPos{D}, model::ABM{<:AbstractGridSpace{D,false}}, r=1; kwargs...) where {D}
nindices = offsets_within_radius_no_0(abmspace(model), r)
stored_ids = abmspace(model).stored_ids
rng = abmrng(model)
Expand Down
Loading