cosmokdtree: a flexible OpenMP-parallelized k-d tree for computational astrophysics applications

Abstract

Modern numerical astrophysics applications present a common demand for efficient spatial querying techniques, such as neighbor searches or density estimations over billions of resolution elements. To address the growing need for fast and versatile tools tailored to these tasks, we present cosmokdtree, a fast and flexible multi-purpose k-d tree implementation in Fortran, parallelized with OpenMP directives. Our library supports arbitrary dimensionality and spatial distributions, and is designed for efficient tree construction and fast query performance for moderate-scale applications (performance tested up to 10^9 points). Python bindings coupled to the module are also provided. All these ingredients yield a k-d tree package suitable for analysis purposes. We benchmark cosmokdtree across a variety of scenarios, including different point distributions, dimensionality, and parallel scaling. We also compare its performance against the widely used scipy implementation and other efficient alternatives such as coretran's and a GPU-based version, showing that cosmokdtree consistently achieves lower construction and query times than CPU alternatives while keeping a reasonable memory usage. Our tree-building phase implementation, executed on a mid-range workstation-class CPU, approaches the performance of GPU-based implementations when run on high-end consumer graphic cards, although data-center GPUs, out of the scope for our comparison, could still deliver substantially higher performance and a broader conclusion cannot be extracted. We further present some applications tackling common problems in astrophysics, namely, the friends-of-friends clustering algorithm and the particle-to-mesh assignment process. The code is publicly released and intended to serve as a flexible multi-purpose tool for computational applications in a wide range of scenarios, particularly in astrophysics.