Package: torch-harmonics
torch-harmonicsOverview|Installation|More information|Getting started|Contributors|Cite us|ReferencesOverviewtorch-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].InstallationDownload directyly from PyPI:pipinstalltorch-harmonicsBuild in your environment using the Python package:gitclonegit@github.com:NVIDIA/torch-harmonics.gitcdtorch-harmonics
pipinstall-e.Alternatively, use the Dockerfile to build your custom container after cloning:gitclonegit@github.com:NVIDIA/torch-harmonics.gitcdtorch-harmonics
dockerbuild.-ttorch_harmonics
dockerrun--gpusall-it--rm--ipc=host--ulimitmemlock=-1--ulimitstack=67108864torch_harmonicsMore about torch-harmonicsSpherical harmonicsThespherical harmonicsare 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 harmonics up to degree 5Spherical harmonic transformThe 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 transformThe 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 $x = \cos \theta$ and then approximated by the sum$$
\hat f_l^m = \sum_{j=1}^{N_\theta}  \hat{f}^m(\arccos x_j) P_l^m(x_j) w_j.
$$Here, $x_j \in [-1,1]$ are the quadrature nodes with the respective quadrature weights $w_j$.Getting startedThe main functionality oftorch_harmonicsis provided in the form oftorch.nn.Modulesfor composability. A minimum example is given by:importtorchimporttorch_harmonicsasthdevice=torch.device('cuda'iftorch.cuda.is_available()else'cpu')nlat=512nlon=2*nlatbatch_size=32signal=torch.randn(batch_size,nlat,nlon)# transform data on an equiangular gridsht=th.RealSHT(nlat,nlon,grid="equiangular").to(device)coeffs=sht(signal)To enable scalable model-parallelism,torch-harmonicsimplements a distributed variant of the SHT located intorch_harmonics.distributed.Detailed usage of torch-harmonics, alongside helpful analysis provided in a series of notebooks:Getting startedQuadratureVisualizing the spherical harmonicsSpectral fitting vs. SHTConditioning of the GramianSolving the Helmholtz equationSolving the shallow water equationsTraining Spherical Fourier Neural OperatorsRemarks on automatic mixed precision (AMP) supportNote that torch-harmonics uses Fourier transforms fromtorch.fftwhich in turn uses kernels from the optimizedcuFFTlibrary. This library supports fourier transforms offloat32andfloat64(i.e.singleanddoubleprecision) tensors for all input sizes. Forfloat16(i.e.halfprecision) andbfloat16inputs however, the dimensions which are transformed are restricted to powers of two. Since data is converted to one of these reduced precision floating point formats whentorch.cuda.amp.autocastis used, torch-harmonics will issue an error when the input shapes are not powers of two. For these cases, we recommend disabling autocast for the harmonics transform specifically:importtorchimporttorch_harmonicsasthsht=th.RealSHT(512,1024,grid="equiangular").cuda()withtorch.cuda.amp.autocast(enabled=True):# do some AMP converted math herex=some_math(x)# convert tensor to float32x=x.to(torch.float32)# now disable autocast specifically for the transform,# making sure that the tensors are not converted# back to reduced precision internallywithtorch.cuda.amp.autocast(enabled=False):xt=sht(x)# continue operating on the transformed tensorxt=some_more_math(xt)Depending on the problem, it might be beneficial to upcast data tofloat64instead offloat32precision for numerical stability.ContributorsBoris Bonev(bbonev@nvidia.com),Thorsten Kurth(tkurth@nvidia.com),Christian Hundt(chundt@nvidia.com),Nikola Kovachki(nkovachki@nvidia.com),Jean Kossaifi(jkossaifi@nvidia.com)Cite usIf you usetorch-harmonicsin 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.