neighbors

triwarp.neighbors

Point-set acceleration structures (BVH/HashGrid) and raw neighbor queries.

The accelerator is a keyword, not a function name. There are two questions here -- "everything within r" and "the k nearest" -- and each is one function: query_ball (with query_ball_count and query_ball_with_offsets for the count alone and the flat CSR form) and query_nearest. Which broad phase runs is a QueryBackend keyword -- "hashgrid" or "bvh" -- or is inferred from a prebuilt structure passed as accelerator. Both backends are exact and return the same answer; the choice is a cost one, and query_nearest's docstring carries guidance on which to pick. The kernel side is one warp-uniform kernel branching on an ACCEL_* selector.

The BVH-only queries keep the structure in their names, because naming it is informative rather than redundant there: these index arbitrary bounds rather than points, which a hash grid cannot, and they are broad phase only -- every bound the query region reaches is returned, with no narrow-phase filter on the primitive's own geometry. query_bvh_ball is the ball and query_bvh_box the per-query box; both return the same flat (indices, offsets) pair, so neither carries a _with_offsets suffix -- unlike query_ball_with_offsets, which needs one to separate it from query_ball's dense form. Prefer the ball whenever the predicate is a ball: it is not merely tighter, its traversal is substantially cheaper per candidate than the box one (see query_bvh_ball).

Also home to geodesic_ball, the surface-aware counterpart to the spatial ball queries here: it returns the same CSR (indices, offsets) shape but walks the mesh edge graph, so it excludes vertices that are close in space yet across a fold of the surface.

nearest_neighbor_distance is the one derived quantity rather than a raw query: the per-point distance to the closest other point, which is what a cloud's scale is normally estimated from.

QueryBackend module-attribute

QueryBackend = Literal['hashgrid', 'bvh']

bvh_from_points

bvh_from_points(
    points: array[vec3], leaf_size: int = 4
) -> wp.Bvh

Build a bounding-volume hierarchy over points for radius queries.

Each leaf stores the same geometry as points (degenerate bounds via a clone), matching the broad-phase pattern used by query_ball and query_nearest.

Parameters:

Name	Type	Description	Default
points	array[vec3]	(n, 3) positions as wp.vec3.	required
leaf_size	int	Maximum primitives per leaf; forwarded to warp.Bvh.	4

Returns:

Type	Description
Bvh	BVH suited for the query_ball* and query_nearest family; pass it as their accelerator, which selects backend="bvh" by its type.

See Also

query_ball query_nearest

Source code in triwarp/neighbors.py

def bvh_from_points(points: wp.array[wp.vec3], leaf_size: int = 4) -> wp.Bvh:
    """
    Build a bounding-volume hierarchy over ``points`` for radius queries.

    Each leaf stores the same geometry as ``points`` (degenerate bounds via a clone),
    matching the broad-phase pattern used by
    [`query_ball`][triwarp.neighbors.query_ball] and
    [`query_nearest`][triwarp.neighbors.query_nearest].

    Parameters
    ----------
    points
        ``(n, 3)`` positions as ``wp.vec3``.
    leaf_size
        Maximum primitives per leaf; forwarded to ``warp.Bvh``.

    Returns
    -------
    warp.Bvh
        BVH suited for the ``query_ball*`` and ``query_nearest`` family; pass it as their
        ``accelerator``, which selects ``backend="bvh"`` by its type.

    See Also
    --------
    [`query_ball`][triwarp.neighbors.query_ball]
    [`query_nearest`][triwarp.neighbors.query_nearest]
    """
    return wp.Bvh(points, points, leaf_size=leaf_size)

hashgrid_from_points

hashgrid_from_points(
    points: array[vec3], radius: float, grid_bins: int = 128
) -> wp.HashGrid

Build a 3D hash grid over points for radius queries.

Parameters:

Name	Type	Description	Default
points	array[vec3]	(n, 3) positions as wp.vec3.	required
radius	float	Cell size passed to warp.HashGrid.build and used by the query_ball* and query_nearest kernels under backend="hashgrid".	required
grid_bins	int	Resolution of the hash grid along each axis.	128

Returns:

Type	Description
HashGrid	Hash grid suited for the query_ball* and query_nearest family, and related kernels. The cell width is recorded on the returned object as cell_width, which query_nearest reads back to size its search (warp.HashGrid itself does not keep it).

See Also

query_ball_count query_ball query_nearest

Source code in triwarp/neighbors.py

def hashgrid_from_points(
    points: wp.array[wp.vec3], radius: float, grid_bins: int = 128
) -> wp.HashGrid:
    """
    Build a 3D hash grid over ``points`` for radius queries.

    Parameters
    ----------
    points
        ``(n, 3)`` positions as ``wp.vec3``.
    radius
        Cell size passed to ``warp.HashGrid.build`` and used by
        the ``query_ball*`` and ``query_nearest`` kernels under ``backend="hashgrid"``.
    grid_bins
        Resolution of the hash grid along each axis.

    Returns
    -------
    warp.HashGrid
        Hash grid suited for the ``query_ball*`` and ``query_nearest`` family,
        and related kernels. The cell width is recorded on the returned object as
        ``cell_width``, which
        [`query_nearest`][triwarp.neighbors.query_nearest] reads back to size its
        search (``warp.HashGrid`` itself does not keep it).

    See Also
    --------
    [`query_ball_count`][triwarp.neighbors.query_ball_count]
    [`query_ball`][triwarp.neighbors.query_ball]
    [`query_nearest`][triwarp.neighbors.query_nearest]
    """
    n = int(points.shape[0])
    grid = wp.HashGrid(grid_bins, grid_bins, grid_bins, device=points.device)
    grid.reserve(n)
    grid.build(points, radius)
    grid.cell_width = float(radius)
    return grid

bvh_from_bounds

bvh_from_bounds(
    lower: array[vec3],
    upper: array[vec3],
    leaf_size: int = 4,
) -> wp.Bvh

Build a bounding-volume hierarchy over axis-aligned bounds.

Each primitive i is represented by lower[i] and upper[i] corner positions, suitable for query_bvh_ball and query_bvh_box broad-phase intersection tests.

Parameters:

Name	Type	Description	Default
lower	array[vec3]	(n,) minimum corner of each bound as wp.vec3.	required
upper	array[vec3]	(n,) maximum corner of each bound as wp.vec3.	required
leaf_size	int	Maximum primitives per leaf; forwarded to warp.Bvh.	4

Returns:

Type	Description
Bvh	BVH suited for AABB intersection queries.

Raises:

Type	Description
RuntimeError	If lower and upper are not all on one device.

Source code in triwarp/neighbors.py

def bvh_from_bounds(
    lower: wp.array[wp.vec3], upper: wp.array[wp.vec3], leaf_size: int = 4
) -> wp.Bvh:
    """
    Build a bounding-volume hierarchy over axis-aligned bounds.

    Each primitive ``i`` is represented by ``lower[i]`` and ``upper[i]`` corner
    positions, suitable for
    [`query_bvh_ball`][triwarp.neighbors.query_bvh_ball] and
    [`query_bvh_box`][triwarp.neighbors.query_bvh_box]
    broad-phase intersection tests.

    Parameters
    ----------
    lower
        ``(n,)`` minimum corner of each bound as ``wp.vec3``.
    upper
        ``(n,)`` maximum corner of each bound as ``wp.vec3``.
    leaf_size
        Maximum primitives per leaf; forwarded to ``warp.Bvh``.

    Returns
    -------
    warp.Bvh
        BVH suited for AABB intersection queries.

    Raises
    ------
    RuntimeError
        If ``lower`` and ``upper`` are not all on one device.
    """
    require_same_device(lower=lower, upper=upper)
    return wp.Bvh(lower, upper, leaf_size=leaf_size)

query_bvh_ball

query_bvh_ball(
    bvh: Bvh,
    queries: array[vec3],
    radius: float,
    *,
    include_total: bool = False,
) -> tuple[wp.array[wp.int32], wp.array[wp.int32]]

Low-level BVH ball query: primitive indices in one flat buffer plus offsets.

For each query center q, tests whether every primitive bound in bvh comes within radius of q -- the exact bound-to-point distance, so the returned set is the ball's and not its enclosing cube's. Like query_bvh_box and unlike query_ball_with_offsets, there is no narrow-phase filter on the primitive's own geometry: every bound the ball reaches is returned.

Parameters:

Name	Type	Description	Default
bvh	Bvh	Pre-built BVH from bvh_from_bounds or bvh_from_points.	required
queries	array[vec3]	(m, 3) query centers stored as wp.vec3.	required
radius	float	Ball radius about each query center.	required
include_total	bool	Return offsets in the length-m + 1 CSR form whose trailing element is the total hit count, instead of the length-m form. Same meaning as on the ball queries.	False

Returns:

Type	Description
(candidate_indices_flat, offsets)	Exactly the packing query_bvh_box returns: offsets is the exclusive prefix sum of per-query hit counts, and query k owns candidate_indices_flat[offsets[k] : offsets[k+1]].

Raises:

Type	Description
RuntimeError	If bvh and queries are not all on one device.

Notes

Prefer this over the cube query wherever the caller's predicate is a ball, and not only for the candidates it does not return: the ball traversal is substantially cheaper per candidate than the box one, enough that the cube query loses even when handed the inscribed cube and therefore strictly fewer candidates. The margin grows with the radius, which is where a cube broad phase hurts most.

See Also

query_bvh_box The same packing for an axis-aligned box, given per query. A box is a different predicate and not a looser ball -- it admits corners the ball excludes -- so reach for it when the neighbourhood really is axis-aligned, and for this when it is a ball. query_ball_with_offsets The same ball, over a point cloud and with a narrow-phase distance filter.

Source code in triwarp/neighbors.py

