NVIDIA/torch-harmonics

NVIDIA/torch-harmonics

Description: Differentiable signal processing on the sphere for PyTorch

Language: Jupyter Notebook

License: NOASSERTION

Stars: 677

Forks: 67

Open issues: 6

Created: 2023-01-10T19:21:10Z

Pushed: 2026-06-10T22:26:23Z

Default branch: main

Fork: no

Archived: no

README:

torch-harmonics

[Overview](#overview) | [Installation](#installation) | [More information](#more-about-torch-harmonics) | [Getting started](#getting-started) | [Contributors](#contributors) | [Cite us](#cite-us) | [References](#references)


Overview

torch-harmonics implements differentiable signal processing on the sphere. This includes differentiable implementations of the spherical harmonic transforms, vector spherical harmonic transforms and discrete-continuous convolutions on the sphere. The package was originally implemented to enable Spherical Fourier Neural Operators (SFNO) [1].

The SHT algorithm uses quadrature rules to compute the projection onto the associated Legendre polynomials and FFTs for the projection onto the harmonic basis. This algorithm tends to outperform others with better asymptotic scaling for most practical purposes [2].

torch-harmonics uses PyTorch primitives to implement these operations, making it fully differentiable. Moreover, the quadrature can be distributed onto multiple ranks making it spatially distributed.

torch-harmonics has been used to implement a variety of differentiable PDE solvers which generated the animations below. Moreover, it has enabled the development of Spherical Fourier Neural Operators [1].

Installation

Prebuilt wheels

Prebuilt Linux wheels with compiled CUDA extensions are available on pypi.nvidia.com. Pick the package matching your CUDA toolkit version:

| CUDA | Package | Min PyTorch | Install command | |------|---------|-------------|-----------------| | 12.6 | torch-harmonics-cu126 | 2.6.0 | pip install torch-harmonics-cu126 --extra-index-url https://pypi.nvidia.com | | 12.8 | torch-harmonics-cu128 | 2.7.0 | pip install torch-harmonics-cu128 --extra-index-url https://pypi.nvidia.com | | 12.9 | torch-harmonics-cu129 | 2.8.0 | pip install torch-harmonics-cu129 --extra-index-url https://pypi.nvidia.com | | 13.0 | torch-harmonics-cu130 | 2.9.1 | pip install torch-harmonics-cu130 --extra-index-url https://pypi.nvidia.com |

If you don't need a specific CUDA version, use one of the rolling aliases:

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

# CPU only
pip install torch-harmonics-cpu-latest --extra-index-url https://pypi.nvidia.com

> Tip: Run nvidia-smi to check your driver's CUDA version.

PyPI

The vanilla torch-harmonics package on PyPI ships a CPU-only prebuilt wheel. This version is built for the newest PyTorch release. For GPU support, use the NVIDIA PyPI packages above.

pip install torch-harmonics

Building from source

If your OS, PyTorch or CUDA toolkit version is not covered by the available wheels, we recomment building torch-harmonics from the GitHub repository. Use --no-build-isolation so that custom CPU and CUDA kernels compile against your existing torch installation:

git clone git@github.com:NVIDIA/torch-harmonics.git
cd torch-harmonics
pip install --no-build-isolation -e .

If CUDA devices are not detected automatically (e.g. inside a container), set the FORCE_CUDA_EXTENSION flag. Set TORCH_CUDA_ARCH_LIST to only the architectures you need to reduce compilation time:

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

:warning: Custom CUDA extensions require architectures >= 7.0.

Alternatively, build a Docker container:

git clone git@github.com:NVIDIA/torch-harmonics.git
cd torch-harmonics
docker build . -t torch_harmonics
docker run --gpus all -it --rm --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 torch_harmonics

More about torch-harmonics

Spherical harmonics

The spherical harmonics are special functions defined on the two-dimensional sphere $S^2$ (embedded in three dimensions). They form an orthonormal basis of the space of square-integrable functions defined on the sphere $L^2(S^2)$ and are comparable to the harmonic functions defined on a circle/torus. The spherical harmonics are defined as

$$ Y_l^m(\theta, \lambda) = \sqrt{\frac{(2l + 1)}{4 \pi} \frac{(l - m)!}{(l + m)!}} P_l^m(\cos \theta) \exp(im\lambda), $$

where $\theta$ and $\lambda$ are colatitude and longitude respectively, and $P_l^m$ the normalized, associated Legendre polynomials.

Spherical harmonic transform

The spherical harmonic transform (SHT)

$$ f_l^m = \int_{S^2} \overline{Y_{l}^{m}}(\theta, \lambda) f(\theta, \lambda) \mathrm{d} \mu(\theta, \lambda) $$

realizes the projection of a signal $f(\theta, \lambda)$ on $S^2$ onto the spherical harmonics basis. The SHT generalizes the Fourier transform on the sphere. Conversely, a truncated series expansion of a function $f$ can be written in terms of spherical harmonics as

$$ f (\theta, \lambda) = \sum_{m=-M}^{M} \exp(im\lambda) \sum_{l=|m|}^{M} \hat f_l^m P_l^m (\cos \theta), $$

where $\hat{f}_l^m$, are the expansion coefficients associated to the mode $m$, $n$.

The implementation of the SHT follows the algorithm as presented in [2]. A direct spherical harmonic transform can be accomplished by a Fourier transform

$$ \hat f^m(\theta) = \frac{1}{2 \pi} \int_{0}^{2\pi} f(\theta, \lambda) \exp(-im\lambda) \mathrm{d} \lambda $$

in longitude and a Legendre transform

$$ \hat f_l^m = \frac{1}{2} \int^{\pi}_0 \hat f^{m} (\theta) P_l^m (\cos \theta) \sin \theta \mathrm{d} \theta $$

in latitude.

Discrete Legendre transform

The second integral, which computed the projection onto the Legendre polynomials is realized with quadrature. On the Gaussian grid, we use Gaussian quadrature in the $\cos \theta$ domain. The integral

$$ \hat f_l^m = \frac{1}{2} \int_{-1}^1 \hat{f}^m(\arccos x) P_l^m (x) \mathrm{d} x $$

is obtained with the substitution…