moljax: GPU-accelerated method of lines for stiff reaction-diffusion PDEs with FFT preconditioning

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.