def query_bvh_ball(
    bvh: wp.Bvh, queries: wp.array[wp.vec3], radius: float, *, include_total: bool = False
) -> tuple[wp.array[wp.int32], wp.array[wp.int32]]:
    """
    Low-level BVH ball query: primitive indices in one flat buffer plus offsets.

    For each query center ``q``, tests whether every primitive bound in ``bvh`` comes within
    ``radius`` of ``q`` -- the exact bound-to-point distance, so the returned set is the ball's and
    not its enclosing cube's. Like
    [`query_bvh_box`][triwarp.neighbors.query_bvh_box] and unlike
    [`query_ball_with_offsets`][triwarp.neighbors.query_ball_with_offsets], there is no
    narrow-phase filter on the primitive's own geometry: every bound the ball reaches is returned.

    Parameters
    ----------
    bvh
        Pre-built BVH from [`bvh_from_bounds`][triwarp.neighbors.bvh_from_bounds]
        or [`bvh_from_points`][triwarp.neighbors.bvh_from_points].
    queries
        ``(m, 3)`` query centers stored as ``wp.vec3``.
    radius
        Ball radius about each query center.
    include_total
        Return ``offsets`` in the length-``m + 1`` CSR form whose trailing element is the total hit
        count, instead of the length-``m`` form. Same meaning as on the ball queries.

    Returns
    -------
    candidate_indices_flat, offsets
        Exactly the packing
        [`query_bvh_box`][triwarp.neighbors.query_bvh_box] returns:
        ``offsets`` is the exclusive prefix sum of per-query hit counts, and query ``k`` owns
        ``candidate_indices_flat[offsets[k] : offsets[k+1]]``.

    Raises
    ------
    RuntimeError
        If ``bvh`` and ``queries`` are not all on one device.

    Notes
    -----
    Prefer this over the cube query wherever the caller's predicate is a ball, and not only for the
    candidates it does not return: the ball traversal is *substantially cheaper per candidate* than
    the box one, enough that the cube query loses even when handed the inscribed cube and therefore
    strictly fewer candidates. The margin grows with the radius, which is where a cube broad phase
    hurts most.

    See Also
    --------
    [`query_bvh_box`][triwarp.neighbors.query_bvh_box]
        The same packing for an axis-aligned **box**, given per query. A box is a different
        predicate and not a looser ball -- it admits corners the ball excludes -- so reach for it
        when the neighbourhood really is axis-aligned, and for this when it is a ball.
    [`query_ball_with_offsets`][triwarp.neighbors.query_ball_with_offsets]
        The same ball, over a **point** cloud and with a narrow-phase distance filter.
    """
    require_same_device(bvh=bvh, queries=queries)
    device = queries.device
    m = int(queries.shape[0])

    if m == 0:
        return (
            wp.empty(0, dtype=wp.int32, device=device),
            wp.zeros(1 if include_total else 0, dtype=wp.int32, device=device),
        )

    hit_counts = wp.empty(m, dtype=wp.int32, device=device)
    # ``wp.uint64(bvh.id)`` explicitly: unlike ``wp.launch``, ``wp.map`` infers a bare Python int
    # scalar's dtype as ``wp.int32`` rather than matching the mapped @wp.func's declared parameter
    # type, and a mismatched dtype is a codegen-time TypeError, not a silent truncation.
    wp.map(
        kernel_neighbors.ball_count_in_bounds,
        wp.uint64(bvh.id),
        queries,
        wp.float32(radius),
        out=hit_counts,
    )

    segment_bounds, total_hits = tw.array.counts_to_offsets(hit_counts, include_total=True)
    offsets = segment_bounds if include_total else segment_bounds[:m]
    if total_hits == 0:
        return wp.empty(0, dtype=wp.int32, device=device), offsets

    candidate_indices_flat = wp.empty(total_hits, dtype=wp.int32, device=device)
    wp.launch(
        kernel_neighbors.query_bvh_ball_neighbors,
        dim=m,
        inputs=[queries, bvh.id, wp.float32(radius), segment_bounds[:m], candidate_indices_flat],
        device=device,
    )

    return candidate_indices_flat, offsets

query_bvh_box

query_bvh_box(
    bvh: Bvh,
    query_lower: array[vec3],
    query_upper: array[vec3],
    *,
    include_total: bool = False,
) -> tuple[wp.array[wp.int32], wp.array[wp.int32]]

Primitives overlapping one axis-aligned box per query, in the same CSR packing.

The box counterpart of the ball queries, and the per-query generalization of query_bvh_ball: each query carries its own (lower, upper) corners rather than sharing one cube size. On a BVH built by bvh_from_points, whose leaf bounds are degenerate, a hit means the point is inside the box, so the answer is exact and no narrow phase is needed; on one built by bvh_from_bounds it is box-versus-box overlap, like every other broad phase here.

Parameters:

Name	Type	Description	Default
bvh	Bvh	Pre-built BVH over points or bounds.	required
query_lower	array[vec3]	(m,) wp.vec3 corners of the query boxes, one pair per query.	required
query_upper	array[vec3]	(m,) wp.vec3 corners of the query boxes, one pair per query.	required
include_total	bool	Return offsets in the length-m + 1 CSR form whose trailing element is the total hit count, instead of the length-m form. Same meaning as on the ball queries.	False

Returns:

Type	Description
(candidate_indices_flat, offsets)	Query k owns candidate_indices_flat[offsets[k] : offsets[k + 1]], with offsets the exclusive prefix sum of per-query hit counts.

Raises:

Type	Description
ValueError	If query_lower and query_upper do not have the same length.
RuntimeError	If bvh, query_lower and query_upper are not all on one device.

Notes

The test is inclusive on every face: a point exactly on a box face (lower or upper) is inside it. A box with any upper < lower component matches nothing, and that is not checked -- the check would cost a host readback per call to reject a caller error whose answer is already empty.

There is no list-returning sibling, so the name carries no _with_offsets suffix: this query's consumers are kernels that recover the owning query from offsets, and a Python list of per-query arrays would add O(m) host slicing to a query whose whole point is that it is batched.

See Also

query_bvh_ball The same packing over the same bounds for a ball, and the one to prefer wherever the predicate is a ball rather than an axis-aligned region -- not only because the box admits corners the ball excludes, but because the box traversal is itself several times more expensive per candidate returned: the cost is dominated by the node test itself, not by how many candidates come back. triwarp.points.half_space_mask The unbounded counterpart: selection by one plane rather than by a box.

Source code in triwarp/neighbors.py

def query_bvh_box(
    bvh: wp.Bvh,
    query_lower: wp.array[wp.vec3],
    query_upper: wp.array[wp.vec3],
    *,
    include_total: bool = False,
) -> tuple[wp.array[wp.int32], wp.array[wp.int32]]:
    """
    Primitives overlapping one axis-aligned box **per query**, in the same CSR packing.

    The box counterpart of the ball queries, and the per-query generalization of
    [`query_bvh_ball`][triwarp.neighbors.query_bvh_ball]: each query
    carries its own ``(lower, upper)`` corners rather than sharing one cube size. On a BVH built by
    [`bvh_from_points`][triwarp.neighbors.bvh_from_points], whose leaf bounds are degenerate, a hit
    means the point is **inside** the box, so the answer is exact and no narrow phase is needed; on
    one built by [`bvh_from_bounds`][triwarp.neighbors.bvh_from_bounds] it is box-versus-box
    overlap, like every other broad phase here.

    Parameters
    ----------
    bvh
        Pre-built BVH over points or bounds.
    query_lower, query_upper
        ``(m,)`` ``wp.vec3`` corners of the query boxes, one pair per query.
    include_total
        Return ``offsets`` in the length-``m + 1`` CSR form whose trailing element is the total hit
        count, instead of the length-``m`` form. Same meaning as on the ball queries.

    Returns
    -------
    candidate_indices_flat, offsets
        Query ``k`` owns ``candidate_indices_flat[offsets[k] : offsets[k + 1]]``, with ``offsets``
        the exclusive prefix sum of per-query hit counts.

    Raises
    ------
    ValueError
        If ``query_lower`` and ``query_upper`` do not have the same length.
    RuntimeError
        If ``bvh``, ``query_lower`` and ``query_upper`` are not all on one device.

    Notes
    -----
    The test is **inclusive** on every face: a point exactly on a box face (lower or upper) is
    inside it. A box with any ``upper < lower`` component matches nothing, and that is not checked
    -- the check would cost a host readback per call to reject a caller error whose answer is
    already empty.

    There is no list-returning sibling, so the name carries no ``_with_offsets`` suffix: this
    query's consumers are kernels that recover the owning query from ``offsets``, and a Python list
    of per-query arrays would add ``O(m)`` host slicing to a query whose whole point is that it is
    batched.

    See Also
    --------
    [`query_bvh_ball`][triwarp.neighbors.query_bvh_ball]
        The same packing over the same bounds for a **ball**, and the one to prefer wherever the
        predicate is a ball rather than an axis-aligned region -- not only because the box admits
        corners the ball excludes, but because the box *traversal* is itself several times more
        expensive per candidate returned: the cost is dominated by the node test itself, not by
        how many candidates come back.
    [`triwarp.points.half_space_mask`][triwarp.points.half_space_mask]
        The unbounded counterpart: selection by one plane rather than by a box.
    """
    require_same_device(bvh=bvh, query_lower=query_lower, query_upper=query_upper)
    device = query_lower.device
    m = int(query_lower.shape[0])
    if int(query_upper.shape[0]) != m:
        raise ValueError("query_lower and query_upper must have the same length")

    if m == 0:
        return (
            wp.empty(0, dtype=wp.int32, device=device),
            wp.zeros(1 if include_total else 0, dtype=wp.int32, device=device),
        )

    hit_counts = wp.empty(m, dtype=wp.int32, device=device)
    # ``wp.uint64(bvh.id)`` explicitly -- see the same cast in ``query_bvh_ball``
    # above.
    wp.map(
        kernel_neighbors.aabb_count_in_bounds,
        wp.uint64(bvh.id),
        query_lower,
        query_upper,
        out=hit_counts,
    )

    segment_bounds, total_hits = tw.array.counts_to_offsets(hit_counts, include_total=True)
    offsets = segment_bounds if include_total else segment_bounds[:m]
    if total_hits == 0:
        return wp.empty(0, dtype=wp.int32, device=device), offsets

    candidate_indices_flat = wp.empty(total_hits, dtype=wp.int32, device=device)
    wp.launch(
        kernel_neighbors.query_bvh_box_neighbors,
        dim=m,
        inputs=[query_lower, query_upper, bvh.id, segment_bounds[:m], candidate_indices_flat],
        device=device,
    )

    return candidate_indices_flat, offsets

query_ball

query_ball(
    points: array[vec3],
    queries: array[vec3],
    r: float,
    *,
    accelerator: HashGrid | Bvh | None = ...,
    backend: QueryBackend | None = ...,
    grid_bins: int = ...,
    leaf_size: int = ...,
    return_sorted: bool = ...,
) -> tuple[
    list[wp.array[wp.int32]], list[wp.array[wp.float32]]
]

query_ball(
    points: array[vec3],
    queries: vec3,
    r: float,
    *,
    accelerator: HashGrid | Bvh | None = ...,
    backend: QueryBackend | None = ...,
    grid_bins: int = ...,
    leaf_size: int = ...,
    return_sorted: bool = ...,
) -> tuple[wp.array[wp.int32], wp.array[wp.float32]]

query_ball(
    points: array[vec3],
    queries: array[vec3] | vec3,
    r: float,
    *,
    accelerator: HashGrid | Bvh | None = None,
    backend: QueryBackend | None = None,
    grid_bins: int = 128,
    leaf_size: int = 4,
    return_sorted: bool = False,
) -> (
    tuple[
        list[wp.array[wp.int32]], list[wp.array[wp.float32]]
    ]
    | tuple[wp.array[wp.int32], wp.array[wp.float32]]
)

Find all data points within distance r of each query center (per-query arrays).

