cuHPX: GPU-Accelerated Differentiable Spherical Harmonic Transforms on HEALPix Grids

2025-10-02

AI Summary

To address computational bottlenecks in HEALPix-based spherical harmonic transforms (SHTs)—arising from irregular pixel geometry, latitude-dependent tiling, and high-resolution requirements—this work introduces the first CUDA-accelerated, differentiable SHT library. Our method integrates GPU-parallelized computation, high-precision spherical harmonic numerical algorithms, optimized memory access locality, and an efficient resampling scheme between HEALPix and equiangular grids. Key contributions include: (i) the first end-to-end differentiable HEALPix SHT implementation, enabling gradient backpropagation and out-of-core computation; (ii) over 20× speedup versus state-of-the-art CPU and GPU SHT libraries, while preserving numerical consistency (absolute error < 1e−12) and rapid spectral convergence. The library has been successfully deployed in geoscience, cosmological simulations, and spherical deep learning, significantly enhancing scalability and efficiency for large-scale spherical data analysis.

Technology Category

• Machine Learning: Hardware-aware ML
• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Heuristic Search

Application Category

• Search and Retrieval-Augmented AI: Multilingual and cross-lingual Web search
• Security and Privacy: Data transparency and provenance
• Systems and Infrastructure for Web, Mobile and WoT: Data management and stream processing for Web, mobile and wireless applications

Abstract

HEALPix (Hierarchical Equal Area isoLatitude Pixelization) is a widely adopted spherical grid system in astrophysics, cosmology, and Earth sciences. Its equal-area, iso-latitude structure makes it particularly well-suited for large-scale data analysis on the sphere. However, implementing high-performance spherical harmonic transforms (SHTs) on HEALPix grids remains challenging due to irregular pixel geometry, latitude-dependent alignments, and the demands for high-resolution transforms at scale. In this work, we present cuHPX, an optimized CUDA library that provides functionality for spherical harmonic analysis and related utilities on HEALPix grids. Beyond delivering substantial performance improvements, cuHPX ensures high numerical accuracy, analytic gradients for integration with deep learning frameworks, out-of-core memory-efficient optimization, and flexible regridding between HEALPix, equiangular, and other common spherical grid formats. Through evaluation, we show that cuHPX achieves rapid spectral convergence and delivers over 20 times speedup compared to existing libraries, while maintaining numerical consistency. By combining accuracy, scalability, and differentiability, cuHPX enables a broad range of applications in climate science, astrophysics, and machine learning, effectively bridging optimized GPU kernels with scientific workflows.

Problem

• Optimizing spherical harmonic transforms on HEALPix grids
• Addressing irregular geometry and high-resolution demands
• Enabling differentiable transforms for deep learning integration

Innovation

• GPU-accelerated CUDA library for HEALPix transforms
• Provides differentiable spherical harmonic analysis capabilities
• Enables flexible regridding between multiple spherical formats