neptune¶
Neptune, a point transformer for neutrino events.
Paper: https://arxiv.org/abs/2510.01733. Reference implementation: https://github.com/felixyu7/neptune.
- class graphnet.models.transformer.neptune.Neptune(*args, **kwargs)[source]¶
Bases:
GNNNeptune: a point transformer for neutrino events.
Tokenizes each event by farthest point sampling (FPS) over its pulses, encodes the tokens with a rotary transformer and pools them to one vector. Use with EdgelessGraph and NodesAsPulses.
Geometric defaults assume IceCube86 standardization (one unit is 0.5 km and 30 μs); rescale position_encoding_schema, rope_scales and the tokenizer’s metric_time_scale for other detectors. The charge column must hold log10 charge, as IceCube86 standardizes it. For faster FPS, pip install “torch-fps==0.6.*”.
Construct Neptune.
- Parameters:
input_feature_names (
List[str]) – Names of the columns of data.x.coordinate_columns (
List[str]) – Names of the (x, y, z) columns.time_column (
str) – Name of the time column.charge_column (
Optional[str], default:None) – Name of the charge column. None for detectors without charge (e.g. Prometheus): every pulse then has unit charge.feature_columns (
Optional[List[str]], default:None) – Further per-pulse features for the tokenizer. Defaults to every column not named above.num_patches (
int, default:128) – Maximum number of tokens per event.d_model (
int, default:768) – Token dimension.depth (
int, default:12) – Number of transformer layers.num_heads (
int, default:12) – Number of attention heads; d_model / num_heads must be even and at least 8 for the 4D rotary embedding.hidden_dim (
int, default:2048) – Inner dimension of the SwiGLU feed-forward network.dropout (
float, default:0.1) – Dropout used throughout the model.drop_path_rate (
float, default:0.0) – Stochastic-depth rate of the last layer, ramped linearly from zero.output_dim (
Optional[int], default:None) – Readout width (nb_outputs); defaults to d_model.pool_type (
str, default:'mean') – “mean” (masked mean) or “attention” (learned query) pooling over tokens.tokenizer_mlp_layers (
Optional[List[int]], default:None) – Hidden sizes of the tokenizer’s per-pulse MLP; defaults to [256, 512, 768].tokenizer_kwargs (
Optional[Dict[str,Any]], default:None) – Further FPSTokenizer arguments, e.g. assign_mode, lloyd_iters or metric_time_scale.layerscale_init (
float, default:1e-05) – Initial value of the LayerScale parameters.position_encoding_schema (
Optional[Dict[int,Union[float,Tuple[float,float]]]], default:None) – FourierEncoder schema over the centroid columns (x, y, z, t), as {column: (highest frequency, span)}. Time is encoded relatively by RoPE; add column 3 to also encode it absolutely.position_encoding_bands (
int, default:20) – Frequencies per encoded column.rope_scales (
Optional[List[float]], default:None) – Highest rotary frequency per axis (x, y, z, t).rope_base (
int, default:60) – Ratio of highest to lowest rotary frequency per axis.compile_encoder (
bool, default:True) – torch.compile the transformer layers, 2-4x faster on GPU, where it also enables block-diagonal packed attention over only the valid tokens. Turn off for CPU-only or debugging runs.args (Any)
kwargs (Any)
- Return type:
object