Same exact search as scipy.spatial.KDTree.query_ball_point with p=2 and eps=0; only the spatial index differs, and which index is a keyword rather than a function name (see backend). High-level wrapper around query_ball_with_offsets.

Unlike SciPy's object array of lists, multi-query results are two Python lists of length m, each element a rank-1 wp.array for that query. A single wp.vec3 query returns one (indices, distances) pair directly (not wrapped in lists). This clones each query's segment out of the internal flat buffer; for one flat buffer plus offsets on device, call query_ball_with_offsets instead.

Parameters:

Name	Type	Description	Default
points	array[vec3]	(n, 3) data points stored as wp.vec3.	required
queries	array[vec3] | vec3	Either (m, 3) query centers as wp.array[wp.vec3], or a single wp.vec3 (treated as one query).	required
r	float	Inclusion radius; cast to float32 in kernels (non-negative).	required
accelerator	HashGrid | Bvh | None	A prebuilt warp.HashGrid or warp.Bvh over points, to reuse across queries. It selects the backend by its own type, so backend is redundant when this is given and raises if it names the other one.	None
backend	QueryBackend | None	Which broad phase to use when accelerator is None: "hashgrid" (the default) enumerates the cells overlapping the query cube, "bvh" descends an AABB tree. The narrow phase and the answer are identical -- this is a cost choice, not a semantic one.	None
grid_bins	int	Grid resolution when building a hash grid. Ignored under backend="bvh" and whenever accelerator is given.	128
leaf_size	int	Maximum primitives per leaf when building a BVH. Ignored under backend="hashgrid" and whenever accelerator is given.	4
return_sorted	bool	If True, neighbors within each query are ordered by increasing distance. If False, order follows the broad phase's traversal (undefined ordering).	False

Returns:

Type	Description
(neighbor_indices, neighbor_distances)	If queries has m rows: list[wp.array[wp.int32]] and list[wp.array[wp.float32]], each of length m. Element k lists neighbors of queries[k] (indices into points and distances ‖points[i] - q‖₂). If queries is a single wp.vec3: two rank-1 arrays (possibly length 0), not lists. Empty points yields empty neighbor arrays and per-query empty slices; duplicate neighbors are not produced.

Raises:

Type	Description
ValueError	If backend is neither name, or contradicts the type of accelerator.
RuntimeError	If points, queries and accelerator are not all on one device.

Notes

SciPy may sort indices when return_sorted is left default on multi-point queries; here sorting only occurs when return_sorted=True, and sorts by distance, not by index. Ball boundaries use float32 arithmetic; extremely tight radii near representable limits may disagree slightly with pure float64 SciPy runs.

See Also

query_ball_with_offsets The flat CSR form, without cloning a segment per query. query_ball_count The counts alone, when the neighbors themselves are not wanted. hashgrid_from_points bvh_from_points scipy.spatial.KDTree.query_ball_point

Source code in triwarp/neighbors.py

def query_ball(
    points: wp.array[wp.vec3],
    queries: wp.array[wp.vec3] | wp.vec3,
    r: float,
    *,
    accelerator: wp.HashGrid | wp.Bvh | None = None,
    backend: QueryBackend | None = None,
    grid_bins: int = 128,
    leaf_size: int = 4,
    return_sorted: bool = False,
) -> (
    tuple[list[wp.array[wp.int32]], list[wp.array[wp.float32]]]
    | tuple[wp.array[wp.int32], wp.array[wp.float32]]
):
    """
    Find all data points within distance ``r`` of each query center (per-query arrays).

    Same exact search as [`scipy.spatial.KDTree.query_ball_point`][] with ``p=2`` and ``eps=0``;
    only the spatial index differs, and which index is a keyword rather than a function name (see
    ``backend``). High-level wrapper around
    [`query_ball_with_offsets`][triwarp.neighbors.query_ball_with_offsets].

    Unlike SciPy's object array of lists, multi-query results are two Python lists of length ``m``,
    each element a rank-1 ``wp.array`` for that query. A single ``wp.vec3`` query returns one
    ``(indices, distances)`` pair directly (not wrapped in lists). This clones each query's segment
    out of the internal flat buffer; for one flat buffer plus offsets on device, call
    [`query_ball_with_offsets`][triwarp.neighbors.query_ball_with_offsets] instead.

    Parameters
    ----------
    points
        ``(n, 3)`` data points stored as ``wp.vec3``.
    queries
        Either ``(m, 3)`` query centers as ``wp.array[wp.vec3]``, or a single ``wp.vec3``
        (treated as one query).
    r
        Inclusion radius; cast to ``float32`` in kernels (non-negative).
    accelerator
        A prebuilt ``warp.HashGrid`` or ``warp.Bvh`` over ``points``, to reuse across queries. It
        selects the backend by its own type, so ``backend`` is redundant when this is given and
        raises if it names the other one.
    backend
        Which broad phase to use when ``accelerator`` is ``None``: ``"hashgrid"`` (the default)
        enumerates the cells overlapping the query cube, ``"bvh"`` descends an AABB tree. The
        narrow phase and the answer are identical -- this is a cost choice, not a semantic one.
    grid_bins
        Grid resolution when building a hash grid. Ignored under ``backend="bvh"`` and whenever
        ``accelerator`` is given.
    leaf_size
        Maximum primitives per leaf when building a BVH. Ignored under ``backend="hashgrid"`` and
        whenever ``accelerator`` is given.
    return_sorted
        If ``True``, neighbors within each query are ordered by increasing distance. If ``False``,
        order follows the broad phase's traversal (undefined ordering).

    Returns
    -------
    neighbor_indices, neighbor_distances
        If ``queries`` has ``m`` rows: ``list[wp.array[wp.int32]]`` and
        ``list[wp.array[wp.float32]]``, each of length ``m``. Element ``k`` lists neighbors
        of ``queries[k]`` (indices into ``points`` and distances ``‖points[i] - q‖₂``).

        If ``queries`` is a single ``wp.vec3``: two rank-1 arrays (possibly length 0), not lists.

        Empty ``points`` yields empty neighbor arrays and per-query empty slices; duplicate
        neighbors are not produced.

    Raises
    ------
    ValueError
        If ``backend`` is neither name, or contradicts the type of ``accelerator``.
    RuntimeError
        If ``points``, ``queries`` and ``accelerator`` are not all on one device.

    Notes
    -----
    SciPy may sort indices when ``return_sorted`` is left default on multi-point queries; here
    sorting only occurs when ``return_sorted=True``, and sorts by distance, not by index. Ball
    boundaries use ``float32`` arithmetic; extremely tight radii near representable limits may
    disagree slightly with pure ``float64`` SciPy runs.

    See Also
    --------
    [`query_ball_with_offsets`][triwarp.neighbors.query_ball_with_offsets]
        The flat CSR form, without cloning a segment per query.
    [`query_ball_count`][triwarp.neighbors.query_ball_count]
        The counts alone, when the neighbors themselves are not wanted.
    [`hashgrid_from_points`][triwarp.neighbors.hashgrid_from_points]
    [`bvh_from_points`][triwarp.neighbors.bvh_from_points]
    [`scipy.spatial.KDTree.query_ball_point`][]
    """
    require_same_device(points=points, queries=queries, accelerator=accelerator)
    device = points.device

    single_query = isinstance(queries, wp.vec3)
    if single_query:
        queries = wp.array([queries], dtype=wp.vec3, device=device)

    neighbor_indices_flat, neighbor_distances_flat, offsets = query_ball_with_offsets(
        points,
        queries,
        r,
        accelerator=accelerator,
        backend=backend,
        grid_bins=grid_bins,
        leaf_size=leaf_size,
        return_sorted=return_sorted,
    )
    neighbor_indices = tw.array.split(neighbor_indices_flat, offsets, copy=True)
    neighbor_distances = tw.array.split(neighbor_distances_flat, offsets, copy=True)

    if single_query:
        return neighbor_indices[0], neighbor_distances[0]
    return neighbor_indices, neighbor_distances

query_ball_count

query_ball_count(
    points: array[vec3],
    queries: array[vec3],
    r: float,
    *,
    accelerator: HashGrid | Bvh | None = None,
    backend: QueryBackend | None = None,
    grid_bins: int = 128,
    leaf_size: int = 4,
) -> wp.array[wp.int32]

Count neighbors of each query within Euclidean distance r.

For each query center q, returns how many entries p in points satisfy ‖p - q‖₂ ≤ r. This matches scipy.spatial.KDTree.query_ball_point with p=2, eps=0, and return_length=True (exact search; only the spatial index differs).

The narrow phase keeps points with float32 Euclidean distance at most r whichever broad phase ran, so the count does not depend on backend.

Parameters:

Name	Type	Description	Default
points	array[vec3]	(n, 3) data points stored as wp.vec3.	required
queries	array[vec3]	(m, 3) query centers stored as wp.vec3.	required
r	float	Inclusion radius; cast to float32 in kernels (non-negative).	required
accelerator	HashGrid | Bvh | None	As in query_ball.	None
backend	HashGrid | Bvh | None	As in query_ball.	None
grid_bins	HashGrid | Bvh | None	As in query_ball.	None
leaf_size	HashGrid | Bvh | None	As in query_ball.	None

Returns:

Type	Description
array[int32]	Length-m device array whose k-th element is the neighbor count for queries[k]. If n == 0, returns zeros.

Raises:

Type	Description
ValueError	If backend is neither name, or contradicts the type of accelerator.
RuntimeError	If points, queries and accelerator are not all on one device.

See Also

query_ball hashgrid_from_points bvh_from_points scipy.spatial.KDTree.query_ball_point

Source code in triwarp/neighbors.py

