jejjohnson / jejjohnson/spectraldiffx
perf: pre-factored LU for ChebyshevHelmholtzSolver
- Dominant language
- Python
- Stars
- 1
- Forks
- 0
- PR merge metrics
- No merged PRs in 30d
Description
## Summary
The `ChebyshevHelmholtzSolver1D.solve()` currently performs a **fresh O(N^3) LU decomposition on every call**, even when the operator matrix `A = D^2 - alpha*I` hasn't changed. For time-stepping (where `solve()` is called thousands of times with the same `alpha`), pre-factoring reduces the per-call cost from O(N^3) to O(N^2).
## Current Implementation
```python
# solvers.py line 107-120 — every call does this:
def solve(self, f, alpha=0.0, bc_left=0.0, bc_right=0.0):
A = D @ D - alpha * jnp.eye(N + 1) # O(N^3) matmul (redundant!)
# ... boundary row modifications ...
return jnp.linalg.solve(A, b) # O(N^3) LU decomposition
```
**Cost per call:** O(N^3) + O(N^3) = O(N^3)
**For 1000 time steps at N=128:** 1000 * 4.3M = 4.3 billion flops wasted on redundant factorizations
## Proposed Implementation
### Option A: Store LU factorization as module attribute
```python
class ChebyshevHelmholtzSolver1D(eqx.Module):
grid: ChebyshevGrid1D
alpha: float = 0.0
_lu: Array = eqx.field(init=False) # LU factors
_piv: Array = eqx.field(init=False) # pivot indices
def __post_init__(self):
D = self.grid.D
N = self.grid.N
A = D @ D - self.alpha * jnp.eye(N + 1)
# Enforce Dirichlet BC rows (homogeneous default)
A = A.at[0, :].set(0.0).at[0, 0].set(1.0)
A = A.at[N, :].set(0.0).at[N, N].set(1.0)
self._lu, self._piv = jax.scipy.linalg.lu_factor(A)
def solve(self, f, bc_left=0.0, bc_right=0.0):
b = f.at[0].set(bc_right).at[self.grid.N].set(bc_left)
return jax.scipy.linalg.lu_solve((self._lu, self._piv), b)
```
**Cost per call:** O(N^2) (forward/back substitution only)
**Speedup:** N/const ~ 10-100x for typical N
### Option B: Lazy factorization with caching
```python
def solve(self, f, alpha=0.0, ...):
# Only re-factor when alpha changes
A = self._get_or_build_A(alpha)
lu, piv = self._get_or_factor(A)
return jax.scipy.linalg.lu_solve((lu, piv), b)
```
**Recommendation:** Option A is simpler and fits the equinox Module pattern (immutable after init). Users who need different `alpha` values create separate solver instances.
## Back-of-Envelope Calculations
| N | LU factor (once) | LU solve (per call) | Speedup (1000 calls) |
|---|-------------------|---------------------|----------------------|
| 32 | 72K flops | 2.2K flops | 33x |
| 64 | 549K flops | 8.5K flops | 65x |
| 128 | 4.3M flops | 33K flops | 130x |
| 256 | 34M flops | 132K flops | 258x |
## Backward Compatibility
The current `alpha` parameter on `solve()` would move to the constructor. This is a breaking change but aligns with the equinox Module pattern where solver configuration is set at construction time.
Alternative: keep `alpha` on `solve()` but cache the factorization and invalidate when `alpha` changes.
## Acceptance Criteria
- [ ] LU factorization computed once at construction
- [ ] `solve()` uses O(N^2) forward/back substitution
- [ ] Benchmark showing speedup for repeated solves
- [ ] Inhomogeneous Dirichlet BCs still work (modify RHS only, not A)
- [ ] JIT-compatible
## Part of
- Epic #61
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.