A Numba-based two-point correlation function (2PCF) calculator using a grid decomposition.
Like Corrfunc, but written in Numba,
with simplicity and hackability in mind.
Aside from the 2PCF calculation, the
particle_grid module is both simple and
fast and may be useful on its own as a way to partition particle sets in 3D.
$ git clone https://github.com/lgarrison/numba-2pcf.git $ cd numba-2pcf $ python -m pip install -e .
from numba_2pcf.cf import numba_2pcf import numpy as np rng = np.random.default_rng(123) N = 10**6 box = 2. pos = rng.random((N,3), dtype=np.float32)*box res = numba_2pcf(pos, box, Rmax=0.05, nbin=10) res.pprint_all()
rmin rmax rmid xi npairs -------------------- -------------------- -------------------- ----------------------- -------- 0.0 0.005000000074505806 0.002500000037252903 -0.004519257448573177 65154 0.005000000074505806 0.010000000149011612 0.00750000011175871 0.0020113763064291135 459070 0.010000000149011612 0.01500000022351742 0.012500000186264515 0.000984359247434119 1244770 0.01500000022351742 0.020000000298023225 0.017500000260770324 -6.616896085054336e-06 2421626 0.020000000298023225 0.02500000037252903 0.022500000335276125 0.00019365366488166558 3993210 0.02500000037252903 0.03000000044703484 0.027500000409781934 5.769329601057471e-05 5956274 0.03000000044703484 0.03500000052154064 0.032500000484287736 0.0006815801672250821 8317788 0.03500000052154064 0.04000000059604645 0.037500000558793545 2.04711840243732e-05 11061240 0.04000000059604645 0.04500000067055226 0.042500000633299354 9.313641918828885e-05 14203926 0.04500000067055226 0.05000000074505806 0.04750000070780516 -0.00011690771042793813 17734818
The goal of this project is not to provide the absolute best performance that
given hardware can produce, but it is a goal to provide as good performance
as Numba will let us reach (while keeping the code readable). So we pay special
attention to things like
float32 particle inputs when possible!),
parallelization, and some early-exit conditions (when we know a pair can’t fall
in any bin).
As a demonstration that this code provides passably good performance,
here’s a dummy test of 107 unclustered data points in a 2 Gpc/h box
(so number density 1.2e-3), with Rmax=200 Mpc/h and bin width of 1 Mpc/h:
from numba_2pcf.cf import numba_2pcf import numpy as np rng = np.random.default_rng(123) N = 10**6 box = 2000 pos = rng.random((N,3), dtype=np.float32)*box %timeit numba_2pcf(pos, box, Rmax=150, nbin=150, corrfunc=False, nthread=24) # 3.5 s %timeit numba_2pcf(pos, box, Rmax=150, nbin=150, corrfunc=True, nthread=24) # 1.3 s
So within a factor of 3 of Corrfunc, and we aren’t even exploiting the
symmetry of the autocorrelation (i.e. we count every pair twice). Not bad!
Testing Against Corrfunc
The code is tested against Corrfunc. And actually, the
numba_2pcf() function takes a flag
corrfunc=True that calls Corrfunc
instead of the Numba implementation to make such testing even easier.
numba_2pcf works a lot like Corrfunc, or any other grid-based 2PCF code: the
3D volume is divided into a grid of cells at least
Rmax in size, where
is the maximum radius of the correlation function measurement. Then, we know
all valid particle pairs must be in neighboring cells. So the task is simply
to loop through each cell in the grid, pairing it with each of its 26 neighbors
(plus itself). We parallelize over cell pairs, and add up all the pair counts
across threads at the end.
This grid decomposition prunes distant pairwise comparisons, so even though
the runtime still formally scales as O(N2), it makes the 2PCF
tractable for many realistic problems in cosmology and large-scale structure.
A numba implementation isn’t likely to beat Corrfunc on speed, but numba
can still be fast enough to be useful (especially when the computation parallelizes
well). The idea is that this code provides a “fast enough” parallel implementation
while still being highly readable — the 2PCF implementation is about 150 lines
of code, and the gridding scheme 100 lines.
particle-jackknife branch contains an implementation of an idea for computing
the xi(r) variance based on the variance of the per-particle xi(r) measurements.
It doesn’t seem to be measuring the right thing, but the code is left for posterity.
This repo was generated from @DFM’s Cookiecutter Template. Thanks, DFM!