neuralop.layers.neighbor_search.NeighborSearch

class neuralop.layers.neighbor_search.NeighborSearch(use_open3d=True, return_norm=False)[source]

Neighborhood search between two arbitrary coordinate meshes.

For each point x in queries, returns a set of the indices of all points y in data within the ball of radius r B_r(x)

Parameters:

use_open3dbool, optional

Whether to use open3d or native PyTorch implementation, by default True NOTE: open3d implementation requires 3d data

return_normbool, optional

Whether to return normalized distances, by default False

Methods

forward(data, queries, radius)

Find the neighbors, in data, of each point in queries within a ball of radius.

forward(data, queries, radius)[source]

Find the neighbors, in data, of each point in queries within a ball of radius. Returns in CRS format.

Parameters:

datatorch.Tensor of shape [n, d]

Search space of possible neighbors NOTE: open3d requires d=3

queriestorch.Tensor of shape [m, d]

Points for which to find neighbors NOTE: open3d requires d=3

radiusfloat

Radius of each ball: B(queries[j], radius)

Returns:

return_dictdict

Dictionary with keys: neighbors_index, neighbors_row_splits

neighbors_index: torch.Tensor with dtype=torch.int64

Index of each neighbor in data for every point in queries. Neighbors are ordered in the same orderings as the points in queries. Open3d and torch_cluster implementations can differ by a permutation of the neighbors for every point.

neighbors_row_splits: torch.Tensor of shape [m+1] with dtype=torch.int64

The value at index j is the sum of the number of neighbors up to query point j-1. First element is 0 and last element is the total number of neighbors.