Welcome to jax-healpy’s documentation!#
jax-healpy is a JAX-based implementation of HEALPix (Hierarchical Equal Area isoLatitude Pixelization) functions, designed for high-performance scientific computing with GPU acceleration support.
This project provides a JAX-native implementation of popular HEALPix functions, enabling automatic differentiation, just-in-time compilation, and seamless integration with modern machine learning and scientific computing workflows.
Note
This project is in beta stage. APIs may change and some features are still under development.
Key Features#
GPU Acceleration: Leverage JAX’s XLA compilation for high-performance computing on CPUs and GPUs
Automatic Differentiation: Full support for forward and reverse-mode automatic differentiation
Vectorized Operations: Efficient batch processing of HEALPix operations
HEALPix Compatibility: Drop-in replacement for many healpy functions
Spherical Harmonics: Integration with s2fft for spherical harmonic transforms
Clustering Tools: Advanced clustering algorithms for astronomical data analysis
User Guide
Examples
API Reference
- Pixel Functions (pixelfunc)
- pixelfunc.py : Healpix pixelization related functions
pix2ang()ang2pix()pix2xyf()xyf2pix()pix2vec()vec2pix()ang2vec()vec2ang()get_interp_weights()get_interp_val()get_all_neighbours()nest2ring()ring2nest()reorder()ud_grade()nside2npix()npix2nside()nside2order()order2nside()order2npix()npix2order()nside2resol()nside2pixarea()isnsideok()isnpixok()get_nside()maptype()- Coordinate Conversions
- Pixel Coordinates
- Scheme Conversions
- Map Resolution Functions
- Interpolation
- Neighbor Functions
- HEALPix Parameters
- Utility Functions
- Constants
- Spherical Harmonic Functions (sphtfunc)
- Clustering Functions
- Query Disc Functions
Development