def query_ball_count(
    points: wp.array[wp.vec3],
    queries: wp.array[wp.vec3],
    r: float,
    *,
    accelerator: wp.HashGrid | wp.Bvh | None = None,
    backend: QueryBackend | None = None,
    grid_bins: int = 128,
    leaf_size: int = 4,
) -> wp.array[wp.int32]:
    """
    Count neighbors of each query within Euclidean distance ``r``.

    For each query center ``q``, returns how many entries ``p`` in ``points`` satisfy
    ``‖p - q‖₂ ≤ r``. This matches [`scipy.spatial.KDTree.query_ball_point`][] with ``p=2``,
    ``eps=0``, and ``return_length=True`` (exact search; only the spatial index differs).

    The narrow phase keeps points with ``float32`` Euclidean distance at most ``r`` whichever
    broad phase ran, so the count does not depend on ``backend``.

    Parameters
    ----------
    points
        ``(n, 3)`` data points stored as ``wp.vec3``.
    queries
        ``(m, 3)`` query centers stored as ``wp.vec3``.
    r
        Inclusion radius; cast to ``float32`` in kernels (non-negative).
    accelerator, backend, grid_bins, leaf_size
        As in [`query_ball`][triwarp.neighbors.query_ball].

    Returns
    -------
    wp.array[wp.int32]
        Length-``m`` device array whose ``k``-th element is the neighbor count for
        ``queries[k]``. If ``n == 0``, returns zeros.

    Raises
    ------
    ValueError
        If ``backend`` is neither name, or contradicts the type of ``accelerator``.
    RuntimeError
        If ``points``, ``queries`` and ``accelerator`` are not all on one device.

    See Also
    --------
    [`query_ball`][triwarp.neighbors.query_ball]
    [`hashgrid_from_points`][triwarp.neighbors.hashgrid_from_points]
    [`bvh_from_points`][triwarp.neighbors.bvh_from_points]
    [`scipy.spatial.KDTree.query_ball_point`][]
    """
    require_same_device(points=points, queries=queries, accelerator=accelerator)
    kind, accelerator = _resolve_accelerator(accelerator, backend)
    device = points.device
    n = int(points.shape[0])
    m = int(queries.shape[0])
    if n == 0:
        return wp.zeros(m, dtype=wp.int32, device=device)

    neighbor_counts = wp.empty(m, dtype=wp.int32, device=device)
    if kind == "hashgrid":
        if accelerator is None:
            accelerator = hashgrid_from_points(points, r, grid_bins)
        accel_selector = kernel_neighbors.ACCEL_HASHGRID
    else:
        if accelerator is None:
            accelerator = bvh_from_points(points, leaf_size)
        accel_selector = kernel_neighbors.ACCEL_BVH

    wp.launch(
        kernel_neighbors.query_ball_count,
        dim=m,
        inputs=[points, queries, accel_selector, accelerator.id, wp.float32(r), neighbor_counts],
        device=device,
    )
    return neighbor_counts

query_ball_with_offsets

query_ball_with_offsets(
    points: array[vec3],
    queries: array[vec3] | vec3,
    r: float,
    *,
    accelerator: HashGrid | Bvh | None = None,
    backend: QueryBackend | None = None,
    grid_bins: int = 128,
    leaf_size: int = 4,
    return_sorted: bool = False,
    include_total: bool = False,
) -> tuple[
    wp.array[wp.int32],
    wp.array[wp.float32],
    wp.array[wp.int32],
]

Low-level ball query: neighbors in one concatenated pair plus per-query offsets.

Same geometry as query_ball (broad phase out to r, float32 test ‖points[i] - q‖₂ ≤ r). Semantics match scipy.spatial.KDTree.query_ball_point with p=2 and eps=0.

Prefer query_ball for a Python list of one array per query; use this when you want a single flat buffer on device (e.g. fused downstream kernels) and CSR-style boundaries without cloning each segment.

Parameters:

Name	Type	Description	Default
points	array[vec3]	(n, 3) data points stored as wp.vec3.	required
queries	array[vec3] | vec3	Either (m, 3) query centers as wp.array[wp.vec3], or a single wp.vec3 (treated as one query).	required
r	float	Inclusion radius; cast to float32 in kernels (non-negative).	required
accelerator	HashGrid | Bvh | None	As in query_ball.	None
backend	HashGrid | Bvh | None	As in query_ball.	None
grid_bins	HashGrid | Bvh | None	As in query_ball.	None
leaf_size	HashGrid | Bvh | None	As in query_ball.	None
return_sorted	bool	If True, neighbors within each query are ordered by increasing distance.	False
include_total	bool	Return offsets in the length-m + 1 CSR form, whose trailing element is the total neighbor count, instead of the length-m form. Free -- the terminator is built either way, since the return_sorted path needs it as segment bounds.	False

Returns:

Type

Description

(neighbor_indices_flat, neighbor_distances_flat, offsets)

Three rank-1 arrays. Let m = queries.shape[0] after any wp.vec3 wrap. offsets has length m (or m + 1 with include_total) and is the exclusive prefix sum of per-query neighbor counts (same layout as wp.utils.array_scan(..., inclusive=False)): query k owns neighbor_indices_flat[offsets[k] : offsets[k+1]] where offsets[m] is neighbor_indices_flat.shape[0] (the total neighbor count), stored only in the CSR form. neighbor_indices_flat and neighbor_distances_flat have that total length and list point indices and distances ‖points[i] - q‖₂ in parallel. Empty points still returns zero offsets; empty neighbor sets yield length-0 flat arrays and zero offsets.

Raises:

Type	Description
ValueError	If backend is neither name, or contradicts the type of accelerator.
RuntimeError	If points, queries and accelerator are not all on one device.

Notes

SciPy may sort indices when return_sorted is left default on multi-point queries; here sorting only occurs when return_sorted=True, and sorts by distance, not by index. Ball boundaries use float32 arithmetic; extremely tight radii near representable limits may disagree slightly with pure float64 SciPy runs.

See Also

query_ball query_ball_count hashgrid_from_points bvh_from_points scipy.spatial.KDTree.query_ball_point

Source code in triwarp/neighbors.py

def query_ball_with_offsets(
    points: wp.array[wp.vec3],
    queries: wp.array[wp.vec3] | wp.vec3,
    r: float,
    *,
    accelerator: wp.HashGrid | wp.Bvh | None = None,
    backend: QueryBackend | None = None,
    grid_bins: int = 128,
    leaf_size: int = 4,
    return_sorted: bool = False,
    include_total: bool = False,
) -> tuple[wp.array[wp.int32], wp.array[wp.float32], wp.array[wp.int32]]:
    """
    Low-level ball query: neighbors in one concatenated pair plus per-query offsets.

    Same geometry as [`query_ball`][triwarp.neighbors.query_ball] (broad phase out to ``r``,
    ``float32`` test ``‖points[i] - q‖₂ ≤ r``). Semantics match
    [`scipy.spatial.KDTree.query_ball_point`][] with ``p=2`` and ``eps=0``.

    Prefer [`query_ball`][triwarp.neighbors.query_ball] for a Python list of one array per query;
    use this when you want a single flat buffer on device (e.g. fused downstream kernels) and
    CSR-style boundaries without cloning each segment.

    Parameters
    ----------
    points
        ``(n, 3)`` data points stored as ``wp.vec3``.
    queries
        Either ``(m, 3)`` query centers as ``wp.array[wp.vec3]``, or a single ``wp.vec3``
        (treated as one query).
    r
        Inclusion radius; cast to ``float32`` in kernels (non-negative).
    accelerator, backend, grid_bins, leaf_size
        As in [`query_ball`][triwarp.neighbors.query_ball].
    return_sorted
        If ``True``, neighbors within each query are ordered by increasing distance.
    include_total
        Return ``offsets`` in the length-``m + 1`` CSR form, whose trailing element is the total
        neighbor count, instead of the length-``m`` form. Free -- the terminator is built either
        way, since the ``return_sorted`` path needs it as segment bounds.

    Returns
    -------
    neighbor_indices_flat, neighbor_distances_flat, offsets
        Three rank-1 arrays. Let ``m = queries.shape[0]`` after any ``wp.vec3`` wrap.

        ``offsets`` has length ``m`` (or ``m + 1`` with ``include_total``) and is the exclusive
        prefix sum of per-query neighbor counts (same layout as
        ``wp.utils.array_scan(..., inclusive=False)``): query ``k`` owns
        ``neighbor_indices_flat[offsets[k] : offsets[k+1]]`` where ``offsets[m]`` is
        ``neighbor_indices_flat.shape[0]`` (the total neighbor count), stored only in the
        CSR form.

        ``neighbor_indices_flat`` and ``neighbor_distances_flat`` have that total length
        and list point indices and distances ``‖points[i] - q‖₂`` in parallel. Empty
        ``points`` still returns zero ``offsets``; empty neighbor sets yield
        length-0 flat arrays and zero ``offsets``.

    Raises
    ------
    ValueError
        If ``backend`` is neither name, or contradicts the type of ``accelerator``.
    RuntimeError
        If ``points``, ``queries`` and ``accelerator`` are not all on one device.

    Notes
    -----
    SciPy may sort indices when ``return_sorted`` is left default on multi-point queries; here
    sorting only occurs when ``return_sorted=True``, and sorts by distance, not by index. Ball
    boundaries use ``float32`` arithmetic; extremely tight radii near representable limits may
    disagree slightly with pure ``float64`` SciPy runs.

    See Also
    --------
    [`query_ball`][triwarp.neighbors.query_ball]
    [`query_ball_count`][triwarp.neighbors.query_ball_count]
    [`hashgrid_from_points`][triwarp.neighbors.hashgrid_from_points]
    [`bvh_from_points`][triwarp.neighbors.bvh_from_points]
    [`scipy.spatial.KDTree.query_ball_point`][]
    """
    require_same_device(points=points, queries=queries, accelerator=accelerator)
    kind, resolved = _resolve_accelerator(accelerator, backend)
    if kind == "hashgrid":
        return _ball_with_offsets(
            points,
            queries,
            r,
            build_accelerator=lambda: (
                resolved if resolved is not None else hashgrid_from_points(points, r, grid_bins)
            ),
            count_fn=lambda pts, qrs, radius, built: query_ball_count(
                pts, qrs, radius, accelerator=built
            ),
            accel=kernel_neighbors.ACCEL_HASHGRID,
            include_total=include_total,
            return_sorted=return_sorted,
        )
    return _ball_with_offsets(
        points,
        queries,
        r,
        build_accelerator=lambda: (
            resolved if resolved is not None else bvh_from_points(points, leaf_size)
        ),
        count_fn=lambda pts, qrs, radius, built: query_ball_count(
            pts, qrs, radius, accelerator=built
        ),
        accel=kernel_neighbors.ACCEL_BVH,
        include_total=include_total,
        return_sorted=return_sorted,
    )

knn_initial_radius

knn_initial_radius(
    points: array[vec3],
    k: int,
    *,
    bounds: tuple[vec3, vec3] | None = None,
) -> float

Radius at which a k-nearest search is expected to succeed on the first try.

Inverts a uniform-density model of points: the smallest ball expected to hold k of n points is the one whose volume is k / n of the bounding box's. This is the default initial_radius of query_nearest, which deepens from it until each row certifies itself, so the value affects speed only and never the result.

A cloud that is flat or collinear has a (near-)zero box volume, which the 3-D formula would invert into a useless radius, so the effective dimension d is taken to be the number of axes whose extent is non-negligible against the largest, and the matching formula is used: d = 3 gives (3 f V / 4pi)^(1/3), d = 2 gives (f A / pi)^(1/2) and d = 1 gives f L / 2, with fill fraction f = k / n.

Parameters:

Name	Type	Description	Default
points	array[vec3]	(n,) data points as wp.vec3.	required
k	int	Number of neighbors the search will ask for; must be >= 1.	required
bounds	tuple[vec3, vec3] | None	Optional (min_bound, max_bound) from aabb. Pass it to reuse a reduction you already ran; otherwise it is computed here (one device reduction plus one readback). Trusted, not validated — but only the estimate depends on it, so a wrong box costs speed and not correctness, unlike the identically named argument of query_nearest.	None

Returns:

