spharmgrid

Spherical harmonic tools for filtering, regridding, and kinematics in atmospheric science with Xarray, JAX, and PyTorch.

spharmgrid (spherical harmonic gridding) implements spherical harmonic filtering, regridding, differential operators, and atmospheric kinematics for global Xarray fields, JAX arrays, and PyTorch tensors. It computes relative vorticity, divergence, streamfunction, velocity potential, Helmholtz decomposition, and inverse wind transforms. The Xarray API uses DUCC (ducc0) for spherical harmonic transforms; the optional JAX API uses S2FFT, and the optional PyTorch API uses torch-harmonics.

Supported grids are full rectangular Gauss–Legendre (GL) and Clenshaw–Curtis (CC) grids.

Install

Install with either pip, uv, or conda:

pip install spharmgrid

uv add spharmgrid

conda install -c conda-forge spharmgrid

Optional extras are:

• spharmgrid[dask] — Dask-backed lazy execution;
• spharmgrid[cf] — optional cf-xarray coordinate discovery;
• spharmgrid[cli] — command-line NetCDF, Zarr, and GRIB I/O;
• spharmgrid[jax] — JAX arrays and S2FFT transforms;
• spharmgrid[torch] — PyTorch tensors and torch-harmonics transforms.

The spharmgrid.jax API requires JAX and S2FFT. See the JAX API documentation for supported GL and CC/MWSS dimensions, x64 and dtype requirements, and examples. For CPU use, spharmgrid[jax] installs JAX and S2FFT. For GPU or TPU use, install the appropriate JAX accelerator build first by following the official JAX installation instructions, then install spharmgrid[jax]; spharmgrid does not bundle or select CUDA or TPU builds. The JAX API currently requires JAX x64 mode and float64 spatial inputs; it does not enable x64 globally.

The spharmgrid.torch API can be installed with spharmgrid[torch]. See the PyTorch API documentation for installation instructions.

Install spharmgrid[cli,dask] to use the transforming CLI commands.

For a standalone command-line installation:

uv tool install "spharmgrid[cli,dask]"

For a project environment, install the CLI and Dask extras with either:

uv add "spharmgrid[cli,dask]"

pip install "spharmgrid[cli,dask]"

Quick start

Importing spharmgrid registers the .sg accessor on Xarray objects. This example applies a T6–42 spectral filter to a DataArray:

import xarray as xr
import spharmgrid

field = xr.open_dataarray("msl.nc")
filtered = field.sg.filter("T6-42")

For workflows with several spectral operations, analyze once and reuse the private coefficient representation:

spectral = field.sg.analyze()
large_scale = spectral.filter("T5-42").synthesize()
laplacian = spectral.laplacian().synthesize()

See the Quick start for regridding, atmospheric wind diagnostics, direct-function equivalents, and further examples.

For differentiable JAX workflows, configure JAX x64 mode before creating arrays and use spharmgrid.jax:

import jax

jax.config.update("jax_enable_x64", True)

import jax.numpy as jnp
import spharmgrid as sg
import spharmgrid.jax as sgj

grid = sg.gaussian_grid(64, 127)
field = jnp.ones((grid.nlat, grid.nlon))
filtered = sgj.filter(field, grid=grid, truncation="T42")

The last two array dimensions are latitude and longitude.

For differentiable PyTorch workflows, install spharmgrid[torch] and use spharmgrid.torch:

import torch
import spharmgrid as sg
import spharmgrid.torch as sgt

grid = sg.gaussian_grid(64, 128)
field = torch.randn(grid.nlat, grid.nlon)
filtered = sgt.filter(field, grid=grid, truncation="T42")

See the PyTorch API documentation for tensor dimensions, reusable torch.nn modules, device/autograd behavior, and PyTorch bandwidth limits.

Documentation

The optional CLI reads NetCDF, Zarr, and GRIB and writes NetCDF and Zarr. See the CLI documentation for installation and usage.

See the full documentation for grid requirements, coordinate handling, CF metadata, atmospheric kinematics, inverse transforms, zero-mode conventions, and command-line use.

Citation

If you use spharmgrid in research, please cite the software release DOI: 10.5281/zenodo.22559210. Citation metadata are also provided in CITATION.cff.

References

See the documentation References for the scientific literature and software cited by spharmgrid.

License

spharmgrid is distributed under the BSD 3-Clause License.