facebookresearch/fmmax

FMMAX is a tool from Facebook Research that simulates how light behaves when it passes through structured, layered materials, like photonic crystals, metasurfaces, or the tiny patterns found on advanced LEDs. If you are designing a nanostructured optical device and need to predict how it will scatter, transmit, or absorb light, this library computes those electromagnetic fields for you. Under the hood, it implements a technique called the Fourier modal method (FMM), which is a standard approach in computational photonics. The method works by breaking a periodic structure into flat layers, representing the light's field within each layer as a sum of Fourier components, and then stitching the layers together using a scattering matrix. This makes it relatively efficient for certain geometries compared to fully numerical solvers. The implementation is built on JAX, which means it can run on GPUs for speed and can automatically compute gradients, useful if you want to optimize a design rather than just analyze one. The project stands out for a few capabilities beyond the basics. It supports Brillouin zone integration, which lets you model localized light sources (like a single emitter or a focused beam) inside a periodic structure, rather than only uniform plane waves. It also includes advanced "vector FMM" formulations that improve simulation accuracy for tricky geometries, and it handles anisotropic and magnetic materials, broadening the range of devices you can model. The likely users are optics researchers and engineers working on things like metasurface-enhanced micro-LEDs, photonic crystal slabs, or other nanostructured optical components. The automatic differentiation support makes it particularly suited for inverse design, where you specify desired optical behavior and let an optimizer search for the structure that produces it. The project is installable via pip and released under the MIT license.