nbodykit: An Open-source, Massively Parallel Toolkit for Large-scale Structure

A Catalog is a Python object derived from a CatalogSource class that holds information about discrete objects¹⁹ in a column-based format. Catalogs implement a random-read interface, which allows users to access arbitrary slices of the data while also taking advantage of the high throughput of a parallel file system. Often, users will initialize Catalog objects by reading data from a file on disk, using a NumPy array already stored in memory, or by generating simulated particles at runtime using one of nbodykit's built-in classes.

A Mesh is a Python object that computes a discrete representation of a continuous quantity on a uniform mesh. It is derived from a MeshSource class and provides a paintable interface, which refers to the process of "painting" the density field values onto the discrete mesh cells. When the user calls the paint() function, the mesh data is returned as a three-dimensional array. Mesh objects can be created directly from a Catalog via the to_mesh() function or by generating simulated fields directly on the mesh.

Algorithms are implemented as Python classes and interact with data by consuming Catalog and Mesh objects as input. The algorithm is executed when the user initializes the class, and the returned instance stores the results as attributes.

Most objects in nbodykit are serializable²⁰ via a save() function. Algorithm classes not only save the result of the algorithm but also save input parameters and metadata. They typically implement both a save() and load() function, such that the algorithm result can be de-serialized into an object of the same type. The two main data containers, catalogs and meshes, can be serialized using nbodykit's intrinsic format, which relies on the massively parallel IO library bigfile (Feng 2017c). nbodykit includes support for reading these serialized results from disk back into Catalog or Mesh objects.

nbodykit is fully parallelized using the Python bindings of the MPI standard available through mpi4py. The MPI standard allows processes running in parallel, each with their own memory, to exchange messages. This mechanism enables independent results to be computed by individual processes and then combined into a single result.

Both the Catalog and Mesh objects are distributed data containers, meaning that the data are spread out evenly across the available processes within an MPI communicator.²¹ Nearly all algorithm calculations are performed on this distributed data, with final results computed via a reduce operation across all processes in the communicator. Rarely throughout the code base, data are instead gathered to a single root process, and operations are performed on these data before re-distributing the results to all processes. This only occurs when wall-clock time will not be a concern for most use cases and the additional complexity of a massively parallel implementation is not merited.

The distributed nature of the Catalog object is implemented by using the random-read interface to access different slices of the tabular data for different processes. The values of a Mesh object are stored internally on a three-dimensional NumPy array, which is distributed evenly across all processes. The domain of the 3D mesh is decomposed across parallel processes using the particle mesh library pmesh, which also provides an interface for computing parallel FFTs of the mesh data using pfft-python. The pfft-python software exhibits excellent scaling with the number of available processes, enabling high-resolution (large number of cells) mesh calculations.

The analysis of LSS data often involves hundreds to thousands of repetitions of a single, less computationally expensive task. Examples include estimating the covariance matrix of a clustering statistic from a set of simulations and best-fit parameter estimation for a model. nbodykit implements a TaskManager utility to allow users to easily iterate over multiple tasks while executing in parallel. Users can specify the desired number of processes assigned to each task, and the TaskManager will iterate through the tasks, ensuring that all processes are being utilized.

We provide support for loading data from disk into Catalog objects for several of the most common data storage formats in LSS data analysis. These formats include plaintext comma-separated value (CSV) data (via pandas; McKinney 2010), binary data stored in a column-based format, HDF5 data (via h5py; Collette 2017), FITS data (via fitsio; Sheldon 2017), and the bigfile data format. We also provide more specialized readers for particle data from the Tree-PM simulations of White (2002) and the legacy binary format of the GADGET simulations (Springel 2005). These Catalog objects use the nbodykit.io module, which includes several "file-like" classes for reading data from disk. These file-like objects implement a read() function that provides the random-read interface which returns a slice of the data for the requested columns. Users can easily support custom file formats by implementing their own subclass and read() interface.

