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

Benjamin Horowitz, Adrian E. Bayer
2026-06-17

SCID:  54.1/7f8eh589
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/
Publication Details
Publication Date
2026-06-17
Journal
Publisher
ISSN
Access Type
Author Information
Authors
Benjamin Horowitz
Adrian E. Bayer
Explore More Research
Use the citation graph to discover related papers and expand your research horizons.
Click any node to explore
Download PDF
100%