Source code for pyLEAFS.spatialhash

"""Generic spatial hash over a point set.

A :class:`SpatialHash` buckets an arbitrary ``(n, D)`` set of points into the
regions of a shared :class:`Grid`, then answers Moore-neighbourhood queries.
It is deliberately agnostic about what the points are --- resources, foragers,
predators --- so the same structure serves every layer. Unlike the resource
field's internal buffer, this is a lightweight index rebuilt from a positions
array on demand.
"""

import numpy as np

from pyLEAFS.grid import Grid


[docs] class SpatialHash: """Bucket a point set into grid regions for neighbourhood queries. Parameters ---------- grid : Grid Shared spatial grid. positions : ndarray, shape (n, D), optional Initial point set. May be rebuilt later with :meth:`build`. Examples -------- >>> g = Grid((5, 5), L=10.0) >>> pts = np.array([[1.0, 1.0], [12.0, 3.0], [1.5, 1.5]]) >>> sh = SpatialHash(g, pts) >>> idx = sh.query(np.array([1.2, 1.2])) >>> sorted(idx.tolist()) [0, 2] """ def __init__(self, grid, positions=None): self.grid = grid if positions is not None: self.build(positions) else: self._positions = np.zeros((0, grid.D)) self._buckets = {}
[docs] def build(self, positions): """(Re)index ``positions`` into per-region buckets. Parameters ---------- positions : ndarray, shape (n, D) Physical positions to index. """ self._positions = np.atleast_2d(np.asarray(positions, dtype=np.float64)) if self._positions.shape[0] == 0: self._buckets = {} return rids = self.grid.region_id(self.grid.region_of(self._positions)) self._buckets = {} for i, rid in enumerate(rids): self._buckets.setdefault(int(rid), []).append(i)
[docs] def query(self, pos): """Indices of points in the Moore neighbourhood of ``pos``. Parameters ---------- pos : ndarray, shape (D,) Physical query position. Returns ------- ndarray of int Indices into the indexed positions array of every point whose region is in the Moore neighbourhood of ``pos``. Not distance filtered --- the caller refines with actual displacements. """ pos = np.asarray(pos, dtype=np.float64).reshape(1, self.grid.D) ridx = self.grid.region_of(pos) nbr_ids = self.grid.neighbour_ids(ridx)[0] out = [] for rid in nbr_ids: b = self._buckets.get(int(rid)) if b: out.extend(b) return np.array(out, dtype=np.int64)
[docs] def query_within(self, pos, radius): """Indices of indexed points within ``radius`` of ``pos``. Refines :meth:`query` by minimum-image distance. Assumes ``radius`` is no larger than one region side so the Moore neighbourhood suffices. """ cand = self.query(pos) if cand.size == 0: return cand disp = self.grid.displacement( np.asarray(pos, dtype=np.float64).reshape(1, self.grid.D), self._positions[cand], ) d2 = np.einsum("ij,ij->i", disp, disp) return cand[d2 <= radius * radius]