noether.modeling.functional.geometric

Attributes

HAS_PYG

HAS_TRITON

Functions

radius_triton(x, y, r[, batch_x, batch_y, ...])

Find all edges where points in x are within radius r of points in y.

radius_pytorch(x, y, r[, batch_x, batch_y, ...])

Fallback implementation of radius using pure PyTorch operations.

knn_triton(x, y, k[, batch_x, batch_y])

Calculates k-nearest neighbors using a Triton kernel.

knn_pytorch(x, y, k[, batch_x, batch_y, cosine])

Fallback implementation of knn using pure PyTorch operations.

radius(x, y, r, max_num_neighbors[, batch_x, batch_y])

Find all points within radius r.

knn(x, y, k[, batch_x, batch_y])

Module Contents

noether.modeling.functional.geometric.HAS_PYG
noether.modeling.functional.geometric.HAS_TRITON = True
noether.modeling.functional.geometric.radius_triton(x, y, r, batch_x=None, batch_y=None, max_num_neighbors=None)

Find all edges where points in x are within radius r of points in y.

Parameters:

  • x (torch.Tensor) – Source points, shape (N, D) or flattened (N*D,) with D inferred

  • y (torch.Tensor) – Target points, shape (M, D) or flattened (M*D,)

  • r (float) – Radius threshold

  • max_num_neighbors (int | None) – Maximum neighbors per target point (default: N)

  • batch_x (torch.Tensor | None) – Batch indices for x points, shape (N,)

  • batch_y (torch.Tensor | None) – Batch indices for y points, shape (M,)

Returns:

  • row 0: source indices (from x)

  • row 1: target indices (from y)

Return type:

Edge index tensor of shape (2, E) where

noether.modeling.functional.geometric.radius_pytorch(x, y, r, batch_x=None, batch_y=None, max_num_neighbors=None)

Fallback implementation of radius using pure PyTorch operations.

Parameters:

  • x (torch.Tensor) – Source points (N, D).

  • y (torch.Tensor) – Query points (M, D).

  • r (float) – Radius to search for.

  • max_num_neighbors (int | None) – Maximum number of neighbors to return.

  • batch_x (torch.Tensor | None) – Batch index for source points.

  • batch_y (torch.Tensor | None) – Batch index for query points.

Returns:

Edge index (2, num_edges).

Return type:

torch.Tensor

noether.modeling.functional.geometric.knn_triton(x, y, k, batch_x=None, batch_y=None)

Calculates k-nearest neighbors using a Triton kernel.

Parameters:

  • x (torch.Tensor) – Reference points (N, D)

  • y (torch.Tensor) – Query points (M, D)

  • k (int) – Number of neighbors

  • batch_x (torch.Tensor | None) – Optional batch indices for x (N,)

  • batch_y (torch.Tensor | None) – Optional batch indices for y (M,)

  • cosine – If True, uses Cosine distance. False = Euclidean.

Returns:

Indices of the k nearest neighbors for each y (M, K)

Return type:

torch.Tensor

noether.modeling.functional.geometric.knn_pytorch(x, y, k, batch_x=None, batch_y=None, cosine=False)

Fallback implementation of knn using pure PyTorch operations.

Parameters:
Returns:

Edge index (2, num_edges).

Return type:

torch.Tensor

noether.modeling.functional.geometric.radius(x, y, r, max_num_neighbors, batch_x=None, batch_y=None)

Find all points within radius r.

Parameters:

  • x (torch.Tensor) – Source points (N, D).

  • y (torch.Tensor) – Query points (M, D).

  • r (float) – Radius to search for.

  • max_num_neighbors (int) – Maximum number of neighbors to return per query.

  • batch_x (torch.Tensor | None) – Batch index for source points.

  • batch_y (torch.Tensor | None) – Batch index for query points.

Returns:

Edge index (2, num_edges). first row: source indices (from y) second row: target indices (from x)

Return type:

torch.Tensor

noether.modeling.functional.geometric.knn(x, y, k, batch_x=None, batch_y=None)
Parameters:
Return type:

torch.Tensor