Formats storing data on disk in a column-based structure yield the best performance results, as the entirety of the data do not need to be parsed to yield the desired slice of data on a given process. This is not true for the CSV storage format. We mitigate performance issues by implementing an enhanced version of the CSV parser in pandas that supports faster parallel random access. Our preferred IO format, bigfile, is massively parallel and stores data in a column-based format.

Finally, the Catalog object supports loading data from multiple files at once, providing a continuous view of the entirety of the data. This becomes particularly powerful when combined with the random-read interface, as arbitrary slices of the combined data can be accessed. For example, a single Catalog object can provide access to arbitrary slices of the output binary snapshots from an N-body simulation (stored over multiple files), often totaling 10–100 GB in size.

nbodykit includes several Catalog classes that generate simulated data at runtime. The simplest of these allows generating random columns of data in parallel using the numpy.random module. We also provide a UniformCatalog class that generates uniformly distributed particles in a box. These classes are useful for testing purposes, as well as for use as unclustered, synthetic data in clustering estimators.

nbodykit also includes functionality for generating more realistic approximations of LSS. LogNormalCatalog generates a set of objects by Poisson sampling a log-normal density field and applies the Zel'dovich approximation to model nonlinear evolution (Coles & Jones 1991; Agrawal et al. 2017). The user can specify the input linear power spectrum and the desired output redshift of the catalog.

Catalog objects can also be created using the mock generation techniques of the Halotools software (Hearin et al. 2017) for populating halos with objects. Halotools includes functionality for populating halos via a wide range of techniques, including the halo occupation distribution (HOD), conditional luminosity function, and abundance matching methods. We refer the reader to Hearin et al. (2017) for further details. nbodykit supports using a generic Halotools model to populate a halo catalog. We also include built-in, specialized support for the HOD models of Zheng et al. (2007), Leauthaud et al. (2011), and Hearin et al. (2016).

Finally, the fastpm-python package implements an nbodykit Catalog object that generates particles via the FastPM approximate N-body simulation scheme (Feng et al. 2016). The FastPM library is massively parallel and exhibits excellent strong scaling with the number of available processes (see Section 6).

nbodykit uses the dask library (Dask Development Team 2016) to represent the data columns of a Catalog object as dask array objects instead of using the more familiar NumPy array. The dask array has two key features that help users work interactively with data and, in particular, large data sets. The first feature is delayed evaluation. When manipulating a dask array, operations are not evaluated immediately but instead stored in a task graph. Users can explicitly evaluate the dask array (returning a NumPy array) via a call to a compute() function. Second, dask arrays are chunked. The array object is internally divided into many smaller arrays, and calculations are performed on these smaller "chunks."

The delayed evaluation of dask arrays is particularly useful during the process of data cleaning, when users manipulate input data before feeding it into the analysis pipeline. Common examples of data cleaning include changing the coordinate system from galactic to Euclidean, converting between unit conventions, and applying masks. When using large data sets, the time to load the full data set into memory can be significant. This delay hinders data exploration and limits the interactive benefits of the Python language. dask arrays allow users to design data-cleaning pipelines on the fly. If the data format on disk supports random-read access, users can easily select and peek at a small subset of data without reading the full data set. This becomes especially useful when prototyping scientific models in an interactive environment, such as a Jupyter notebook.

The chunked nature of the dask array allows array computations to be performed on large data sets that do not fit into memory because the chunk size defines the amount of data loaded into memory at any given time. It effectively extends the maximum size of usable data sets from the size of the memory to the size of the disk storage. This feature also simplifies the process of dealing with large data sets in interactive environments.

The Mesh object implements a paint() function, which is responsible for generating the field values on the mesh and returning an array-like object to the user. Meshes provide an equal treatment of configuration and Fourier space, and users can specify whether the painted array is defined in configuration or Fourier space. In the former case, a RealField is returned and in the latter, a ComplexField. These objects are implemented by the pmesh package and are subclasses of the NumPy ndarray class. They are related via a real-to-complex parallel FFT operation, implemented using pfft-python via the r2c() and c2r() functions.

The paint() function paints mass-weighted (or equivalently, number-weighted) quantities to the mesh. The field that is painted is