Type	Description
float	The estimated radius, or math.inf when no finite estimate is meaningful — n == 0, k >= n (every point is a neighbor) or a degenerate box (all points coincident). The query kernels read math.inf as "one complete scan", which is the right answer in each of those cases.

See Also

query_nearest aabb

Source code in triwarp/neighbors.py

def knn_initial_radius(
    points: wp.array[wp.vec3], k: int, *, bounds: tuple[wp.vec3, wp.vec3] | None = None
) -> float:
    """
    Radius at which a ``k``-nearest search is expected to succeed on the first try.

    Inverts a uniform-density model of ``points``: the smallest ball expected to hold ``k`` of
    ``n`` points is the one whose volume is ``k / n`` of the bounding box's. This is the default
    ``initial_radius`` of [`query_nearest`][triwarp.neighbors.query_nearest], which deepens from it
    until each row certifies itself, so the value affects **speed only** and never the result.

    A cloud that is flat or collinear has a (near-)zero box volume, which the 3-D formula would
    invert into a useless radius, so the effective dimension ``d`` is taken to be the number of
    axes whose extent is non-negligible against the largest, and the matching formula is used:
    ``d = 3`` gives ``(3 f V / 4pi)^(1/3)``, ``d = 2`` gives ``(f A / pi)^(1/2)`` and ``d = 1``
    gives ``f L / 2``, with fill fraction ``f = k / n``.

    Parameters
    ----------
    points
        ``(n,)`` data points as ``wp.vec3``.
    k
        Number of neighbors the search will ask for; must be ``>= 1``.
    bounds
        Optional ``(min_bound, max_bound)`` from
        [`aabb`][triwarp.bounds.aabb]. Pass it to reuse a reduction you already ran;
        otherwise it is computed here (one device reduction plus one readback). Trusted, not
        validated — but only the estimate depends on it, so a wrong box costs speed and not
        correctness, unlike the identically named argument of
        [`query_nearest`][triwarp.neighbors.query_nearest].

    Returns
    -------
    float
        The estimated radius, or ``math.inf`` when no finite estimate is meaningful — ``n == 0``,
        ``k >= n`` (every point is a neighbor) or a degenerate box (all points coincident). The
        query kernels read ``math.inf`` as "one complete scan", which is the right answer in each
        of those cases.

    See Also
    --------
    [`query_nearest`][triwarp.neighbors.query_nearest]
    [`aabb`][triwarp.bounds.aabb]
    """
    n = int(points.shape[0])
    if n == 0 or k >= n:
        return math.inf

    if bounds is None:
        bounds = tw.bounds.aabb(points)
    min_bound, max_bound = bounds
    extents = sorted((float(max_bound[axis] - min_bound[axis]) for axis in range(3)), reverse=True)
    if extents[0] <= 0.0:
        return math.inf

    fill = k / n
    dimension = sum(extent > _FLAT_AXIS_FRACTION * extents[0] for extent in extents)
    if dimension == 3:
        volume = extents[0] * extents[1] * extents[2]
        return (3.0 * fill * volume / (4.0 * math.pi)) ** (1.0 / 3.0)
    if dimension == 2:
        return math.sqrt(fill * extents[0] * extents[1] / math.pi)
    return 0.5 * fill * extents[0]

query_nearest

query_nearest(
    points: array[vec3],
    queries: vec3,
    k: int,
    *,
    accelerator: HashGrid | Bvh | None = ...,
    backend: QueryBackend | None = ...,
    max_radius: float = ...,
    grid_bins: int = ...,
    leaf_size: int = ...,
    initial_radius: float | None = ...,
    bounds: tuple[vec3, vec3] | None = ...,
) -> tuple[twt.Array1dInt32, twt.Array1dFloat32]

query_nearest(
    points: array[vec3],
    queries: array[vec3] | vec3,
    k: Literal[1] = 1,
    *,
    accelerator: HashGrid | Bvh | None = ...,
    backend: QueryBackend | None = ...,
    max_radius: float = ...,
    grid_bins: int = ...,
    leaf_size: int = ...,
    initial_radius: float | None = ...,
    bounds: tuple[vec3, vec3] | None = ...,
) -> tuple[twt.Array1dInt32, twt.Array1dFloat32]

query_nearest(
    points: array[vec3],
    queries: array[vec3],
    k: int,
    *,
    accelerator: HashGrid | Bvh | None = ...,
    backend: QueryBackend | None = ...,
    max_radius: float = ...,
    grid_bins: int = ...,
    leaf_size: int = ...,
    initial_radius: float | None = ...,
    bounds: tuple[vec3, vec3] | None = ...,
) -> tuple[twt.Array2dInt32, twt.Array2dFloat32]

query_nearest(
    points: array[vec3],
    queries: array[vec3] | vec3,
    k: int = 1,
    *,
    accelerator: HashGrid | Bvh | None = None,
    backend: QueryBackend | None = None,
    max_radius: float = math.inf,
    grid_bins: int = 128,
    leaf_size: int = 4,
    initial_radius: float | None = None,
    bounds: tuple[vec3, vec3] | None = None,
) -> tuple[
    twt.Array2dInt32 | twt.Array1dInt32,
    twt.Array2dFloat32 | twt.Array1dFloat32,
]

For each query center, find the k nearest data points.

Iterative deepening: each query grows a search radius from initial_radius until its k-th neighbour is closer than the radius, which certifies the row -- so the answer is exact and does not depend on backend. The candidate row is register-resident for k <= 64.

Which backend to use

"hashgrid" (the default) is the faster broad phase when the query scale matches the cloud's density, because initial_radius sets the cell width: the walk visits (2 ceil(r / cell) + 1) ** 3 cells, so a radius f times too large costs f ** 3, and the search falls back to an exact linear scan once r outgrows a few cells. That also makes it sensitive to a cloud whose density is not uniform, and to queries drawn from a different distribution than the points, since the cell width was sized for the data rather than for where the queries land.

"bvh" has no cell width to get wrong, so it is the one to reach for when the query scale is unknown, the cloud is non-uniform, or the queries sit outside it. Its own cost grows with k faster than the grid's, so at small k the grid tends to win.

But at moderate k the tree can win on a perfectly uniform cloud too, which is not what "an awkward one" suggests. The grid's cost is non-monotonic in cloud size at moderate k: a row whose true k-th distance runs past _knn_widest_grid_radius abandons the cell walk for an exact linear scan of the whole cloud, and how much of the cloud falls in that tail moves with n. The tree's cost stays monotonic over the same range. Both answers are still exact; only the cost moves.

Both are exact; this is a cost choice only. Reuse the structure across calls by passing it as accelerator when several queries share one cloud.

Parameters:

Name	Type	Description	Default
points	array[vec3]	(n, 3) data points stored as wp.vec3.	required
queries	array[vec3] | vec3	Either (m, 3) query centers as wp.array[wp.vec3], or a single wp.vec3.	required
k	int	Number of neighbours per query, >= 1.	1
accelerator	HashGrid | Bvh | None	A prebuilt warp.HashGrid or warp.Bvh over points, to reuse across queries. It selects the backend by its own type, so backend is redundant when this is given and raises if it names the other one. A hash grid passed here keeps its own cell_width rather than one derived from initial_radius.	None
backend	QueryBackend | None	"hashgrid" (the default) or "bvh", when accelerator is None. See the note above.	None
max_radius	float	Stop deepening past this distance and leave the remaining slots unfilled (index -1, distance inf). Defaults to unbounded.	inf
grid_bins	int	Grid resolution when building a hash grid. Ignored under backend="bvh" and whenever accelerator is given.	128
leaf_size	int	Maximum primitives per leaf when building a BVH. Ignored under backend="hashgrid" and whenever accelerator is given.	4
initial_radius	float | None	First search radius. Defaults to knn_initial_radius, which estimates it from the cloud's density -- so pass it explicitly when the queries are at a different scale, since under "hashgrid" this also sets the cell width.	None
bounds	tuple[vec3, vec3] | None	(min_bound, max_bound) of points, to skip triwarp.bounds.aabb and its readback. Trusted, not validated, and not only a shortcut: it also fixes the per-query radius at which a scan is provably complete, so a box that does not contain points can end the search early and leave slots unfilled (-1 and inf) rather than merely running slower. Pass the real bounding box, or leave it None.	None

Returns:

Type

Description

(neighbor_indices, neighbor_distances)

(m, k) wp.int32 indices into points and (m, k) wp.float32 distances, each row sorted by increasing distance. A slot no neighbour was found for holds -1 and inf. The trailing axis collapses when there is only one of it, as scipy.spatial.KDTree.query does: at k == 1 the two are rank-1 of length m, and for a single wp.vec3 query they are rank-1 of length k. The rank depends only on k and the form of queries, never on whether an answer was found — an empty points or an empty queries returns the same rank a populated one would.

Raises:

Type	Description
ValueError	If k < 1, max_radius < 0, initial_radius < 0, or backend is neither name or contradicts the type of accelerator.
RuntimeError	If points, queries and accelerator are not all on one device.

See Also

knn_initial_radius The default first radius, and the one parameter worth passing by hand. query_ball A fixed radius rather than a fixed count. nearest_neighbor_distance The k=2 self-query, which is the cloud's spacing. hashgrid_from_points bvh_from_points scipy.spatial.KDTree.query

Source code in triwarp/neighbors.py

