jejjohnson / jejjohnson/spectraldiffx
Epic: Chebyshev Performance Optimizations
- Dominant language
- Python
- Stars
- 1
- Forks
- 0
- PR merge metrics
- No merged PRs in 30d
Description
## Summary
The current Chebyshev implementation uses dense matrices and dense linear algebra throughout. While mathematically clean and correct, several well-known techniques can dramatically improve performance — especially for the Helmholtz solver (O(N^3) per call) and differentiation (O(N^2) per call).
## Sub-Issues
| # | Title | Speedup | Effort | Priority |
|---|-------|---------|--------|----------|
| #62 | Pre-factored LU for Helmholtz solver | 2-10x for repeated solves | Small | P0 |
| #64 | FFT-based Chebyshev differentiation | O(N^2) -> O(N log N) | Medium | P1 |
| #65 | Scan-based matrix-vector products for time-stepping | Reduced overhead, XLA fusion | Small | P1 |
| #66 | Ultra-spherical / banded Chebyshev solvers | O(N^3) -> O(N) solves | Large | P2 |
## Current Bottlenecks
### Differentiation: O(N^2) dense matmul
\`\`\`python
# Current: D @ u where D is (N+1) x (N+1) fully dense
du_dx = D @ u # O(N^2) per derivative
d2u_dx2 = D2 @ u # O(N^2), D2 precomputed
\`\`\`
### Helmholtz solver: O(N^3) dense LU per call
\`\`\`python
# Current: fresh LU decomposition every call
A = D @ D - alpha * I # O(N^3) matmul
u = jnp.linalg.solve(A, b) # O(N^3) LU
\`\`\`
| N | Matrix size | Dense LU flops | Time @ 1 TFLOP/s |
|---|------------|----------------|-------------------|
| 32 | 33 x 33 | ~72K | 72 ns |
| 64 | 65 x 65 | ~549K | 0.5 us |
| 128 | 129 x 129 | ~4.3M | 4.3 us |
| 256 | 257 x 257 | ~34M | 34 us |
| 512 | 513 x 513 | ~270M | 270 us |
The O(N^3) scaling becomes dominant for N > 128, especially in time-stepping loops where solve() is called thousands of times.
## References
- Trefethen, "Spectral Methods in MATLAB" (2000), Chapters 6-8
- Olver & Townsend, "A fast and well-conditioned spectral method" (2013) — ultra-spherical
- Driscoll, Hale & Trefethen, "Chebfun and numerical quadrature" (2014)
- Burns et al., "Dedalus: A flexible framework for numerical simulations with spectral methods" (2020)
## Depends On
- #28 (Chebyshev extensions) — these optimizations complement the feature additions
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.