torch-harmonics: Differentiable Spherical Harmonic Transforms for PyTorch

What torch-harmonics Solves for Spherical Data

Many physical simulation problems are naturally defined on a sphere: atmospheric modeling, climate prediction, and acoustic propagation on a globe. These problems require processing signals on the two-dimensional surface S2, embedded in three dimensions. Standard FFT-based signal processing applies to flat domains; applying it to spherical data requires a different mathematical basis.

Spherical harmonics are the natural orthogonal basis for functions defined on S2. torch-harmonics implements the forward and inverse spherical harmonic transform (SHT), the vector spherical harmonic transform (vSHT), and discrete-continuous convolutions on the sphere. All of these operations are implemented using PyTorch primitives, which makes them fully differentiable. Gradients flow back through the transforms during training, which is necessary for using these operations as layers inside a neural network.

The library was originally developed to enable Spherical Fourier Neural Operators (SFNO), an architecture for global weather prediction that uses spherical convolutions. The README also lists differentiable PDE solvers, including shallow water equations, Ginzburg-Landau, and Allen-Cahn, as applications built with the library.

The SHT Algorithm: Quadrature and FFTs

The spherical harmonic transform in torch-harmonics uses two mathematical components. For the projection onto the associated Legendre polynomials, the algorithm uses quadrature rules. For the projection onto the harmonic basis, it uses the Fast Fourier Transform. The README states that this algorithm tends to outperform alternatives with better asymptotic scaling for most practical grid sizes.

The quadrature step can be distributed across multiple GPU ranks, which means the spatial dimension of the transform can be split across devices. This distributed quadrature is how the library supports training with spatially distributed data, such as a global atmospheric field that does not fit on a single GPU.

The library compiles custom CUDA extensions for acceleration. The setup.py checks for GPU devices at build time; if CUDA devices are not detected inside a container, the environment variable TORCH_HARMONICS_BUILD_CUDA_EXTENSION=1 forces the CUDA build. Custom CUDA extensions require GPU architectures with compute capability 7.0 or higher.

Installing torch-harmonics: CPU, CUDA Wheels, and Source

Three installation paths are available. The simplest is the CPU-only wheel on PyPI, which requires only pip:

pip install torch-harmonics

This version does not include compiled CUDA kernels and is built against the newest PyTorch release at the time of packaging.

For GPU support, NVIDIA provides prebuilt wheels on pypi.nvidia.com with CUDA-version-specific packages. The rolling alias for the latest CUDA build is:

pip install torch-harmonics-cuda-latest --extra-index-url https://pypi.nvidia.com

Specific CUDA versions are available as torch-harmonics-cu126, torch-harmonics-cu128, torch-harmonics-cu129, and torch-harmonics-cu130, paired with specific PyTorch versions. The README recommends running nvidia-smi to check the installed driver's CUDA version before choosing a package.

Building from source is available for configurations not covered by the prebuilt wheels:

git clone [email protected]:NVIDIA/torch-harmonics.git
cd torch-harmonics
pip install --no-build-isolation -e .

The --no-build-isolation flag ensures the custom CPU and CUDA kernels compile against the existing torch installation rather than a fresh build-time copy. A Docker build is also documented in the repository's Dockerfile.

Controlling CUDA Architecture Compilation

The source build compiles CUDA kernels for the GPU architectures named in TORCH_CUDA_ARCH_LIST. Setting this variable to only the architectures in use reduces build time:

export TORCH_CUDA_ARCH_LIST="8.0 8.6 9.0a 10.0a+PTX"
pip install --no-build-isolation -e .

The README highlights one compilation subtlety: appending the suffix 'a' to an architecture identifier (for example, 9.0a instead of 9.0) enables architecture-specific tensor core instructions such as wgmma on Hopper GPUs. The README states that DISCO convolution layers benefit significantly from this. The trade-off is that a binary compiled with 9.0a runs only on the exact GPU generation it was built for, not on the broader architecture family. Omitting the 'a' suffix produces a portable binary that runs on the whole generation but without those tensor core instructions.

This compilation choice matters when deploying a trained model on infrastructure that may include different GPU generations. A portable build handles heterogeneous clusters; an architecture-specific build maximizes performance on a known target.

Limitations: CUDA Requirements and Licensing

The custom CUDA extensions require GPU compute capability 7.0 or higher. Older GPUs (compute capability 6.x or below) cannot use the compiled kernels. The CPU-only installation works on any machine but does not benefit from GPU acceleration.

The library requires Python 3.10 or newer and torch >= 2.6.0, as declared in pyproject.toml. Projects using older PyTorch versions must upgrade before adding torch-harmonics as a dependency.

The license listed in the repository is BSD-3-Clause. The third clause prohibits using the names of the copyright holders or contributors to endorse or promote products derived from the library without prior written permission. This is a stricter requirement than the MIT license: redistributors building products on top of torch-harmonics cannot reference NVIDIA or the contributing authors in marketing materials without that permission. The copyright header appears in the repository files and must be retained in source and binary redistributions.

torch-harmonics vs. SHTools: Differentiability vs. Breadth

SHTools is a Python and Fortran library for spherical harmonic transforms used in geophysics. It covers a wide range of spherical harmonic conventions, gravity and magnetic field analysis, and rotation operations. It is designed for numerical analysis and scientific computing where the priority is correct, flexible output rather than gradient propagation.

torch-harmonics makes a different design choice. It implements fewer conventions but ensures that every operation is differentiable, making it usable as a layer in PyTorch models trained with backpropagation. An atmospheric model trained with gradient descent would use torch-harmonics; a geophysical analysis pipeline that reads satellite gravity data and outputs spectral coefficients would be more at home with SHTools.

The practical boundary is the PyTorch dependency. torch-harmonics is not a standalone library: it requires torch >= 2.6.0 and is tightly coupled to PyTorch's tensor operations. Workflows outside of PyTorch cannot use it.

Editorial conclusion

torch-harmonics is the right choice when a PyTorch training pipeline needs spherical harmonic transforms and the gradients must flow through them: weather modeling, global PDE solvers, and spherical neural operators. It is the wrong tool for traditional numerical analysis that does not involve gradient-based optimization, or for anyone without a GPU with compute capability 7.0 or higher to take advantage of the compiled CUDA extensions. The BSD-3-Clause license allows redistribution in commercial products, but the copyright notice and disclaimer must be reproduced in any distributed form.

Frequently asked questions

Official sources

1. Issues
2. NVIDIA/torch-harmonics on GitHub
3. README
4. Releases