def query_nearest(
    points: wp.array[wp.vec3],
    queries: wp.array[wp.vec3] | wp.vec3,
    k: int = 1,
    *,
    accelerator: wp.HashGrid | wp.Bvh | None = None,
    backend: QueryBackend | None = None,
    max_radius: float = math.inf,
    grid_bins: int = 128,
    leaf_size: int = 4,
    initial_radius: float | None = None,
    bounds: tuple[wp.vec3, wp.vec3] | None = None,
) -> tuple[twt.Array2dInt32 | twt.Array1dInt32, twt.Array2dFloat32 | twt.Array1dFloat32]:
    """
    For each query center, find the ``k`` nearest data points.

    Iterative deepening: each query grows a search radius from ``initial_radius`` until its ``k``-th
    neighbour is closer than the radius, which certifies the row -- so the answer is exact and does
    not depend on ``backend``. The candidate row is register-resident for ``k <= 64``.

    !!! note "Which backend to use"
        ``"hashgrid"`` (the default) is the faster broad phase when the query scale matches the
        cloud's density, because ``initial_radius`` sets the cell width: the walk visits
        ``(2 ceil(r / cell) + 1) ** 3`` cells, so a radius **f** times too large costs ``f ** 3``,
        and the search falls back to an exact linear scan once ``r`` outgrows a few cells. That also
        makes it sensitive to a cloud whose density is not uniform, and to queries drawn from a
        different distribution than the points, since the cell width was sized for the data rather
        than for where the queries land.

        ``"bvh"`` has no cell width to get wrong, so it is the one to reach for when the query scale
        is unknown, the cloud is non-uniform, or the queries sit outside it. Its own cost grows
        with ``k`` faster than the grid's, so at small ``k`` the grid tends to win.

        **But at moderate ``k`` the tree can win on a perfectly uniform cloud too, which is not what
        "an awkward one" suggests.** The grid's cost is *non-monotonic* in cloud size at moderate
        ``k``: a row whose true ``k``-th distance runs past ``_knn_widest_grid_radius`` abandons
        the cell walk for an exact linear scan of the whole cloud, and how much of the cloud falls
        in that tail moves with ``n``. The tree's cost stays monotonic over the same range. Both
        answers are still exact; only the cost moves.

        Both are exact; this is a cost choice only. Reuse the structure across calls by passing it
        as ``accelerator`` when several queries share one cloud.

    Parameters
    ----------
    points
        ``(n, 3)`` data points stored as ``wp.vec3``.
    queries
        Either ``(m, 3)`` query centers as ``wp.array[wp.vec3]``, or a single ``wp.vec3``.
    k
        Number of neighbours per query, ``>= 1``.
    accelerator
        A prebuilt ``warp.HashGrid`` or ``warp.Bvh`` over ``points``, to reuse across queries. It
        selects the backend by its own type, so ``backend`` is redundant when this is given and
        raises if it names the other one. A hash grid passed here keeps its own ``cell_width``
        rather than one derived from ``initial_radius``.
    backend
        ``"hashgrid"`` (the default) or ``"bvh"``, when ``accelerator`` is ``None``. See the note
        above.
    max_radius
        Stop deepening past this distance and leave the remaining slots unfilled (index ``-1``,
        distance ``inf``). Defaults to unbounded.
    grid_bins
        Grid resolution when building a hash grid. Ignored under ``backend="bvh"`` and whenever
        ``accelerator`` is given.
    leaf_size
        Maximum primitives per leaf when building a BVH. Ignored under ``backend="hashgrid"`` and
        whenever ``accelerator`` is given.
    initial_radius
        First search radius. Defaults to
        [`knn_initial_radius`][triwarp.neighbors.knn_initial_radius], which estimates it from the
        *cloud's* density -- so pass it explicitly when the queries are at a different scale, since
        under ``"hashgrid"`` this also sets the cell width.
    bounds
        ``(min_bound, max_bound)`` of ``points``, to skip
        [`triwarp.bounds.aabb`][triwarp.bounds.aabb] and its readback. **Trusted, not validated, and
        not only a shortcut**: it also fixes the per-query radius at which a scan is provably
        complete, so a box that does not contain ``points`` can end the search early and leave
        slots unfilled (``-1`` and ``inf``) rather than merely running slower. Pass the real
        bounding box, or leave it ``None``.

    Returns
    -------
    neighbor_indices, neighbor_distances
        ``(m, k)`` ``wp.int32`` indices into ``points`` and ``(m, k)`` ``wp.float32`` distances,
        each row sorted by increasing distance. A slot no neighbour was found for holds ``-1`` and
        ``inf``.

        The trailing axis collapses when there is only one of it, as
        [`scipy.spatial.KDTree.query`][] does: at ``k == 1`` the two are rank-1 of length ``m``, and
        for a single ``wp.vec3`` query they are rank-1 of length ``k``. The rank depends only on
        ``k`` and the form of ``queries``, never on whether an answer was found — an empty
        ``points`` or an empty ``queries`` returns the same rank a populated one would.

    Raises
    ------
    ValueError
        If ``k < 1``, ``max_radius < 0``, ``initial_radius < 0``, or ``backend`` is neither name or
        contradicts the type of ``accelerator``.
    RuntimeError
        If ``points``, ``queries`` and ``accelerator`` are not all on one device.

    See Also
    --------
    [`knn_initial_radius`][triwarp.neighbors.knn_initial_radius]
        The default first radius, and the one parameter worth passing by hand.
    [`query_ball`][triwarp.neighbors.query_ball]
        A fixed radius rather than a fixed count.
    [`nearest_neighbor_distance`][triwarp.neighbors.nearest_neighbor_distance]
        The ``k=2`` self-query, which is the cloud's spacing.
    [`hashgrid_from_points`][triwarp.neighbors.hashgrid_from_points]
    [`bvh_from_points`][triwarp.neighbors.bvh_from_points]
    [`scipy.spatial.KDTree.query`][]
    """
    require_same_device(points=points, queries=queries, accelerator=accelerator)
    kind, resolved = _resolve_accelerator(accelerator, backend)
    _validate_nearest(k, max_radius, initial_radius)

    device = points.device
    single_query = isinstance(queries, wp.vec3)
    if single_query:
        queries = wp.array([queries], dtype=wp.vec3, device=device)

    m = int(queries.shape[0])
    n = int(points.shape[0])

    # Both degenerate inputs answer "every slot unfilled" over `(m, k)`, which is empty of its own
    # accord when there are no queries -- and both must go through the same shaping as the general
    # path, or `k == 1` comes back rank-2 here and rank-1 everywhere else.
    if m == 0 or n == 0:
        return _empty_nearest(m, k, single_query, device)

    if bounds is None:
        bounds = tw.bounds.aabb(points)
    min_bound, max_bound = bounds
    if initial_radius is None:
        initial_radius = knn_initial_radius(points, k, bounds=bounds)

    if kind == "bvh":
        bvh = resolved if resolved is not None else bvh_from_points(points, leaf_size)
        kernel = kernel_neighbors.bvh_nearest_kernel(k)
        accel_id = bvh.id
        # The BVH follows an unbounded radius, so it needs no linear-scan cutover argument.
        widest: list[wp.float32] = []
    else:
        if resolved is None:
            cell_size = _knn_cell_size(initial_radius, min_bound, max_bound, grid_bins)
            grid = hashgrid_from_points(points, cell_size, grid_bins)
        else:
            grid = resolved
            cell_size = float(getattr(grid, "cell_width", initial_radius))
        kernel = kernel_neighbors.hashgrid_nearest_kernel(k)
        accel_id = grid.id
        widest = [wp.float32(_knn_widest_grid_radius(cell_size, n))]

    # ``wp.empty``, not ``wp.full``: every row is written in full by the kernel, so pre-filling
    # here would be two wasted launches.
    neighbor_indices = twt.empty_2d((m, k), wp.int32, device=device)
    neighbor_distances = twt.empty_2d((m, k), wp.float32, device=device)
    wp.launch(
        kernel,
        dim=m,
        inputs=[
            points,
            queries,
            accel_id,
            wp.int32(k),
            wp.float32(max_radius),
            wp.float32(initial_radius),
            *widest,
            min_bound,
            max_bound,
            neighbor_indices,
            neighbor_distances,
        ],
        device=device,
    )
    return _shape_nearest(neighbor_indices, neighbor_distances, k, single_query)

query_weighted_nearest

query_weighted_nearest(
    points: array[vec3],
    weights: array[float32],
    queries: array[vec3],
    *,
    max_weight: float | None = None,
    accelerator: Bvh | None = None,
    leaf_size: int = 4,
) -> tuple[wp.array[wp.int32], wp.array[wp.float32]]

Nearest site under the weighted distance |p - q| - w(p).

Each site carries a radius, and the winner is the one whose surface is closest rather than whose centre is -- the additively weighted (Apollonius) nearest-neighbour query, which is what picks the influencing site when the sites have different scales: a sphere set, a level-of-detail cluster, a set of samples with per-sample confidence. Plain query_nearest is the w = 0 case, and the two genuinely differ -- on a random 40-site cloud, 1 query in 6 had a different winner.

Parameters:

Name	Type	Description	Default
points	array[vec3]	(n,) site positions as wp.vec3.	required
weights	array[float32]	(n,) per-site weights, subtracted from the distance. Larger wins ties of distance; may be negative, which pushes a site away.	required
queries	array[vec3]	(m,) query positions as wp.vec3.	required
max_weight	float | None	An upper bound on weights, which is what makes the search prunable. None reduces weights on the device and reads the maximum back (one readback), so pass it when the bound is already known -- a radius cap, or a previous call's reduction.	None
accelerator	Bvh | None	A prebuilt wp.Bvh over points, to reuse across queries. Spelled the way the four query_ball / query_nearest entry points spell it, but with no backend beside it: this query has only the BVH broad phase, so there is nothing to select between.	None
leaf_size	int	Maximum primitives per BVH leaf when one is built here.	4

Returns:

Type	Description
(index, weighted_distance)	(m,) winning site per query and its |p - q| - w(p), which is negative wherever a query lies inside a site's radius. A query with no site at all (an empty cloud) reports -1 and inf.

Raises:

Type	Description
ValueError	If weights does not have one entry per point.
RuntimeError	If points, weights, queries and accelerator are not all on one device.

Notes

max_weight is trusted, not checked: a bound smaller than some weight can silently prune the true winner, and verifying it would cost the very reduction the parameter exists to avoid. The default is therefore the safe one.

Exactness is the same argument the k-NN queries use, with the weight folded in: a site outside the cube of half-extent r is farther than r, so it scores worse than r - max_weight, and a best score at or below that bound cannot be beaten. Unlike a plain nearest query this means the search radius must exceed the answer's distance by the weight range, so a wide weight distribution costs more scans than a narrow one.

See Also

query_nearest The unweighted query, and the k > 1 form. This one answers k = 1 only, because no caller has needed more.

Source code in triwarp/neighbors.py

