Jaxwell: GPU-accelerated, differentiable 3D iterative FDFD electromagnetic solver
Jaxwell is JAX + Maxwell: an iterative solver for solving the finite-difference frequency-domain Maxwell equations on NVIDIA GPUs. Jaxwell is differentiable and fits seamlessly in the JAX ecosystem, enabling nanophotonic inverse design problems to be cast as ML training jobs and take advantage of the tsunami of innovations in ML-specific hardware, software, and algorithms.
Jaxwell is a finite-difference frequency-domain solver that finds solutions to the time-harmonic form of Maxwell's equations, specifically:
for the electric field E via the API
x, err = jaxwell.solve(params, z, b)where E → x, ω²ε → z, -iωJ → b, params controls how the solve proceeds iteratively, and err is the error in the solution.
Following meep, Jaxwell uses
dimensionless units,
and assumes μ = 1 everywhere,
and implements Shin's stretched-coordinate perfectly matched layers (SC-PML)
for absorbing boundary conditions.
You can install this version of Jaxwell with pip install git+https://github.com/djps/jaxwell.git
but the easiest way to get started is to go straight to the example from the original repository
colaboratory notebook.
References:
- PMLs and diagonalization: [Shin2012] W. Shin and S. Fan. “Choice of the perfectly matched layer boundary condition for frequency-domain Maxwell's equations solvers.” Journal of Computational Physics 231 (2012): 3406–31
- COCG algorithm: [Gu2014] X. Gu, T. Huang, L. Li, H. Li, T. Sogabe and M. Clemens, "Quasi-Minimal Residual Variants of the COCG and COCR Methods for Complex Symmetric Linear Systems in Electromagnetic Simulations," in IEEE Transactions on Microwave Theory and Techniques, vol. 62, no. 12, pp. 2859-2867, Dec. 2014
- BiCGStabL algorithm: [Sleijpen1993] G. Sleijpen and D. R. Fokkema. "BiCGstab(l) for linear equations involving unsymmetric matrices with complex spectrum." Electronic Transactions on Numerical Analysis vol 1, pp. 11-32 (1993).