$$\begin{eqnarray}&&F({\boldsymbol{x}})=[1+\delta ^{\prime} ({\boldsymbol{x}})]V({\boldsymbol{x}}),\end{eqnarray} \tag{ 1 }$$

where $V({\boldsymbol{x}})$ represents the field value painted to the mesh, and $\delta ^{\prime} ({\boldsymbol{x}})=n^{\prime} ({\boldsymbol{x}})/\bar{n}^{\prime} -1$ is the weighted overdensity field. It is related to the unweighted number density as $n^{\prime} ({\boldsymbol{x}})\,=W({\boldsymbol{x}})n({\boldsymbol{x}})$, where $W({\boldsymbol{x}})$ are the weights.

In nbodykit, users can control the behavior of both $V({\boldsymbol{x}})$ and $W({\boldsymbol{x}})$. In the default case, both quantities are unity, and the field painted to the mesh is $1+\delta$. As an illustration, $V({\boldsymbol{x}})$ can be specified as a velocity component to paint the momentum field (mass-weighted velocity). We also provide a mechanism by which users can further transform the painted field on the mesh. The apply() function can be used to apply a function to the mesh, either in configuration or in Fourier space. Multiple functions can be applied to the mesh, and the operations are performed when paint() is called.

All Catalog objects include a to_mesh() function, which creates a Mesh object using the specified number of cells per mesh side. This function allows users to configure exactly how the catalog is interpolated onto the mesh. Users can choose from several different mass assignment windows, including the Cloud-In-Cell (CIC), Triangular Shaped Cloud (TSC), and Piecewise Cubic Spline (PCS) schemes (Hockney & Eastwood 1981). The Daubechies wavelet (Daubechies 1992) and its symmetric counterpart ("Symlets;" see, e.g., PyWavelets²³ ) are also available. By default, the CIC window is used. The interlacing technique (Hockney & Eastwood 1981; Sefusatti et al. 2016) can reduce the effects of aliasing in Fourier space. In this scheme, the Catalog object is interpolated onto two separate meshes separated by half of a cell size. When the fields are combined in Fourier space, the leading-order contribution to aliasing is eliminated.

Users can also configure whether or not the window is compensated, which divides the density field in Fourier space by (Hockney & Eastwood 1981)

$$\begin{eqnarray}&&W({\boldsymbol{k}})={{\rm{\Pi }}}_{i}{[\mathrm{sinc}(\pi {k}_{i}/2{k}_{{\rm{N}}})]}^{p},\end{eqnarray} \tag{ 2 }$$

where $i\in \{x,y,z\}$; $p=2,3,4$ for CIC, TSC, and PCS, respectively; and $\mathrm{sinc}(x)\equiv \sin (x)/x$. The Nyquist frequency of the mesh is given by ${k}_{{\rm{N}}}=\pi N/L$, where L is the box size, and N is the number of cells per box side.

We provide comparisons of the various interpolation windows and correction methods in this section. First, Figure 2 illustrates the effects of interlacing when using the CIC, TSC, and PCS schemes. This comparison is similar to the detailed analysis presented in Sefusatti et al. (2016). Second, we show the effectiveness of the wavelet windows at reducing aliasing in Figure 3. For both figures, we paint a LogNormalCatalog of $5\times {10}^{7}$ objects to a mesh of 512³ cells in a box of side length $2500\,{h}^{-1}\,\mathrm{Mpc}$. We compare the measured power spectrum to a "reference" power spectrum, computed using a mesh of 1024³ cells and the PCS window. When using the CIC, TSC, and PCS windows, we de-convolve the interpolation window using Equation (2), while we apply no such corrections when using wavelet-based windows.

Zoom In Zoom Out Reset image size

Zoom In Zoom Out Reset image size

Figure 2 confirms the results of Sefusatti et al. (2016)—the interlacing technique performs very well at reducing the effects of aliasing on the measured power spectrum. We achieve subpercent accuracy up to the Nyquist frequency when combining interlacing with the CIC, TSC, and PCS windows. In general, higher order windows perform better, with the PCS scheme achieving a precision of at least ∼10⁻⁵ up to the Nyquist frequency.