def query_weighted_nearest(
    points: wp.array[wp.vec3],
    weights: wp.array[wp.float32],
    queries: wp.array[wp.vec3],
    *,
    max_weight: float | None = None,
    accelerator: wp.Bvh | None = None,
    leaf_size: int = 4,
) -> tuple[wp.array[wp.int32], wp.array[wp.float32]]:
    """
    Nearest site under the weighted distance ``|p - q| - w(p)``.

    Each site carries a radius, and the winner is the one whose *surface* is closest rather than
    whose centre is -- the additively weighted (Apollonius) nearest-neighbour query, which is what
    picks the influencing site when the sites have different scales: a sphere set, a level-of-detail
    cluster, a set of samples with per-sample confidence. Plain
    [`query_nearest`][triwarp.neighbors.query_nearest] is the ``w = 0`` case, and the two
    genuinely differ -- on a random 40-site cloud, 1 query in 6 had a different winner.

    Parameters
    ----------
    points
        ``(n,)`` site positions as ``wp.vec3``.
    weights
        ``(n,)`` per-site weights, subtracted from the distance. Larger wins ties of distance; may
        be negative, which pushes a site away.
    queries
        ``(m,)`` query positions as ``wp.vec3``.
    max_weight
        An **upper bound** on ``weights``, which is what makes the search prunable. ``None`` reduces
        ``weights`` on the device and reads the maximum back (one readback), so pass it when the
        bound is already known -- a radius cap, or a previous call's reduction.
    accelerator
        A prebuilt ``wp.Bvh`` over ``points``, to reuse across queries. Spelled the way the four
        [`query_ball`][triwarp.neighbors.query_ball] /
        [`query_nearest`][triwarp.neighbors.query_nearest] entry points spell it, but with no
        ``backend`` beside it: this query has only the BVH broad phase, so there is nothing to
        select between.
    leaf_size
        Maximum primitives per BVH leaf when one is built here.

    Returns
    -------
    index, weighted_distance
        ``(m,)`` winning site per query and its ``|p - q| - w(p)``, which is **negative** wherever a
        query lies inside a site's radius. A query with no site at all (an empty cloud) reports
        ``-1`` and ``inf``.

    Raises
    ------
    ValueError
        If ``weights`` does not have one entry per point.
    RuntimeError
        If ``points``, ``weights``, ``queries`` and ``accelerator`` are not all on one device.

    Notes
    -----
    ``max_weight`` is trusted, not checked: a bound *smaller* than some weight can silently prune
    the true winner, and verifying it would cost the very reduction the parameter exists to avoid.
    The default is therefore the safe one.

    Exactness is the same argument the k-NN queries use, with the weight folded in: a site outside
    the cube of half-extent ``r`` is farther than ``r``, so it scores worse than ``r - max_weight``,
    and a best score at or below that bound cannot be beaten. Unlike a plain nearest query this
    means the search radius must exceed the answer's distance *by the weight range*, so a wide
    weight distribution costs more scans than a narrow one.

    See Also
    --------
    [`query_nearest`][triwarp.neighbors.query_nearest]
        The unweighted query, and the ``k > 1`` form. This one answers ``k = 1`` only, because no
        caller has needed more.
    """
    require_same_device(points=points, weights=weights, queries=queries, accelerator=accelerator)
    device = points.device
    n = int(points.shape[0])
    m = int(queries.shape[0])
    if int(weights.shape[0]) != n:
        raise ValueError("weights must have one entry per point")

    if m == 0:
        return (
            wp.empty(0, dtype=wp.int32, device=device),
            wp.empty(0, dtype=wp.float32, device=device),
        )
    if n == 0:
        return (
            wp.full(m, -1, dtype=wp.int32, device=device),
            wp.full(m, float("inf"), dtype=wp.float32, device=device),
        )

    if accelerator is None:
        accelerator = bvh_from_points(points, leaf_size=leaf_size)
    if max_weight is None:
        max_weight = float(cast(float, tw.reduce.max(weights)))
    min_bound, max_bound = tw.bounds.aabb(points)
    # First radius: the mean spacing's own estimate plus the weight bound, since a query cannot be
    # certified below it however close its winner is.
    initial_radius = knn_initial_radius(points, 1, bounds=(min_bound, max_bound)) + max(
        max_weight, 0.0
    )

    out_indices = wp.empty(m, dtype=wp.int32, device=device)
    out_distances = wp.empty(m, dtype=wp.float32, device=device)
    wp.launch(
        kernel_neighbors.query_weighted_nearest_neighbors,
        dim=m,
        inputs=[
            points,
            weights,
            queries,
            accelerator.id,
            wp.float32(max_weight),
            wp.float32(initial_radius),
            min_bound,
            max_bound,
            out_indices,
            out_distances,
        ],
        device=device,
    )
    return out_indices, out_distances

nearest_neighbor_distance

nearest_neighbor_distance(
    points: array[vec3],
) -> wp.array[wp.float32]

Distance from each point to the closest other point of the same cloud.

The standard scale estimate for an unstructured cloud, and what a sampling-driven parameter is normally derived from: the mean of this array is the cloud's mean spacing, which is how ball_pivoting guesses its ball radius and how screened_poisson caps its octree depth. A self-query for two neighbours, keeping the second — the first is the point itself, at distance zero.

Parameters:

Name	Type	Description	Default
points	array[vec3]	(n,) point positions as wp.array[wp.vec3].	required

Returns:

Type	Description
array[float32]	Length-n distances on points.device. Zero where two points coincide exactly.

Notes

A cloud of fewer than two points has no answer, and this reports inf for every entry — the value query_nearest already uses for a slot it could not fill. Open3D's compute_nearest_neighbor_distance reports 0.0 in that one case; on any cloud of two or more points the two agree, because every point then has a nearest neighbour. The distances are float32 (wp.length), so they can differ from a float64 reference in the last digits.

See Also

query_nearest The underlying query; call it directly to reuse a BVH across several queries of one cloud. triwarp.points.farthest_point_sample

Source code in triwarp/neighbors.py

def nearest_neighbor_distance(points: wp.array[wp.vec3]) -> wp.array[wp.float32]:
    """
    Distance from each point to the closest *other* point of the same cloud.

    The standard scale estimate for an unstructured cloud, and what a sampling-driven parameter is
    normally derived from: the mean of this array is the cloud's mean spacing, which is how
    [`ball_pivoting`][triwarp.reconstruction.ball_pivoting] guesses its ball radius and how
    [`screened_poisson`][triwarp.reconstruction.screened_poisson] caps its octree depth. A
    self-query for two neighbours, keeping the second — the first is the point itself, at distance
    zero.

    Parameters
    ----------
    points
        ``(n,)`` point positions as ``wp.array[wp.vec3]``.

    Returns
    -------
    wp.array[wp.float32]
        Length-``n`` distances on ``points.device``. Zero where two points coincide exactly.

    Notes
    -----
    A cloud of fewer than two points has no answer, and this reports ``inf`` for every entry —
    the value [`query_nearest`][triwarp.neighbors.query_nearest] already uses for a slot it
    could not fill. Open3D's ``compute_nearest_neighbor_distance`` reports ``0.0`` in that one
    case; on any cloud of two or more points the two agree, because every point then has a nearest
    neighbour. The distances are ``float32`` (``wp.length``), so they can differ from a
    ``float64`` reference in the last digits.

    See Also
    --------
    [`query_nearest`][triwarp.neighbors.query_nearest]
        The underlying query; call it directly to reuse a BVH across several queries of one cloud.
    [`triwarp.points.farthest_point_sample`][triwarp.points.farthest_point_sample]
    """
    device = points.device
    n = int(points.shape[0])
    if n < 2:
        return wp.full(n, math.inf, dtype=wp.float32, device=device)

    _indices, distances = query_nearest(points, points, k=2, backend="bvh")
    # Column 1 of the ``(n, 2)`` table, which is a *strided* view -- so it is copied into a dense
    # buffer rather than returned, both because callers expect a plain ``wp.array`` and because a
    # strided array is the shape that silently corrupts a downstream Python-scope gather.
    nearest = wp.empty(n, dtype=wp.float32, device=device)
    wp.copy(nearest, distances[:, 1])
    return nearest

closest_pair

closest_pair(points: array[vec3]) -> tuple[int, int, float]

Find the two closest points of a cloud, and the distance between them.

The global minimum of nearest_neighbor_distance, with the partner recovered -- so it is the same k=2 self-query, reduced instead of returned. Answers "does this cloud contain a near-duplicate, and where" in one call, which is the question a tolerance for triwarp.points.point_duplicate_mask is normally chosen from.

Parameters:

Name	Type	Description	Default
points	array[vec3]	(n,) point positions as wp.array[wp.vec3], n >= 2.	required

Returns:

Type	Description
(index_a, index_b, distance)	The two point indices — index_a < n, index_b its nearest neighbour — and their Euclidean distance, all as Python scalars.

Raises:

Type	Description
ValueError	If points holds fewer than two points, which have no pair.

Notes

Ties are broken toward the lowest index_a: the reduction is a min over one int64 per point holding the distance in the high half and the index in the low, so a shorter distance wins and an equal distance defers to the smaller index. Exact duplicates therefore report the first duplicated point at distance 0.0.

Two host readbacks, both unavoidable and both single-element: the reduction's result, and index_b from the query table. Everything up to them stays on the device, so this does not move the (n, 2) table across the bus.

See Also

nearest_neighbor_distance The per-point form, when every distance is wanted rather than the smallest. triwarp.points.point_duplicate_mask Exact coincidence rather than proximity, and a mask rather than one pair.

Source code in triwarp/neighbors.py

def closest_pair(points: wp.array[wp.vec3]) -> tuple[int, int, float]:
    """
    Find the two closest points of a cloud, and the distance between them.

    The global minimum of
    [`nearest_neighbor_distance`][triwarp.neighbors.nearest_neighbor_distance], with the partner
    recovered -- so it is the same ``k=2`` self-query, reduced instead of returned.
    Answers "does this cloud contain a near-duplicate, and where" in one call, which is the question
    a tolerance for
    [`triwarp.points.point_duplicate_mask`][triwarp.points.point_duplicate_mask] is normally chosen
    from.

    Parameters
    ----------
    points
        ``(n,)`` point positions as ``wp.array[wp.vec3]``, ``n >= 2``.

    Returns
    -------
    index_a, index_b, distance
        The two point indices — ``index_a < n``, ``index_b`` its nearest neighbour — and their
        Euclidean distance, all as Python scalars.

    Raises
    ------
    ValueError
        If ``points`` holds fewer than two points, which have no pair.

    Notes
    -----
    Ties are broken toward the **lowest** ``index_a``: the reduction is a ``min`` over one ``int64``
    per point holding the distance in the high half and the index in the low, so a shorter distance
    wins and an equal distance defers to the smaller index. Exact duplicates therefore report the
    first duplicated point at distance ``0.0``.

    Two host readbacks, both unavoidable and both single-element: the reduction's result, and
    ``index_b`` from the query table. Everything up to them stays on the device, so this does not
    move the ``(n, 2)`` table across the bus.

    See Also
    --------
    [`nearest_neighbor_distance`][triwarp.neighbors.nearest_neighbor_distance]
        The per-point form, when every distance is wanted rather than the smallest.
    [`triwarp.points.point_duplicate_mask`][triwarp.points.point_duplicate_mask]
        Exact coincidence rather than proximity, and a mask rather than one pair.
    """
    n = int(points.shape[0])
    if n < 2:
        raise ValueError("closest_pair needs at least two points")

    indices, distances = query_nearest(points, points, k=2, backend="bvh")
    keys = wp.empty(n, dtype=wp.int64, device=points.device)
    wp.launch(
        kernel_points.nearest_pair_keys, dim=n, inputs=[distances, keys], device=points.device
    )

    key = int(cast(int, tw.reduce.min(keys)))
    index_a = key & 0xFFFFFFFF
    # The high half is the distance's own float32 bits, so it decodes on the host for free rather
    # than costing a third readback.
    distance = float(np.array([key >> 32], dtype=np.uint32).view(np.float32)[0])
    index_b = int(read_scalar(indices.flatten(), 2 * index_a + 1))
    return index_a, index_b, distance

geodesic_ball

geodesic_ball(
    vertices: array[vec3],
    faces: array[int32],
    radius: float,
    min_count: int = 6,
) -> tuple[
    wp.array[wp.int32],
    wp.array[wp.int32],
    wp.array[wp.int32],
]

