Source-linked AI summary
nbodykit: an open-source, massively parallel toolkit for large-scale structure
Nick Hand, Yu Feng, Florian Beutler, Yin Li, Chirag Modi, Uros Seljak, Zachary Slepian
TL;DR
Large-scale structure analyses need standardized software that can handle increasing data volumes and wall-clock demands. nbodykit provides a modular, massively parallel Python toolkit with distributed data containers and canonical LSS algorithms, and its benchmarks demonstrate strong scaling for repeated analyses. The authors present it as a community foundation for future LSS data sets, while noting that its current feature set is not all-inclusive.
Problem
Increasing LSS data volumes, modeling complexity, and wall-clock times create a need for standardized, scalable analysis software.
Method
nbodykit combines parallel Python implementations of canonical LSS algorithms with distributed data containers and a modular component-based architecture.
Results
The toolkit demonstrates excellent scaling behavior in performance benchmarks and supports interactive use on HPC systems.
Takeaways & Limitations
nbodykit is intended as a foundation for researchers to build complex LSS applications and address future survey data-analysis demands.
Takeaways & Limitations
The implemented features are a sampling of commonly used tools rather than an all-inclusive set.
Abstract
from arXiv · showhide
We present nbodykit, an open-source, massively parallel Python toolkit for analyzing large-scale structure (LSS) data. Using Python bindings of the Message Passing Interface (MPI), we provide parallel implementations of many commonly used algorithms in LSS. nbodykit is both an interactive and scalable piece of scientific software, performing well in a supercomputing environment while still taking advantage of the interactive tools provided by the Python ecosystem. Existing functionality includes estimators of the power spectrum, 2 and 3-point correlation functions, a Friends-of-Friends grouping algorithm, mock catalog creation via the halo occupation distribution technique, and approximate N-body simulations via the FastPM scheme. The package also provides a set of distributed data containers, insulated from the algorithms themselves, that enable nbodykit to provide a unified treatment of both simulation and observational data sets. nbodykit can be easily deployed in a high performance computing environment, overcoming some of the traditional difficulties of using Python on supercomputers. We provide performance benchmarks illustrating the scalability of the software. The modular, component-based approach of nbodykit allows researchers to easily build complex applications using its tools. The package is extensively documented at http://nbodykit.readthedocs.io, which also includes an interactive set of example recipes for new users to explore. As open-source software, we hope nbodykit provides a common framework for the community to use and develop in confronting the analysis challenges of future LSS surveys.
1. INTRODUCTION
nbodykit addresses the need for standardized, scalable LSS analysis software by combining parallel Python implementations with distributed data containers and an interactive, modular design.
- Growing LSS data volumes and modeling complexity motivate software that reduces wall-clock time and scales to next-generation surveys.
- nbodykit is an open-source Python toolkit providing fully parallel implementations of canonical LSS algorithms for HPC use.
- Its distributed, extensible data containers support large volumes and diverse formats while remaining insulated from algorithms for simulation and observational data.
- The component-based design supports interactive Python use and integration into larger, problem-specific LSS applications.
- The project is developed as open-source software, with documentation, tutorials, and interactive example recipes for new users.
2. OVERVIEW
nbodykit exposes modular Python components for loading, representing, and analyzing LSS data, combining interactive use with specialized compiled dependencies.
- Dependencies: The package uses compiled extensions for specialized, computationally expensive tasks while retaining a pure-Python high-level interface.
- 2.3. A Component-Based Approach: A modular, component-based architecture lets users combine Python classes and functions to build complex applications.
- 2.3.1. Catalog: Catalog objects store discrete objects in columnar form and support arbitrary data slices through a random-read interface.
- 2.3.2. Mesh: Mesh objects represent continuous quantities on uniform grids and can paint density values into three-dimensional arrays.
- Algorithms: Algorithms consume Catalog and Mesh objects as inputs, execute on class initialization, and store results as attributes.
2.4. Parallelism
nbodykit parallelizes data and algorithm execution with MPI, distributing containers across processes and combining results through reductions or task-level parallelism.
- Algorithm results may be consistent across processes or distributed evenly among them.
- MPI enables independent process computations whose results can be combined into a single result.
- Catalog and Mesh data are distributed evenly across communicator processes, with most calculations performed locally before reduction.
- Mesh domains use pmesh decomposition and parallel FFTs, enabling high-resolution calculations with excellent scaling across processes.
- TaskManager assigns user-specified process counts to repeated tasks such as covariance estimation and parameter fitting.
3. CAPABILITIES
This section introduces the detailed capability overview, covering cosmology calculations, data containers, and implemented algorithms.
- The capabilities overview covers cosmology calculations, Catalog and Mesh classes, and the implemented algorithms.
3.1. Cosmology
The cosmology module represents background cosmological models and computes common LSS quantities using CLASS-based functionality. It also provides compatibility with astropy while differing in its treatment of massive neutrinos.
- Cosmology representation and calculations: The module represents cosmological parameter sets and computes theoretical LSS quantities dependent on the background cosmological model.Its calculations use the CLASS Boltzmann solver through Python bindings provided by classylss.
- Cosmology representation and calculations: The Cosmology class is initialized with a unique set of parameters using CLASS syntax and exposes methods for model-dependent quantities.Examples include distance versus redshift, H(z), linear and nonlinear power spectra, and density and velocity transfer functions.
- Power-spectrum tools: The module supplies LinearPower for computing linear power spectra versus redshift and wavenumber using CLASS transfer functions or Eisenstein–Hu approximations.The Eisenstein–Hu option includes a no-wiggle transfer function.
- Interoperability and scope: nbodykit chooses CLASS over astropy.cosmology to leverage a Boltzmann solver for LSS calculations while providing syntax compatibility and transformations between cosmology classes.The two implementations nevertheless differ in their treatment of massive neutrinos.
3.2. Catalogs
nbodykit provides Catalog objects for loading, generating, cleaning, and manipulating large LSS data sets. Its distributed and chunked data model supports parallel access, synthetic catalog creation, and interactive workflows beyond the limits of in-memory processing.
- Catalog creation and data access: Catalog objects load common LSS formats including CSV, columnar binary, HDF5, FITS, and bigfile data.Specialized readers also support particle data from Tree-PM simulations.
- Catalog creation and data access: Columnar formats provide better performance because processes can read desired slices without parsing the entire data set.nbodykit accelerates CSV access with a parallel random-access parser, while bigfile is its preferred massively parallel columnar format.
- Catalog creation and data access: A Catalog can combine multiple files into one continuous view and expose arbitrary slices of large simulation snapshots.The example snapshots often total 10-100 GB.
- Synthetic catalogs: Runtime catalog classes generate random or uniformly distributed particles for testing and synthetic unclustered data.UniformCatalog generates uniformly distributed particles in a box for clustering-estimator use.
- Synthetic catalogs: LogNormalCatalog and Halotools integrations provide more realistic mock catalogs through log-normal fields, the Zeldovich approximation, and halo-population models.Users can specify an input linear power spectrum and output redshift for LogNormalCatalog, while Halotools supports HOD, conditional luminosity function, and abundance matching methods.
- Synthetic catalogs: The FastPM integration generates particles with a massively parallel approximate N-body scheme that exhibits strong scaling with available processes.It returns an nbodykit Catalog object for downstream analysis.
- Interactive data manipulation: Dask arrays represent Catalog columns with delayed evaluation, allowing users to construct data-cleaning task graphs before execution.This supports interactive transformations such as coordinate conversion, unit conversion, and masking without immediately loading all data.
- Interactive data manipulation: Chunked dask-array computation processes data sets larger than memory by limiting the amount loaded at once.The usable data-set scale therefore extends from memory capacity toward disk-storage capacity.
3.3. Meshes
nbodykit’s Mesh objects convert catalog data into distributed, configurable density fields for configuration- or Fourier-space analysis. The section also evaluates interpolation windows, showing that interlacing provides high power-spectrum accuracy while wavelets trade accuracy for speed.
- Mesh objects: Mesh objects represent continuous quantities on uniform three-dimensional grids and can be created directly from Catalog objects or simulated fields.The paint() function returns the mesh field as an array-like object.
- Painting a Mesh: The paint() function supports weighted fields, allowing users to control both the painted quantity V(x) and weights W(x), including mass-weighted velocity.With default unity values, the painted field is 1 + δ.
- From Catalog to Mesh: Catalogs support CIC, TSC, PCS, Daubechies, and Symlet interpolation windows, with optional Fourier-space compensation for the standard windows.The default interpolation window is CIC; equation 2 specifies the compensation window.
- From Catalog to Mesh: Sub-percent accuracy is achieved up to the Nyquist frequency when interlacing is combined with CIC, TSC, or PCS windows.PCS reaches precision of at least ∼10^-5 up to the Nyquist frequency.
- From Catalog to Mesh: Wavelet windows achieve 2% power-spectrum accuracy to k ≈0.7kN without corrections, but do not match interlacing precision and can introduce scale dependence.Daubechies windows show scale dependence from symmetry breaking, whereas Symlets avoid that issue but remain less accurate than interlacing.
- From Catalog to Mesh: TSC and PCS are only marginally slower than CIC, with slowdowns of ∼10% and ∼40%, respectively, while wavelet windows are significantly slower.The authors recommend TSC or PCS in most instances while advising convergence tests for application-specific parameters.
3.4. Algorithms
nbodykit provides parallel implementations of common LSS analysis algorithms for both simulations and observational surveys. Its functionality spans power spectra, correlation functions, grouping, fiber-collision modeling, and supporting density or histogram calculations.
- Overview: The algorithms module provides a foundation of parallel LSS methods rather than an exhaustive collection, supporting observational and N-body simulation data.The implementations are designed for common analysis workflows.
- Power spectrum: FFTPower computes one- or two-dimensional power spectra and multipoles, while ConvolvedFFTPower computes survey power-spectrum multipoles from angular-position and redshift data.The simulation method operates directly on Fourier modes of the overdensity field.
- Correlation functions: Parallel pair-counting and 2PCF classes measure clustering across three-dimensional, angular, line-of-sight, and projected separations.The classes use Corrfunc routines adapted for MPI and support natural or Landy-Szalay estimators.
- 3-point correlation function: The isotropic 3PCF classes compute configuration-space multipoles using a spherical-harmonic estimator scaling as O(N^2).Accuracy was verified against the implementation described by Slepian and Eisenstein; anisotropic and FFT-based variants were not yet implemented.
- Grouping: FOF identifies spatially nearby clusters with a linking length, while CylindricalGroups identifies groups using cylindrical geometry.Both are implemented in parallel and rely on KD-tree or neighbor-querying methods.
- Supporting algorithms: FiberCollisions simulates spectroscopic fiber assignment, and supporting classes estimate nearest-neighbor density proxies and redshift histograms.These tools address survey effects and auxiliary catalog calculations.
4. DEVELOPMENT WORKFLOW
nbodykit uses public version control, extensive automated testing, and documented installation paths to support development across local and HPC systems. Its deployment tooling addresses a Python-specific MPI launch bottleneck caused by shared file systems.
- Development workflow: nbodykit is developed with git and hosted in a public repository.The public development model supports access to the source code.
- Testing: Hundreds of unit tests cover serial and parallel execution through the runtests package, which adds MPI support to the pytest framework.The test driver rebuilds and installs the package before running tests.
- Testing: Continuous integration tests Linux and Mac OS X across Python 2.7, 3.5, and 3.6, blocking merges when tests fail or coverage decreases.The workflow uses Travis CI and Coveralls.
- Installation: nbodykit provides Anaconda installation for personal systems, while HPC deployments generally require recompiling dependencies with machine-specific compilers and MPI settings.Build infrastructure is publicly available and can be reused beyond NERSC systems.
- HPC deployment: Python MPI launches can stall when thousands of processes query shared filesystems for modules, effectively preventing Python applications from running on HPC machines.This is identified as the remaining deployment barrier.
- HPC deployment: python-mpi-bcast bypasses the filesystem bottleneck by broadcasting bundled runtime dependencies to compute nodes via MPI.Users can deploy it by modifying job scripts non-invasively.
- Documentation: Sphinx documentation includes API references, walkthroughs, Jupyter Cookbook recipes, and Binder access for interactive exploration without local installation.The recipes range from simple tasks to complex workflows.
5. IN ACTION
The galaxy clustering emulator combines nbodykit components to generate a galaxy power spectrum from an initial Gaussian density field. It evolves matter, identifies halos, populates galaxies, and computes P(k) for intermediate products.
- 5. IN ACTION: The emulator produces a galaxy power spectrum from a background cosmological model using a sequence of nbodykit components.The workflow begins with Gaussian initial conditions and records power spectra for the initial, matter, halo, and galaxy fields.
- 5. IN ACTION: LinearMesh creates a Gaussian density-field realization in Fourier space from an input power spectrum.
- 5. IN ACTION: FastPM evolves the initial conditions to z = 0 using a quasi-N-body particle-mesh scheme.
- 5. IN ACTION: FOF identifies halos from the evolved matter field, which are then populated with galaxies using the Zheng et al. (2007) HOD and Halotools.
- 5. IN ACTION: The component-based application can be implemented in only ∼30 lines of code while allowing intermediate data products to be output and serialized.
6. PERFORMANCE BENCHMARKS
nbodykit’s four benchmarked algorithms show excellent strong scaling with MPI ranks across small and large data configurations. FFT-based methods have nearly identical runtimes across sample sizes, while pair-counting performance depends on implementation and communication costs.
- 6. PERFORMANCE BENCHMARKS: The four benchmarked algorithms show excellent strong scaling as the number of MPI ranks increases.The tests cover FFTPower, ConvolvedFFTPower, SimulationBoxPairCount, and SimulationBox3PCF on NERSC Cori Haswell nodes.
- 6. PERFORMANCE BENCHMARKS: FFT-based algorithms have nearly identical performance for small and large samples because the dominant FFT uses the same 1024^3 mesh.
- 6. PERFORMANCE BENCHMARKS: ConvolvedFFTPower takes roughly fifteen times as long as FFTPower because it performs more FFT operations.The benchmark configuration uses 15 total FFT operations for ConvolvedFFTPower.
- 6. PERFORMANCE BENCHMARKS: SimulationBoxPairCount is significantly faster than SimulationBox3PCF because it relies on optimized Corrfunc rather than kdcount.
- 6. PERFORMANCE BENCHMARKS: MPI communication costs hinder SimulationBoxPairCount scaling on the small data set, while the emulator’s FastPM step scales strongly.
- 6. PERFORMANCE BENCHMARKS: The emulator’s halo-population step remains roughly constant with increasing core counts because its Halotools implementation is not massively parallel.This step becomes significant only as the number of cores approaches ∼1024.
7. CONCLUSIONS
The paper presents nbodykit as an open-source, massively parallel Python toolkit that combines scalable LSS algorithms, distributed data containers, and modular workflows. Its benchmarks and emulator demonstrate strong scaling, while the authors position the software as a community foundation that remains intentionally non-exhaustive.
- 7. CONCLUSIONS: nbodykit v0.3.0 provides parallel implementations of canonical LSS algorithms and distributed data containers for common astronomical formats.
- 7. CONCLUSIONS: The modular, component-based design balances interactive Jupyter use with scalable HPC execution and supports construction of complex applications.
- 7. CONCLUSIONS: The galaxy clustering emulator provides a complete forward model for the galaxy power spectrum starting from an initial Gaussian density field.
- 7. CONCLUSIONS: The toolkit demonstrates excellent scaling behavior in both the emulator and commonly used algorithms.
- 7. CONCLUSIONS: The open-source project emphasizes reproducibility through automated testing, documentation, and broad installability.
- 7. CONCLUSIONS: The implemented feature set is not all-inclusive, but is intended as a sampling of commonly used tools that the community can extend.