jejjohnson / jejjohnson/spectraldiffx

perf: FFT-based Chebyshev differentiation (O(N log N))

Open
#64 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

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.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.