A differentiable spherical harmonic transform for PyTorch.
Project description
torch-harmonics
is a differentiable implementation of the Spherical Harmonic transform in PyTorch. It was originally implemented to enable Spherical Fourier Neural Operators (SFNO). It 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.
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 (SFNOs) [1].
Installation
Download directyly from PyPI:
pip install torch-harmonics
Build in your environment using the Python package:
git clone git@github.com:NVIDIA/torch-harmonics.git
cd torch-harmonics
pip install -e .
Alternatively, use the Dockerfile to build your custom container after cloning:
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
Contributors
- Boris Bonev (bbonev@nvidia.com)
- Thorsten Kurth (tkurth@nvidia.com)
- Christian Hundt (chundt@nvidia.com)
- Nikola Kovachki (nkovachki@nvidia.com)
- Jean Kossaifi (jkossaifi@nvidia.com)
Implementation
The implementation follows the algorithm as presented in [2].
Spherical harmonic transform
The truncated series expansion of a function $f$ defined on the surface of a sphere can be written as
$$ f(\theta, \lambda) = \sum_{m=-M}^{M} \exp(im\lambda) \sum_{n=|m|}^{M} F_n^m \bar{P}_n^m (\cos \theta), $$
where $\theta$ is the colatitude, $\lambda$ the longitude, $\bar{P}_n^m$ the normalized, associated Legendre polynomials and $F_n^m$, the expansion coefficient associated to the mode $(m,n)$.
A direct spherical harmonic transform can be accomplished by a Fourier transform
$$ 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
$$ F_n^m = \frac{1}{2} \int_{-1}^1 F^m(\theta) \bar{P}_n^m(\cos \theta) \mathrm{d} \cos \theta $$
in latitude.
Discrete Legendre transform
in order to apply the Legendre transfor, we shall use Gauss-Legendre points in the latitudinal direction. The integral
$$ F_n^m = \int_{0}^\pi F^m(\theta) \bar{P}_n^m(\cos \theta) \sin \theta \mathrm{d} \theta $$
is approximated by the sum
$$ F_n^m = \sum_{j=1}^{N_\theta} F^m(\theta_j) \bar{P}_n^m(\cos \theta_j) w_j $$
Usage
Getting started
The main functionality of torch_harmonics
is provided in the form of torch.nn.Modules
for composability. A minimum example is given by:
import torch
import torch_harmonics as th
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
nlat = 512
nlon = 2*nlat
batch_size = 32
signal = torch.randn(batch_size, nlat, nlon)
# transform data on an equiangular grid
sht = th.RealSHT(nlat, nlon, grid="equiangular").to(device)
coeffs = sht(signal)
To enable scalable model-parallelism, torch-harmonics
implements a distributed variant of the SHT located in torch_harmonics.distributed
.
Cite us
If you use torch-harmonics
in an academic paper, please cite [1]
@misc{bonev2023spherical,
title={Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere},
author={Boris Bonev and Thorsten Kurth and Christian Hundt and Jaideep Pathak and Maximilian Baust and Karthik Kashinath and Anima Anandkumar},
year={2023},
eprint={2306.03838},
archivePrefix={arXiv},
primaryClass={cs.LG}
}
References
[1] Bonev B., Kurth T., Hundt C., Pathak, J., Baust M., Kashinath K., Anandkumar A.; Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere; arXiv 2306.0383, 2023.
[2] Schaeffer N.; Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations; G3: Geochemistry, Geophysics, Geosystems, 2013.
[3] Wang B., Wang L., Xie Z.; Accurate calculation of spherical and vector spherical harmonic expansions via spectral element grids; Adv Comput Math, 2018.
Project details
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Hashes for torch_harmonics-0.6.2-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | dfeceb7c3d4ea5d61ba6f629c7881333d7fb0289925c4cab6c4c4da3377b7003 |
|
MD5 | 094118ade146db5a4f54c9692a26f011 |
|
BLAKE2b-256 | e4810bbee74716fcb547f5df156424d4e90646991134080adec490f89f5c59e0 |