torch-harmonics

Differentiable signal processing on the sphere for PyTorch.

torch-harmonics implements differentiable spherical harmonic transforms (SHT), discrete-continuous (DISCO) convolutions, spherical attention, and related operators as PyTorch modules. All operators are autograd-compatible and run on CPU and GPU, with optional custom CUDA kernels for the performance-critical paths.

Getting started

• Installation
  • Prebuilt wheels (GPU and CPU)
  • PyPI (CPU only)
  • Building from source
  • Docker
  • Requirements
  • Building the documentation
• Benchmarking
  • Running benchmarks
  • Saving results
  • Comparing against reference numbers
  • Reference results
• Tutorials
  • Getting started with torch-harmonics
  • Visualizing the spherical harmonics
  • Quadrature
  • Polar and azimuthal derivatives
  • Analyzing the gradients of the SHT
  • Helmholtz equation
  • Differentiable shallow water equations
  • Spherical Fourier neural operators
  • Stanford 2D-3D-S dataset
  • Filter basis functions
  • Resampling signals on the 2-sphere
  • Conditioning of the SHT
  • Spherical attention equivariance test

User guide

• Spherical harmonic transforms
• Spectral convolutions on the sphere
• DISCO convolutions on the sphere
• Attention on the sphere
• Distributed spherical transforms

API reference

• Serial layers
  • Spherical harmonic transforms
  • Convolutions
  • Filter basis
  • Attention mechanism
  • Resampling and quadrature
  • Random fields
• Distributed helpers
  • Setup and teardown
  • Process group accessors
  • Data partitioning
• Distributed layers
  • Coverage
  • Layer reference
• Distributed primitives
  • Transpose (all-to-all)
  • Copy / Reduce
  • Scatter / Gather
  • Halo exchange
  • Tensor reshaping
• Utilities
  • Quadrature
  • Plotting
  • Truncation
  • Debugging

Quick example

import torch
import torch_harmonics as th

# forward / inverse real spherical harmonic transform on an equiangular grid
sht = th.RealSHT(nlat=128, nlon=256, grid="equiangular")
isht = th.InverseRealSHT(nlat=128, nlon=256, grid="equiangular")

signal = torch.randn(1, 128, 256)
coeffs = sht(signal)          # -> spherical harmonic coefficients
reconstructed = isht(coeffs)  # -> back to grid space

Bibliography

• References

Indices

• Index

• Module Index