Per-vertex geodesic-ball neighborhoods, and the reference neighbor that frames each one.

For each vertex this is a breadth-first traversal of the mesh edge graph, enqueueing a neighbor only when it lies within Euclidean radius of the center — i.e. the connected component of the center within the radius ball. This is a geodesic ball rather than a pure Euclidean one, so it excludes vertices that are spatially close but lie across a fold of the surface (e.g. the opposite wall of a torus tube), which a Euclidean hash-grid query would wrongly include and which corrupts the quadric fit. When fewer than min_count vertices are reachable, the nearest out-of-ball vertices are appended (libigl's extra_candidates path).

Also returns the per-vertex reference neighbor used to build the tangent frame: the lowest-indexed edge neighbor, matching libigl's adjacency_list[i][0]. libigl's symmetrized shape operator is frame-dependent, so reproducing its principal values ( principal_curvature with frame_independent=False) requires this exact frame; the default frame-independent computation does not depend on it. Isolated vertices reference themselves.

The traversal runs entirely on device. Vertex adjacency is built as a CSR graph via edges_unique + edges_to_csr, then a single-pass BFS collects each ball into its per-source queue row (the queue prefix is the result) and a scan + gather compacts the rows into the CSR neighbor buffer. Each source uses fixed-capacity scratch of _PER_SOURCE_MAX_NEIGHBORS neighbors (triwarp.kernels.algorithms.bfs, currently 512); if a vertex collects more than that the surplus is dropped and a warning is emitted.

Note

Distances and tie-breaking are computed in float32 (set-equivalent to the libigl reference; borderline ties between equidistant neighbors may resolve differently but leave the order-independent quadric fit unchanged).

Parameters:

Name	Type	Description	Default
vertices	array[vec3]	(n_vertices,) mesh vertex positions as wp.vec3.	required
faces	array[int32]	Length-3 * n_faces flat triangle index buffer as wp.int32.	required
radius	float	Geodesic-ball radius in world units.	required
min_count	int	Minimum neighbors per vertex; the nearest out-of-ball vertices backfill any shortfall.	6

Returns:

Type	Description
tuple[array[int32], array[int32], array[int32]]	(neighbor_indices, offsets, reference_neighbors). offsets is the length-n_vertices + 1 exclusive prefix sum of per-vertex neighbor counts (CSR row bounds); vertex i owns neighbor_indices[offsets[i] : offsets[i + 1]] and offsets[n_vertices] is the total. reference_neighbors has length n_vertices.

Raises:

Type	Description
RuntimeError	If vertices and faces are not all on one device.

Source code in triwarp/neighbors.py

def geodesic_ball(
    vertices: wp.array[wp.vec3], faces: wp.array[wp.int32], radius: float, min_count: int = 6
) -> tuple[wp.array[wp.int32], wp.array[wp.int32], wp.array[wp.int32]]:
    """
    Per-vertex geodesic-ball neighborhoods, and the reference neighbor that frames each one.

    For each vertex this is a breadth-first traversal of the mesh edge graph, enqueueing a neighbor
    only when it lies within Euclidean ``radius`` of the center — i.e. the connected component of
    the center within the radius ball. This is a *geodesic* ball rather than a pure Euclidean one,
    so it excludes vertices that are spatially close but lie across a fold of the surface (e.g. the
    opposite wall of a torus tube), which a Euclidean hash-grid query would wrongly include and
    which corrupts the quadric fit. When fewer than ``min_count`` vertices are reachable, the
    nearest out-of-ball vertices are appended (libigl's ``extra_candidates`` path).

    Also returns the per-vertex reference neighbor used to build the tangent frame: the
    lowest-indexed edge neighbor, matching libigl's ``adjacency_list[i][0]``. libigl's symmetrized
    shape operator is frame-dependent, so reproducing its principal values (
    [`principal_curvature`][triwarp.curvature.principal_curvature] with ``frame_independent=False``)
    requires this exact frame; the default frame-independent computation does not depend on it.
    Isolated vertices reference themselves.

    The traversal runs entirely on device. Vertex adjacency is built as a CSR graph via
    [`edges_unique`][triwarp.edges.edges_unique] +
    [`edges_to_csr`][triwarp.graph.edges_to_csr], then a single-pass BFS collects each ball into
    its per-source queue row (the queue prefix *is* the result) and a scan + gather compacts the
    rows into the CSR neighbor buffer. Each source uses fixed-capacity scratch of
    ``_PER_SOURCE_MAX_NEIGHBORS`` neighbors (``triwarp.kernels.algorithms.bfs``, currently 512);
    if a vertex collects more than that the surplus is dropped and a warning is emitted.

    !!! note

        Distances and tie-breaking are computed in ``float32`` (set-equivalent to the libigl
        reference; borderline ties between equidistant neighbors may resolve differently but leave
        the order-independent quadric fit unchanged).

    Parameters
    ----------
    vertices
        ``(n_vertices,)`` mesh vertex positions as ``wp.vec3``.
    faces
        Length-``3 * n_faces`` flat triangle index buffer as ``wp.int32``.
    radius
        Geodesic-ball radius in world units.
    min_count
        Minimum neighbors per vertex; the nearest out-of-ball vertices backfill any shortfall.

    Returns
    -------
    tuple[wp.array[wp.int32], wp.array[wp.int32], wp.array[wp.int32]]
        ``(neighbor_indices, offsets, reference_neighbors)``. ``offsets`` is the
        length-``n_vertices + 1`` exclusive prefix sum of per-vertex neighbor counts (CSR row
        bounds); vertex ``i`` owns ``neighbor_indices[offsets[i] : offsets[i + 1]]`` and
        ``offsets[n_vertices]`` is the total. ``reference_neighbors`` has length ``n_vertices``.

    Raises
    ------
    RuntimeError
        If ``vertices`` and ``faces`` are not all on one device.
    """
    require_same_device(vertices=vertices, faces=faces)
    device = vertices.device
    n = int(vertices.shape[0])
    if n == 0:
        # A single zero rather than an empty buffer: the CSR row-bounds form is ``n + 1`` long.
        return (
            wp.empty(0, dtype=wp.int32, device=device),
            wp.zeros(1, dtype=wp.int32, device=device),
            wp.empty(0, dtype=wp.int32, device=device),
        )

    unique_edges, _ = tw.edges.edges_unique(faces, n_vertices=n)
    adjacency = tw.graph.edges_to_csr(n, unique_edges)
    adj_offsets = adjacency.offsets
    adj_columns = adjacency.columns

    reference_neighbors = wp.empty(n, dtype=wp.int32, device=device)
    wp.launch(
        kernel_neighbors.geodesic_ball_reference_neighbors,
        dim=n,
        inputs=[adj_offsets, adj_columns, reference_neighbors],
        device=device,
    )

    # Per-source scratch lives in shared global-memory pools sized for one chunk of sources
    # (queue rows, an open-addressing visited row pre-filled with -1 per launch, and a small
    # nearest-fallback pool) instead of ~8 KB of per-thread local arrays.
    chunk = min(n, 1 << 15)
    queue_pool = wp.empty(
        (chunk, kernel_bfs._PER_SOURCE_MAX_NEIGHBORS), dtype=wp.int32, device=device
    )
    visited_pool = wp.empty(
        (chunk, kernel_bfs._VISITED_HASH_CAPACITY), dtype=wp.int32, device=device
    )
    ext_dist_pool = twt.empty_2d((chunk, kernel_bfs._EXTRAS_CAPACITY), wp.float32, device=device)
    ext_idx_pool = twt.empty_2d((chunk, kernel_bfs._EXTRAS_CAPACITY), wp.int32, device=device)

    overflow = wp.zeros(1, dtype=wp.int32, device=device)
    counts = wp.empty(n, dtype=wp.int32, device=device)
    local_offsets = wp.empty(chunk, dtype=wp.int32, device=device)
    chunk_total_buf = wp.empty(1, dtype=wp.int32, device=device)
    chunk_flats: list[wp.array[wp.int32]] = []
    for start in range(0, n, chunk):
        m = min(chunk, n - start)
        visited_pool.fill_(-1)
        wp.launch(
            kernel_neighbors.query_geodesic_ball_collect,
            dim=m,
            inputs=[
                vertices,
                adj_offsets,
                adj_columns,
                wp.float32(radius),
                wp.int32(min_count),
                wp.int32(start),
                queue_pool,
                visited_pool,
                ext_dist_pool,
                ext_idx_pool,
                counts,
                overflow,
            ],
            device=device,
        )
        # Gather this chunk's queue rows before the next chunk reuses the pools: chunk-local
        # exclusive scan of counts, one 4-byte readback for the chunk total, then a coalesced
        # 2D copy into the chunk's flat buffer.
        wp.utils.array_scan(counts[start : start + m], out_array=local_offsets[:m], inclusive=False)
        wp.map(
            wp.add, local_offsets[m - 1 : m], counts[start + m - 1 : start + m], out=chunk_total_buf
        )
        chunk_total = int(read_scalar(chunk_total_buf, 0))
        flat_chunk = wp.empty(chunk_total, dtype=wp.int32, device=device)
        if chunk_total > 0:
            wp.launch(
                kernel_neighbors.gather_queue_rows,
                dim=(m, kernel_bfs._PER_SOURCE_MAX_NEIGHBORS),
                inputs=[queue_pool, counts, local_offsets, wp.int32(start), flat_chunk],
                device=device,
            )
        chunk_flats.append(flat_chunk)

    n_overflow = int(read_scalar(overflow, 0))
    if n_overflow > 0:
        warnings.warn(
            f"geodesic_ball: {n_overflow} neighborhood capacity breaches "
            f"(fixed cap {kernel_bfs._PER_SOURCE_MAX_NEIGHBORS}); surplus neighbors dropped.",
            stacklevel=2,
        )

    # CSR row bounds in the length-``n + 1`` form. The leading zero from ``wp.zeros`` is the first
    # exclusive offset and the inclusive scan fills the rest, so ``offsets[n]`` holds the total --
    # ``array.counts_to_offsets``' convention, open-coded here because that helper reads the total
    # back to the host and no caller here wants it. The terminator belongs at this producer rather
    # than in each consumer's kernel: both of them read ``offsets[i + 1]`` as the row end.
    offsets = wp.zeros(n + 1, dtype=wp.int32, device=device)
    wp.utils.array_scan(counts, out_array=offsets[1:], inclusive=True)

    if len(chunk_flats) == 1:
        # Single chunk (n <= chunk): the chunk buffer already is the global CSR neighbor buffer.
        return chunk_flats[0], offsets, reference_neighbors

    # Chunk order equals ascending source order, so concatenation lines up with the global scan.
    # Same per-segment ``wp.copy`` loop either way -- that is the packing floor -- one call for it.
    return tw.array.concatenate(chunk_flats), offsets, reference_neighbors