jejjohnson / jejjohnson/spectraldiffx
perf: FFT-based Chebyshev differentiation (O(N log N))
- Dominant language
- Python
- Stars
- 1
- Forks
- 0
- PR merge metrics
- No merged PRs in 30d
Description
## Summary
Replace the O(N^2) dense matrix-vector product `D @ u` with an O(N log N) FFT-based algorithm for Chebyshev differentiation. This is the standard approach in Trefethen (2000, Chapter 8) and what production codes (Dedalus, SpectralDNS, Chebfun) use.
## Mathematical Background
### Key insight: Chebyshev ↔ Fourier via cosine substitution
A function u(x) on Chebyshev-Gauss-Lobatto points x_j = cos(j*pi/N) can be written as:
```
u(x_j) = u(cos(theta_j)) = g(theta_j) where theta_j = j*pi/N
```
The function g(theta) is periodic on [0, 2*pi] and even, so it has a cosine series:
```
g(theta) = sum_{k=0}^{N} a_k * cos(k*theta)
```
where the a_k are the **Chebyshev coefficients** of u, computed via DCT-I:
```
a_k = DCT-I(u)_k # O(N log N)
```
### Differentiation in coefficient space
The derivative du/dx in Chebyshev coefficient space satisfies the recurrence:
```
c_k * a'_k = 2*k * a_k + a'_{k+2} for k = N-1, N-2, ..., 1, 0
```
where c_0 = 2, c_k = 1 for k >= 1, and a'_N = a'_{N-1} = 0 (starting values).
This is a **backward recurrence** (O(N) operations), not a matrix multiply.
### Full algorithm
```
Step 1: u(x_j) → a_k via DCT-I # O(N log N)
Step 2: a_k → a'_k via backward recurrence # O(N)
Step 3: a'_k → u'(x_j) via IDCT-I # O(N log N)
Total: O(N log N)
```
vs. current: `D @ u` = O(N^2)
## Pseudocode
```python
def cheb_diff_fft(u: Array) -> Array:
"""Chebyshev derivative via FFT. O(N log N)."""
N = len(u) - 1
# Step 1: Physical values → Chebyshev coefficients via DCT-I
# The DCT-I of u at GL points gives the Chebyshev coefficients
a = dct(u, type=1) / N # normalize
# a[0] /= 2; a[N] /= 2 (endpoint corrections)
# Step 2: Differentiate in coefficient space (backward recurrence)
a_prime = jnp.zeros(N + 1)
a_prime = a_prime.at[N - 1].set(2 * N * a[N])
# Backward recurrence: c_k * a'_k = 2*k*a_k + a'_{k+2}
def body(carry, k):
a_prime_kp2 = carry
a_prime_k = 2 * k * a[k] + a_prime_kp2
# For k > 0, c_k = 1; for k = 0, c_0 = 2
return a_prime_k, a_prime_k
# Use jax.lax.scan for the backward loop (GPU-friendly)
_, a_prime_vals = jax.lax.scan(body, a_prime[N-1], jnp.arange(N-2, -1, -1))
a_prime = a_prime.at[:N-1].set(a_prime_vals)
a_prime = a_prime.at[0].set(a_prime[0] / 2) # c_0 = 2 correction
# Step 3: Chebyshev coefficients → physical values via IDCT-I
u_prime = idct(a_prime, type=1)
return u_prime
```
### Higher derivatives
For d^2u/dx^2, apply the recurrence twice. Or use the second-derivative recurrence directly:
```
c_k * a''_k = 2*k * a'_k + a''_{k+2}
```
This avoids the intermediate IDCT/DCT round-trip.
## Back-of-Envelope: Speedup
| N | Dense D@u | FFT-based | Speedup |
|---|-----------|-----------|---------|
| 32 | 1,089 flops | ~160 flops | 7x |
| 64 | 4,225 flops | ~384 flops | 11x |
| 128 | 16,641 flops | ~896 flops | 19x |
| 256 | 66,049 flops | ~2,048 flops | 32x |
| 512 | 263,169 flops | ~4,608 flops | 57x |
(FFT cost ~ 5*N*log2(N), recurrence ~ 3*N)
## Implementation Plan
1. Add `cheb_diff_fft_1d(u)` as a standalone function
2. Add `cheb_diff2_fft_1d(u)` for second derivative (avoids two recurrences)
3. Option on `SpectralDerivative1D/2D`: `method="fft"` vs `method="matrix"` (default)
4. For 2D: apply FFT-based diff along each axis independently
## JAX Considerations
- The backward recurrence can use `jax.lax.scan` with reversed indices — this compiles to efficient GPU code
- DCT-I is available via `jax.scipy.fft.dct(type=1)` (or our own `dct` wrapper)
- The entire pipeline (DCT → scan → IDCT) fuses into a single XLA computation
## Acceptance Criteria
- [ ] `cheb_diff_fft_1d` matches `D @ u` to machine precision
- [ ] `cheb_diff2_fft_1d` matches `D2 @ u` to machine precision
- [ ] Benchmark showing O(N log N) scaling
- [ ] Integration with SpectralDerivative1D/2D as optional `method` parameter
- [ ] JIT-compatible, vmap-compatible
## References
- Trefethen, "Spectral Methods in MATLAB" (2000), Chapter 8, Program 20
- Press et al., "Numerical Recipes" (2007), Section 5.9
- Canuto et al., "Spectral Methods: Fundamentals" (2006), Section 2.4
## Part of
- Epic #61
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.