jFoF: GPU Friends-of-Friends halo finding with gradient propagation

ABSTRACT We present jFoF, a fully graphics processing unit (GPU) native Friends-of-Friends (FoF) halo finder designed for both high-performance simulation analysis and differentiable modelling. Implemented in jax, jFoF achieves end-to-end acceleration by performing all neighbour searches, label propagation, and group construction directly on GPUs, eliminating costly host-device transfers. We introduce two complementary neighbour-search strategies, a standard k-d tree and a novel linked-cell grid, and demonstrate that jFoF attains up to an order-of-magnitude speed-up compared to optimized central processing unit (CPU) implementations while maintaining consistent halo catalogues. Beyond performance, jFoF enables gradient propagation through discrete halo-finding operations via both frozen-assignment and topological optimization modes. Using a topological optimization approach via a REINFORCE-style estimator, our approach allows smooth optimization of halo connectivity and membership, bridging continuous simulation fields with discrete structure catalogues. These capabilities make jFoF a foundation for differentiable inference, enabling end-to-end, gradient-based optimization of structure formation models within GPU-accelerated astrophysical pipelines. We make our code publicly available at https://github.com/bhorowitz/jFOF/