jejjohnson / jejjohnson/spectraldiffx

perf: ultra-spherical (banded) Chebyshev solvers — O(N) solves

Open
#66 1 comment 0 reactions 0 assignees View on GitHub
enhancement
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.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.