bispectrum

An open-source, fully unit-tested PyTorch library that implements selective G-bispectra for seven group actions as differentiable torch.nn.Modules, ready to plug into ML pipelines and deep learning architectures.

The G-bispectrum is a principled complete invariant of a signal — it retains all information up to the group action. Selectivity reduces cost from O(|G|²) to O(|G|) for finite groups, and from O(L³) to Θ(L²) coefficients for SO(3) on S².

Supported Groups

SO2onS1 is the continuous-n limit of CnonCn and shares its implementation.

API

Every module exposes a uniform interface:

• forward(f) — selective (default) or full bispectral invariants
• fourier(f) — group Fourier coefficients
• invert(beta) — signal reconstruction up to group-action indeterminacy (where available)

Modules default to O(|G|) selective coefficients; pass selective=False for the full O(|G|²) set. CG matrices, DFT kernels, and Bessel roots are precomputed as non-learnable buffers. Dependencies: PyTorch, NumPy, and torch_harmonics (for SO3onS2).

Benchmarks

Median wall-clock on a single NVIDIA H100 80 GB GPU (batch=16, torch.utils.benchmark):

Installation

pip install bispectrum

Or with uv:

uv pip install bispectrum

Development

git clone https://github.com/geometric-intelligence/bispectrum.git
cd bispectrum
uv pip install -e ".[dev]"
pre-commit install

Pre-commit Hooks

This project uses pre-commit for code quality checks. After installing dev dependencies, the hooks run automatically on each commit.

# Run hooks on all files (useful for first-time setup or CI)
pre-commit run --all-files

# Run a specific hook
pre-commit run ruff --all-files

# Update hooks to latest versions
pre-commit autoupdate

Usage

Cyclic group on itself

from bispectrum import CnonCn

bsp = CnonCn(n=8)
f = torch.randn(4, 8)       # signal on Z/8Z
beta = bsp(f)                # shape (4, 8), complex64
f_rec = bsp.invert(beta)     # reconstructed up to cyclic shift

SO(3) on the 2-sphere

from bispectrum import SO3onS2

# Selective bispectrum: O(L²) entries with generic completeness
bsp = SO3onS2(lmax=5, nlat=64, nlon=128, selective=True)
f = torch.randn(1, 64, 128)  # signal on S²
beta = bsp(f)                 # shape (1, 35), complex64

# Full bispectrum: O(L³) entries
bsp_full = SO3onS2(lmax=5, nlat=64, nlon=128, selective=False)
beta_full = bsp_full(f)       # shape (1, 69), complex64

Octahedral group

from bispectrum import OctaonOcta

bsp = OctaonOcta()
f = torch.randn(4, 24)       # signal on O (|O| = 24)
beta = bsp(f)                 # shape (4, 172), complex64
f_rec = bsp.invert(beta)      # reconstructed up to group action

License

MIT