fps

Farthest point sampling (FPS) for point-cloud tokenization.

Ported from the Neptune reference implementation (https://github.com/felixyu7/neptune, MIT licence). Calls go to the optional torch_fps package (https://github.com/felixyu7/torch-fps; Triton kernels on CUDA, C++ on CPU) when it is installed; otherwise the pure-PyTorch implementation below is used.

Shared semantics: a point is valid if its mask is True and its coordinates are finite; distances accumulate in float32; FPS and kNN ties select the lowest index, nearest-centroid assignment ties may pick either centroid; once a row runs out of valid points the last selection repeats; kNN rows with fewer than k_neighbors valid points pad with the centroid index. The backends agree up to float32 rounding, which can flip exact ties.

graphnet.models.components.fps.farthest_point_sampling(points, valid_mask, k, *, start_idx, random_start, generator, validate, assume_finite)[source]

Select k maximally spread-out points per batch row.

Parameters:

  • points (Tensor) – Float [B, N, D] coordinates.

  • valid_mask (Tensor) – Bool [B, N]; False marks padded or invalid points.

  • k (int) – Number of points to select per row; at most the valid count.

  • start_idx (Optional[Tensor], default: None) – Optional [B] first index per row, on points.device.

  • random_start (bool, default: True) – Draw a random valid start when start_idx is None; otherwise start at the first valid point.

  • generator (Optional[Generator], default: None) – Optional generator for the random start.

  • validate (bool, default: True) – Check k <= valid count and start_idx (one host sync).

  • assume_finite (bool, default: False) – Caller guarantees finite coordinates for valid points.

Return type:

Tensor

Returns:

Long [B, k] selected indices.

graphnet.models.components.fps.farthest_point_sampling_with_knn(points, valid_mask, k, k_neighbors, *, start_idx, random_start, generator, validate, assume_finite)[source]

Run FPS and gather each centroid’s k_neighbors nearest points.

Arguments are as in farthest_point_sampling(), plus k_neighbors (0 < k_neighbors <= N).

Return type:

Tuple[Tensor, Tensor]

Returns:

[B, k] centroid indices and [B, k, k_neighbors] neighbour indices, closest first.

Parameters:

  • points (Tensor)

  • valid_mask (Tensor)

  • k (int)

  • k_neighbors (int)

  • start_idx (Tensor | None)

  • random_start (bool)

  • generator (Generator | None)

  • validate (bool)

  • assume_finite (bool)

graphnet.models.components.fps.farthest_point_sampling_with_assign(points, valid_mask, k, *, start_idx, random_start, generator, validate, assume_finite)[source]

Run FPS and assign every point to its nearest centroid.

Arguments are as in farthest_point_sampling().

Return type:

Tuple[Tensor, Tensor]

Returns:

[B, k] centroid indices and [B, N] assignments in [0, k) (unspecified at invalid points).

Parameters:

  • points (Tensor)

  • valid_mask (Tensor)

  • k (int)

  • start_idx (Tensor | None)

  • random_start (bool)

  • generator (Generator | None)

  • validate (bool)

  • assume_finite (bool)

graphnet.models.components.fps.nearest_assign(p_flat, cents, starts, counts, n_max, k)[source]

Assign batch-segmented flat points to their nearest centroid.

Parameters:

  • p_flat (Tensor) – Float32 [N, D] points, grouped by event.

  • cents (Tensor) – Float32 [B * k, D] centroids.

  • starts (Tensor) – Int64 [B] offset of each event in p_flat.

  • counts (Tensor) – Int64 [B] number of points per event.

  • n_max (int) – Largest value in counts.

  • k (int) – Number of centroids per event.

Return type:

Tensor

Returns:

Int64 [N] indices in [0, k).