tokenizers

Learnable tokenizers turning point clouds into transformer tokens.

class graphnet.models.components.tokenizers.FPSTokenizer(feature_dim, max_tokens, token_dim, mlp_layers, k_neighbors, dropout, knn_pool, rel_pos_hidden, charge_weighted_mean, charge_col, assign_mode, metric_time_scale, lloyd_iters)[source]

Bases: LightningModule

Tokenize a batch of point clouds by farthest point sampling.

Ported from the Neptune reference implementation (https://github.com/felixyu7/neptune, MIT licence).

A per-point MLP lifts the features. Events with at most max_tokens points keep one token per point; larger events are reduced to max_tokens centroids by FPS in (x, y, z, t) and pooled per token. Four summary scalars (multiplicity, total charge, time spread, RMS radius) are appended and a second MLP refines every token.

Construct FPSTokenizer.

Parameters:

  • feature_dim (int) – Number of input features per point.

  • max_tokens (int, default: 128) – Maximum number of tokens per event.

  • token_dim (int, default: 768) – Output token dimension.

  • mlp_layers (Optional[List[int]], default: None) – Hidden sizes of the per-point MLP; defaults to [256, 512, 768].

  • k_neighbors (int, default: 8) – Neighbours per centroid when assign_mode=”knn”.

  • dropout (float, default: 0.0) – Dropout inside both MLPs.

  • knn_pool (str, default: 'max') – “max”, or “max_mean” to also concatenate a weighted mean.

  • rel_pos_hidden (int, default: 64) – Hidden size of the relative-offset encoder.

  • charge_weighted_mean (bool, default: True) – Weight means and spreads by log1p charge.

  • charge_col (int, default: 0) – Feature column holding log1p charge.

  • assign_mode (str, default: 'voronoi') – “voronoi” pools each point into its nearest centroid; “knn” pools the k_neighbors nearest points of each centroid.

  • metric_time_scale (float, default: 18.0) – Weight of time relative to space in the FPS and assignment metric only. The default, 18, is 0.3 km/μs in IceCube86 units.

  • lloyd_iters (int, default: 0) – Charge-weighted Lloyd (k-means) steps refining the Voronoi centroids, which then become cell means; cells can end up empty. Ignored for “knn”.

N_EXTRA = 4
metric_scale: Tensor
forward(coords, features, batch_ids, times, batch_size)[source]

Tokenize a ragged batch of point clouds.

Parameters:

  • coords (Tensor) – [N, 3] positions (x, y, z).

  • features (Tensor) – [N, F] per-point features.

  • batch_ids (Tensor) – [N] event index of each point.

  • times (Tensor) – [N, 1] time of each point.

  • batch_size (Optional[int], default: None) – Number of events, so that trailing empty events keep their row; inferred from batch_ids if None.

Return type:

Tuple[Tensor, Tensor, Tensor]

Returns:

(tokens [B, max_tokens, token_dim], centroids [B, max_tokens, 4], masks [B, max_tokens]).