Spatial indexing and neighbor contracts

The question this method answers¶

How can a researcher avoid comparing every point with every other point while still knowing whether a reported neighbor set is complete? Spatial indexing first groups points into cells, gathers a bounded candidate neighborhood, and then applies the physical distance criterion.

Before computation: what should be true?¶

Define the coordinate frame, boundary convention, cell size, cutoff, and fixed capacities. For the exact query, cell_size >= cutoff, Bcap must hold every searched cell, and k_max must hold every accepted neighbor per particle. The 27-cell gather makes a candidate axis of length 27 * Bcap, so k_max <= 27 * Bcap is required before the fixed-size top-k selection can run. The current geometry is open and clamped, not periodic.

Morton binning requires Nbins_per_dim to be a positive power of two no larger than 1024. Dense linear cells accept arbitrary positive dimensions and are the owner used by exact fixed-radius gathering.

Define the mathematical objects¶

A cell is a bounded spatial region assigned an integer ID. A neighborhood is a set of nearby cells inspected for possible interactions. Topology describes which cell or point identities are connected; it can change discontinuously when a point crosses a cell face or distance cutoff.

A Morton or Z-order code interleaves integer coordinate bits into one integer Morton (1966). It often preserves locality in memory, but it is not a distance metric. A mask marks which slots in a fixed-capacity array hold real indices; a sentinel occupies unused slots. An overflow flag says that capacity was insufficient and the stored set may be incomplete.

Derive the method¶

If coordinate bits are $x_b,y_b,z_b\in\{0,1\}$, a three-dimensional Morton code interleaves them as

The exact fixed-radius target for focal point $i$ is

With cubic cells at least as wide as the cutoff, any point in $\mathcal{N}_i$ lies in the focal cell or one of its 26 adjacent cells. Gathering that 27-cell stencil, masking invalid slots, and applying the exact distance test therefore reproduces the brute-force set when neither cell nor neighbor capacity overflows. The lower strict inequality excludes self and coincident points; the upper inclusive inequality keeps points exactly on the cutoff.

What the algorithm actually does¶

assign_particles_to_bins maps a symmetric cube to Morton IDs and clamps off-box positions to boundary bins. assign_to_cells_linear produces dense row-major IDs for arbitrary (nx, ny, nz). Clamping keeps indices valid but does not create periodic boundaries.

fill_bins deterministically retains Bcap hash-ranked members per cell. fill_bins_exact stores the same fixed shape and also returns did_overflow. gather_candidates_from_bins, stencil variants, and approx_knn_candidates return bounded candidates that still require a physical filter and recall audit.

gather_pairs_within_radius assigns dense cells, gathers a masked 27-cell stencil, computes distances, and returns (neighbors, mask, did_overflow) with shapes (N, k_max), (N, k_max), and (). It is exact only when did_overflow is false and all stated preconditions, including the candidate-axis bound, hold.

The candidate axis has length 27 * Bcap; therefore k_max <= 27 * Bcap is required independently of whether k_max can hold every true neighbor. If omitted, Bcap=None selects min(N, max(k_max, 64)). Because that fallback cannot exceed $N$, it does not guarantee k_max <= 27 * Bcap when k_max > 27 * N; the caller must still check the structural bound.

The dims=None path reads positions on the host and is eager-only. For a supported transformed call, cell_size, cutoff, k_max, Bcap, and dims must be static or closed over for jax.jit. In particular, cell_size and cutoff pass through Python float(...) in the precondition guard, while the remaining controls determine array shapes, top-k size, or Python grid structure.

The result is exact only when did_overflow is false and those geometry and capacity conditions all hold.

What JAX differentiates¶

Cell assignment uses floor, clipping, integer encoding, sorting, masks, and top-k selection. These are host-side, discrete preprocessing or discrete JAX operations, not a smooth map from positions to neighbor identity. A JIT-compatible spatial query does not thereby have a meaningful topology derivative.

Once a neighbor set is fixed, a downstream smooth distance, force, or density kernel can be differentiated with respect to floating positions or values. That conditional derivative excludes points where cell membership, cutoff inclusion, ranking, or capacity status changes.

Using it in Jaxstro¶

import jax.numpy as jnp

from jaxstro.spatial import gather_pairs_within_radius

positions = jnp.array(
    [
        [0.0, 0.0, 0.0],
        [0.25, 0.0, 0.0],
        [0.5, 0.0, 0.0],
        [1.25, 0.0, 0.0],
        [0.0, 0.5, 0.0],
        [0.0, 0.0, 0.0],
    ]
)

neighbors, mask, did_overflow = gather_pairs_within_radius(
    positions,
    origin=jnp.array([0.0, 0.0, 0.0]),
    cell_size=0.5,
    cutoff=0.5,
    k_max=5,
    Bcap=6,
    dims=(4, 2, 2),
)

focal_neighbors = set(map(int, neighbors[0][mask[0]].tolist()))
assert focal_neighbors == {1, 2, 4}
assert not bool(did_overflow)

How to audit the result¶

1. Compute all pairwise distances for a small cloud and compare every returned row with the brute-force set 0 < |x_i - x_j| <= cutoff.

2. Include coincident points, exact-cutoff points, boundaries, and empty cells.

3. Check neighbor symmetry when the scientific relation should be symmetric.

4. Force cell overflow and neighbor overflow, then verify did_overflow changes.

5. Compare approximate candidate recall with brute-force neighbors on uniform, boundary-heavy, and clustered clouds.

6. Record boundary convention, cell size, dimensions, Bcap, and k_max.

Figure 1 visualizes why candidate recall and exact cutoff filtering are distinct contracts.

Where the claim stops¶

Morton locality does not imply physical distance. Candidate heuristics do not guarantee exact k-nearest-neighbor recall. Fixed-radius exactness is conditional on open clamped geometry, stencil coverage, and both capacities. None of these queries defines periodic wrapping, differentiable topology, or a complete many-body interaction model.

Connected ideas¶