Neighbor Search

Neighbor Search

neojax.nn.neighbor_search.Neighbors

Padded fixed-width neighbor lists.

Stores for each of the m query points the indices of up to k neighboring data points. All shapes are static, for compatibility with jax.

Entries with mask == False are padding and must be ignored. Their index is a valid (but meaningless) index into the data points and their sq_dist is zero.

Internal Attributes

• index (Int[Array, "m k"]): Indices of the neighbors in the data points.
• mask (Bool[Array, "m k"]): True for entries that are actual neighbors, False for padding.
• sq_dist (Float[Array, "m k"]): Squared distances between the query points and their neighbors. Zero for padding.

---

neojax.nn.neighbor_search.radius_neighbor_search(data: Float[Array, 'n d'], queries: Float[Array, 'm d'], radius: float, max_neighbors: int, chunk_size: int | None = None) -> Neighbors

Finds the neighbors of each query point within a fixed radius.

For each point x in queries, returns the indices of the points y in data with \(\|x - y\| \leq r\), i.e. points on the boundary of the ball are included. A point of data that coincides with the query is a neighbor.

The result has a fixed width of max_neighbors. If more than max_neighbors points are inside the ball, the max_neighbors nearest are kept. If less are inside, the remaining entries are masked out.

Parameters:

Name	Type	Description	Default
data	Float[Array, 'n d']	Search space of possible neighbors, shaped (n, d).	required
queries	Float[Array, 'm d']	Points to find the neighbors of, shaped (m, d).	required
radius	float	Radius of the balls around the query points.	required
max_neighbors	int	Maximum number of neighbors per query point. Clipped to n if larger than the number of data points.	required
chunk_size	int | None	Number of query points to process at once. The search computes (chunk_size, n) pairwise distances at a time, so a smaller chunk size lowers peak memory at the cost of speed. Default is None, all query points at once.	None

Returns:

Type	Description
Neighbors	The neighbor lists, with index, mask and sq_dist shaped (m, k) where k = min(max_neighbors, n).

Raises:

Type	Description
ValueError	If radius is not positive.
ValueError	If max_neighbors or chunk_size is not a positive integer.

Complexity

The search is exact but brute force, \(\mathcal{O}(m \cdot n)\) in compute and \(\mathcal{O}(c \cdot n)\) in memory for chunk size \(c\). This is typically fine for the problem sizes considered.

Choosing max_neighbors

Neighbors beyond max_neighbors are dropped silently. Check neighbors.mask.all(axis=-1) for a representative sample: if it is True for some query, that query's ball may be truncated.