jejjohnson / jejjohnson/spectraldiffx
perf: ultra-spherical (banded) Chebyshev solvers — O(N) solves
- Dominant language
- Python
- Stars
- 1
- Forks
- 0
- PR merge metrics
- No merged PRs in 30d
Description
## Summary
The current Chebyshev Helmholtz solver uses dense O(N^3) LU decomposition. The **ultra-spherical spectral method** (Olver & Townsend 2013) reformulates the problem so that the operator matrix is **almost-banded** (banded + low-rank boundary rows), enabling O(N) solves via the Woodbury formula.
This is the approach used by Dedalus, ApproxFun.jl, and Chebfun for large-N Chebyshev problems. It represents the theoretical optimal complexity for spectral methods.
## Mathematical Background
### The problem with the standard approach
The standard Chebyshev tau method works in the T_k (Chebyshev-T) basis:
```
u(x) = sum_{k=0}^{N} a_k T_k(x)
```
The second derivative d^2u/dx^2 in this basis produces a **dense** matrix because:
```
T_k''(x) = k * sum_{j=0,2,...}^{k-2} (something) * T_j(x)
```
Each mode k couples to ALL lower modes j < k. Result: dense upper-triangular D^2 matrix.
### Ultra-spherical key insight
Instead of expressing d^2u/dx^2 in the T_k basis, express it in the **C_k^(2)** (ultraspherical / Gegenbauer) basis:
```
d^2/dx^2 [T_k(x)] = k(k-1)/2 * C_{k-2}^(2)(x)
```
This is a **diagonal** operation! The second derivative maps T_k to a single C_{k-2}^(2)} mode.
The full operator (d^2/dx^2 - alpha) in mixed bases becomes:
```
D2: T_k → C_{k-2}^(2)} # diagonal (bandwidth 0)
S1: T_k → C_k^(1)} # tridiagonal (bandwidth 1)
S2: C_k^(1)} → C_k^(2)} # tridiagonal (bandwidth 1)
```
The operator matrix `A = D2 - alpha * S2 * S1` is **pentadiagonal** (bandwidth 2).
### Adding boundary conditions
Boundary conditions add 2 dense rows (one per BC):
```
T_k(1) = 1 for all k → dense row [1, 1, 1, ..., 1]
T_k(-1) = (-1)^k → dense row [1, -1, 1, -1, ...]
```
### Almost-banded structure
The full system is:
```
[ BC row 1 (dense) ] [ a_0 ] [ bc_left ]
[ BC row 2 (dense) ] [ a_1 ] [ bc_right ]
[ ] [ a_2 ] [ ]
[ pentadiagonal A ] [ ... ] = [ rhs ]
[ ] [ a_N-2 ] [ ]
[ ] [ a_N-1 ] [ ]
[ ] [ a_N ] [ ]
```
This is: **banded matrix + 2 dense rows** = almost-banded.
### Woodbury formula solve
The almost-banded system can be solved in O(N) using the Woodbury identity:
```
(B + U*V^T)^{-1} = B^{-1} - B^{-1}*U*(I + V^T*B^{-1}*U)^{-1}*V^T*B^{-1}
```
where B is the banded part (O(N) to solve) and U*V^T represents the 2 dense BC rows (rank-2 update, O(1) inner matrix).
**Total cost:** O(N) for the banded solve + O(N) for the rank-2 correction = **O(N)**.
## Pseudocode
```python
def ultraspherical_helmholtz_solve(rhs_coeffs, alpha, bc_left, bc_right, N):
"""Solve (d^2/dx^2 - alpha) u = f with Dirichlet BCs in O(N).
All operations are in coefficient space (no grid-point transforms needed).
"""
# Step 1: Convert RHS from T basis to C^(2) basis
# S0->1 and S1->2 are tridiagonal conversion matrices
f_c2 = S12 @ (S01 @ rhs_coeffs) # O(N), tridiagonal matmuls
# Step 2: Build pentadiagonal operator (banded, stored as diagonals)
# D2[k, k-2] = k(k-1)/2 (diagonal in ultraspherical)
# alpha * S12 @ S01 contributes bandwidth-2 entries
# Total bandwidth: 2 (pentadiagonal)
bands = build_pentadiagonal(N, alpha) # O(N) construction
# Step 3: Solve banded system (ignoring BCs) — O(N)
a_particular = solve_banded(bands, f_c2) # O(N) via Thomas-like algorithm
# Step 4: Woodbury correction for BCs
# Solve banded system for each BC row: B^{-1} * e_bc
z1 = solve_banded(bands, bc_row_1) # O(N)
z2 = solve_banded(bands, bc_row_2) # O(N)
# 2x2 Schur complement solve
# [V^T * z1, V^T * z2] is 2x2
schur = [[dot(bc1, z1), dot(bc1, z2)],
[dot(bc2, z1), dot(bc2, z2)]]
correction = solve_2x2(schur, [bc_left - dot(bc1, a_particular),
bc_right - dot(bc2, a_particular)])
# Step 5: Combine
a_solution = a_particular + correction[0] * z1 + correction[1] * z2
# Step 6: Convert back to T basis if needed (for grid-point evaluation)
# Or keep in coefficient space for further spectral operations
return a_solution
```
## Back-of-Envelope: Speedup vs Dense LU
| N | Dense LU O(N^3) | Ultra-spherical O(N) | Speedup |
|---|-----------------|---------------------|---------|
| 32 | 72K flops | ~500 flops | 144x |
| 64 | 549K flops | ~1K flops | 549x |
| 128 | 4.3M flops | ~2K flops | 2,150x |
| 256 | 34M flops | ~4K flops | 8,500x |
| 512 | 270M flops | ~8K flops | 33,750x |
| 1024 | 2.1G flops | ~15K flops | 140,000x |
The improvement is transformative for large N and makes high-resolution Chebyshev practical.
## Implementation Complexity
This is the most complex optimization in the epic:
1. **Basis conversion matrices** S01 (T→C^(1)) and S12 (C^(1)→C^(2)): tridiagonal, well-known formulas
2. **Banded storage format**: JAX doesn't have built-in banded solvers — need to implement Thomas algorithm or use `jax.scipy.linalg.solve_banded` (if available)
3. **Woodbury correction**: standard linear algebra, O(1) overhead
4. **Coefficient ↔ grid-point transforms**: DCT-based, O(N log N)
### JAX Challenges
- JAX doesn't expose LAPACK's `dgbsv` (banded solver) directly
- Options: (a) implement Thomas algorithm in JAX, (b) use `jax.experimental.sparse`, (c) call LAPACK via `jax.pure_callback`
- The banded solve with `jax.lax.scan` (forward elimination + back substitution) is natural and GPU-friendly
## Acceptance Criteria
- [ ] Ultra-spherical conversion matrices (T → C^(1) → C^(2))
- [ ] Pentadiagonal operator construction
- [ ] O(N) banded solver (Thomas algorithm via scan)
- [ ] Woodbury correction for boundary conditions
- [ ] Benchmark: O(N) scaling confirmed
- [ ] Matches dense solver to machine precision
- [ ] Supports Helmholtz (alpha > 0) and Poisson (alpha = 0)
- [ ] 1D implementation first, 2D via tensor product
## References
- Olver & Townsend, "A fast and well-conditioned spectral method" SIAM Review 55(3), 2013
- Townsend & Olver, "The automatic solution of PDEs using ultraspherical spectral methods" 2014
- Burns et al., "Dedalus" (2020) — production implementation
- Julien & Watson, "Efficient multi-dimensional solution of PDEs using Chebyshev spectral methods" JCP 132, 2013
## Part of
- Epic #61
- Depends on: #64 (FFT-based derivatives for coefficient ↔ grid transforms)
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.