Figure 3 compares the performance of the Daubechies and Symlet wavelets to the CIC, TSC, and PCS windows. As in Figure 2, we plot the ratio of the power spectrum computed using meshes of size 512³ and 1024³ cells. We apply Equation (2) for the CIC, TSC, and PCS windows but do not apply any corrections when using the wavelet windows. For this comparison, we do not use interlacing. We are able to confirm the results of Cui et al. (2008) and Yang et al. (2009), which claim 2% accuracy on the power spectrum up to $k\approx 0.7{k}_{{\rm{N}}}$ when using the DB6 window without any additional corrections. However, the wavelet windows fail to match the precision achieved when using interlacing, even when using the largest wavelet size tested here (a = 20). Furthermore, the Daubechies windows introduce scale dependence on large scales due to symmetry breaking (see the inset of Figure 3). The symmetric Symlet wavelets do not suffer from this issue but also cannot match the accuracy achieved when using interlacing.

Figure 3 also displays the relative speeds of each of the windows discussed in this section (bottom panel). These timing tests were performed using 64 cores on the NERSC Cori Phase I system. The wavelet windows are all significantly slower than the CIC, TSC, and PCS windows. The TSC and PCS methods are only marginally slower than the default CIC scheme, with slowdowns of ∼10% and ∼40%, respectively. The CIC, TSC, and PCS windows rely on optimized implementations in pmesh, while the wavelet windows use a slower lookup table implementation. Due to the precision of the interlacing technique and the relative speed of the TSC and PCS windows, we recommend using these options in most instances. However, it is generally best to determine the optimal set of parameters for a particular application by running convergence tests with different parameter configurations.

We demonstrate the use of Mesh objects through an example in Figure 4, which gives a short code snippet that creates a Mesh object from an existing Catalog, saves the configuration space density field to disk, and then reloads the data into memory. The snippet also demonstrates the preview() function, which can create a lower resolution projection of the full mesh field. This can be useful to quickly inspect mesh fields interactively, which would otherwise be difficult due to memory limitations. We show the preview of the density field from a log-normal catalog in the bottom panel of Figure 4, where the LSS is clearly evident, even in the low-resolution projection.

Zoom In Zoom Out Reset image size

For simulation boxes with periodic boundary conditions, the FFTPower algorithm measures the power directly from the square of the Fourier modes of the overdensity field. The 1D or 2D power spectrum, P(k) or $P(k,\mu )$, can be computed, as well as the power spectrum multipoles ${P}_{{\ell }}(k)$. Here, μ represents the angle cosine between the pair separation vector and the line of sight. For observational data, in the form of R.A., decl., and redshift, the power spectrum multipoles of the density field can be computed using the ConvolvedFFTPower algorithm. The output of this algorithm can be accurately modeled using a theoretical power spectrum convolved with the survey window function (see, e.g., Beutler et al. 2017; Hand et al. 2017b). The implementation uses the FFT-based estimator described in Hand et al. (2017a), which requires $2{\ell }+1$ FFTs to compute a given multipole of order ℓ. This estimator improves the FFT-based estimator presented by Bianchi et al. (2015) and Scoccimarro (2015), building on the ideas of previous power spectrum estimators (Feldman et al. 1994; Yamamoto et al. 2006) and in particular, the treatment of the anisotropic 2PCF using spherical harmonics of Slepian & Eisenstein (2015b). We also provide ProjectedFFTPower for computing the power spectrum of a field in a simulation box, projected along the specified axes. Such an observable can be useful for, e.g., Lyα or weak lensing data analysis (although modifications must be made for realistic cases beyond its current idealized form). The correctness of these algorithms has been verified using independent implementations from within the Baryon Oscillation Spectroscopic Survey (BOSS) collaboration.

