Documentation

[azrael417]

Summary

Adds a full Sphinx documentation site for torch-harmonics covering API reference, user guides, tutorials, and build instructions.

API reference

• Serial layers: all public modules (RealSHT, InverseRealSHT, SpectralConvS2, DiscreteContinuousConvS2, AttentionS2, NeighborhoodAttentionS2, ResampleS2, QuadratureS2, GaussianRandomFieldS2, etc.)
• Distributed layers: all distributed counterparts with a coverage table mapping serial → distributed availability
• Distributed helpers: setup/teardown (init, finalize), process-group queries, compute_split_shapes, split_tensor_along_dim
• Distributed primitives: docstrings added to all 16 internal communication ops (transpose, copy/reduce, scatter/gather, halo exchange, tensor reshaping), documented for advanced users building custom distributed operators
• Utilities: quadrature helpers (precompute_latitudes, precompute_longitudes) with grid type comparison table, plotting (plot_sphere, imshow_sphere), truncation (truncate_sht), and the config debug toggle

User guides

• Spherical harmonic transforms, spectral convolutions, DISCO convolutions, spherical attention (pre-executed notebooks)
• Distributed spherical transforms (Markdown guide with process-grid diagrams, code examples, and internal SHT walkthrough)

Tutorials

• All 13 notebooks included and organized into three groups: Fundamentals, Applications, and Advanced topics

Build & install documentation

• Prebuilt wheels (NVIDIA PyPI + PyPI), building from source, Docker
• Build environment variables table (TORCH_HARMONICS_BUILD_CUDA_EXTENSION, NATIVE_CPU_ARCH, ENABLE_OPENMP, DEBUG, PROFILE, TORCH_CUDA_ARCH_LIST)
• Tensor core tip: use 9.0a / 10.0a arch suffixes for layers like DISCO that benefit from wgmma instructions

Benchmarking

• New documentation page under Getting started covering benchmarks/run.py usage, saving baselines, and cross-architecture comparison
• Added --reference-arch flag to benchmark runner for comparing results across GPU generations

Docstring quality

• All docstrings with :math: LaTeX use r""" raw strings
• Parameter references use param (RST literal) consistently
• Expanded docstrings across core modules (sht.py, quadrature.py, spectral_convolution.py, convolution.py, attention.py, resample.py, truncation.py, random_fields.py, filter_basis.py)

Other changes

• Renamed check_plotting_dependencies → _check_plotting_dependencies (internal helper)
• Dockerfiles updated to use 9.0a / 10.0a arch suffixes for tensor core support
• Fixed non-consecutive header levels in partial_derivatives.ipynb and shallow_water_equations.ipynb
• Added title cells to conditioning_sht.ipynb and stanford_2d3ds_dataset.ipynb
• Standardized all notebook titles to sentence case for consistency (e.g. "Differentiable Shallow Water Equations" → "Differentiable shallow water equations") and renamed "Derivatives of θ and φ" to "Polar and azimuthal derivatives".

Test plan

• make html in docs/ builds cleanly with no warnings
• Spot-check rendered pages: API reference, distributed guide, tutorials index, install page
• pre-commit run --all-files passes
• Existing test suite unaffected (python -m pytest --ignore-glob="**/test_distributed_*" ./tests/)
• Verify benchmark --reference-arch flag works: python benchmarks/run.py --reference-csv <file> --reference-arch <arch>

[bonevbs]

A few comments:

• for the plots in the examples can we use plot_sphere instead of imshow?
• can we rename the headings "Attention" to "Attention Mechanism"?
• in the bibliography, can we change the order of papers based on relevance? I would start with Nathanael Schaeffer's paper which is the fundamental SHT algorithm we implement.