moljax: GPU-accelerated method of lines for stiff reaction-diffusion PDEs with FFT preconditioning
Description
We present moljax, an open-source JAX library for GPU-accelerated method-of-lines simulation of stiff reaction-diffusion PDEs on structured grids. The library combines three capabilities absent from existing Python PDE tools: (i) JIT-compiled adaptive time stepping with accept/reject control flow on GPU; (ii) matrix-free Newton–Krylov solvers using AD for exact Jacobian–vector products; and (iii) FFT/DST/DCT spectral operators with physics-aware preconditioning. Controlled benchmarks on Gray–Scott, Schnakenberg, and Brusselator systems show IMEX and ETDRK4 integrators achieving 10–17 × speedup over explicit RK4 on the same GPU and spatial discretization. Work-precision analysis reveals that ETDRK4 is the only method achieving monotone pointwise convergence for pattern-forming systems. A tubular reactor benchmark demonstrates 18–40 × speedup over Diffrax using identical spatial discretization. Performance claims are supported by controlled comparisons that hold hardware and spatial discretization fixed where attribution is intended. Code: https://github.com/gogipav14/moljax.