nbodykit includes functionality for counting pairs of objects and computing their correlation function in configuration space. We leverage the blazing speed²⁴ of the publicly available Corrfunc chaining mesh code for these calculations (Sinha & Garrison 2017). We adapt its highly optimized pair-counting routines to perform calculations using MPI. We perform a domain decomposition on the input data such that the objects on a particular MPI rank are spatially confined to include all pairs within the maximum separation. For non-uniform density fields, the domain decomposition results in a particle load that is balanced across MPI ranks.²⁵ The relevant pair-counting algorithms are SimulationBoxPairCount and SurveyDataPairCount. These classes can count pairs of objects as a function of the 3D separation r, the separation r and angle to the line of sight μ, the angular separation θ, and the projected distances perpendicular r_p and parallel π to the line of sight.

Users can compute the correlation function of a Catalog using the SimulationBox2PCF and SurveyData2PCFclasses, which internally rely on the previously described pair-counting classes. For data with periodic boundary conditions, we use analytic randoms to estimate the correlation function using the so-called "natural" estimator: ${DD}/{RR}-1$. A Catalog object holding synthetic randoms can be supplied, in which case the Landy–Szalay estimator (Landy & Szalay 1993) is employed: $({DD}-2{DR}+{RR})/{RR}$. The variations of the correlation function that can be computed by these two classes are as follows:

• 1.  
  as a function of three-dimensional separation, $\xi (r)$;
• 2.  
  accounting for the angle to the line of sight, $\xi (r,\mu )$ and $\xi ({r}_{p},\pi )$;
• 3.  
  as a function of angular separation, $w(\theta )$;
• 4.  
  projected over the line-of-sight separations, ${w}_{p}({r}_{p})$.

The correctness of the pair-counting and correlation function algorithms described here was independently verified using the kdcount and Halotools software.

The SimulationBox3PCF and SurveyData3PCF classes compute the multipoles of the isotropic 3PCF in configuration space. The algorithm follows the implementation described in Slepian & Eisenstein (2015b), which scales as ${ \mathcal O }({N}^{2})$, where N is the number of objects. Their improved estimator relies on a spherical harmonic decomposition to achieve a similar scaling with N as two-point clustering estimators. We note that the FFT-based implementation of this algorithm (presented in Slepian & Eisenstein 2016) and the anisotropic version described in Slepian & Eisenstein (2018) have not yet been implemented, although there are plans to do so in the future. We have verified the accuracy of the isotropic 3PCF classes against the implementation used in Slepian & Eisenstein (2015b). An implementation of this algorithm including anisotropy written in C++ and optimized for HPC machines was recently presented in Friesen et al. (2017).

The FOF class implements the well-known FOF algorithm, which identifies clusters of points that are spatially less distant than a threshold linking length. It uses a parallel implementation of the algorithm described in Feng & Modi (2017), which utilizes KD trees and the kdcount software. FOF groups can be identified as a function of three-dimensional or angular separation. We also provide functions for transforming the output of the FOF algorithm to a Catalog of halo objects (a HaloCatalog) in a manner compatible with the Halotools software. While the FOF algorithm is commonly used in LSS analysis, it can lead to spurious artifacts; see, for example, Everitt et al. (2009) and Jain et al. (2004).

nbodykit can also identify clusters of objects using a cylindrical rather than spherical geometry. We implement a parallel version of the algorithm described in Okumura et al. (2017) in the CylindricalGroups class. Our implementation relies on the neighbor querying capability of kdcount and the group-by methods of pandas.

Finally, the FiberCollisions class simulates the process of assigning spectroscopic fibers to objects in a fiber-fed redshift survey such as BOSS or eBOSS (Dawson et al. 2013, 2016). This procedure results in so-called "fiber collisions" when two objects are separated by an angular width on the sky that is smaller than the fiber size. We follow the procedure outlined in Guo et al. (2012) to assign fibers to an input catalog of objects. We identify angular FOF groups using a linking length equal to the fiber collision scale and assign fibers to the objects in such a way as to minimize the number of objects that do not receive a fiber.

nbodykit also includes algorithms that generally serve a supporting role in other algorithms. The KDDensity class estimates a proxy density quantity for an input set of objects using the inverse cube of the distance to an object's nearest neighbor. The RedshiftHistogram class computes the mean number density as a function of redshift, n(z), from an input catalog of objects. We plan to generalize this algorithm to be a more universal histogram calculator that could, for example, compute mass or luminosity functions.