Performance comparison of AMJax against PyAMG on 2D Poisson problems of increasing size (n = 50, 100, 200, 500).
Experimental setup
| Parameter | Value |
|---|---|
| Problem | 2D Poisson on an n×n grid |
| Tolerance | 10⁻⁸ |
| Max solver iterations | 500 |
| Max cycle iterations | 250 (per solver call) |
| Grid sizes | n = 50, 100, 200, 500 |
| AMG hierarchies | Ruge-Stüben, Smoothed Aggregation, Root Node, Pairwise |
| Cycle types | V, F, W |
| Precision | f32, f64 |
| Devices | AMJax on GPU, PyAMG on CPU |
| Modes | single (1 RHS), vmap (k RHS batched) |
| Batch size (k) | 32, 64 |
| Methods | AMJax, AMJax + PCG, PyAMG, PyAMG + PCG |
| Timing | after JIT warm-up (first call excluded) |
| Residual metric | ||b − Ax|| / ||b|| |