jejjohnson / jejjohnson/spectraldiffx

perf: pre-factored LU for ChebyshevHelmholtzSolver

Open
#62 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 `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.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.