facebookresearch / fmmax

Fourier modal method with Jax

FMMAX: Fourier Modal Method with Jax

FMMAX is a an implementation of the Fourier modal method (FMM) in JAX.

The FMM -- also known as rigorous coupled wave analysis (RCWA) -- is a semianalytical method that solves Maxwell's equations in periodic stratified media, where in-plane directions are treated with a truncated Fourier basis and the normal direction is handled by a scattering matrix approach [1999 Whittaker, 2012 Liu, 2020 Jin]. This allows certain classes of structures to be modeled with relatively low computational cost.

Our use of JAX enables GPU acceleration and automatic differentiation of FMM simulations. Besides these features, FMMAX is differentiated from other codes by its support for Brillouin zone integration, advanced vector FMM formulations which improve convergence, and anisotropic and magnetic materials.

Brillouin zone integration

Brillouin zone integration [2022 Lopez-Fraguas] allows modeling of localized sources in periodic structures. Check out the crystal example to see how we model a Gaussian beam incident upon a photonic crystal slab, or an isolated dipole embedded within the slab. The Gaussian beam fields are shown below.

Vector FMM formulations

Vector FMM formulations introduce local coordinate systems at each point in the unit cell, which are normal and tangent to all interfaces. This allows normal and tangent field components to be treated differently and improves convergence. FMMAX implements several vector formulations of the FMM, with automatic vector field generation based on functional minimization similar to [2012 Liu]. We implement the Pol, Normal, and Jones methods of that reference, and introduce a new Jones direct method which we have found to have superior convergence. These are supported also with anisotropic and magnetic materials. The vector_fields example computes vector fields by these methods for an example structure.

Anisotropic, magnetic materials

Our support of anisotropic, magnetic materials allows modeling of uniaxial perfectly matched layers. This is demonstrated in the metal_dipole example, which simulates in vaccuum located above a metal substrate. The resulting electric fields are whown below.

FMM Conventions

• The speed of light, vacuum permittivity, and vacuum permeability are all 1.
• Fields evolve in time as $\exp(-i \omega t)$.
• If $\mathbf{u}$ and $\mathbf{v}$ are the primitive lattice vectors, the unit cell is defined by the parallelogram with vertices at $\mathbf{0}$, $\mathbf{u}$, $\mathbf{u} + \mathbf{v}$, and $\mathbf{v}$.
• For quantities defined on a grid (such as the permittivity distribution of a patterned layer) the value at grid index (0, 0) corresponds to the value at physical location $\mathbf{0}$.
• The scattering matrix block $\mathbf{S}_{11}$ relates incident and transmitted forward-going fields, and other blocks have corresponding definitions. This differs from the convention e.g. in photonic integrated circuits.

Batching

Batched calculations are supported, and should be used where possible to avoid looping. The batch axes are the leading axes, except for the wave amplitudes and electromagnetic fields, where a trailing batch axis is assumed. This allows e.g. computing the transmission through a structure for multiple polarizations via a matrix-matrix operation (transmitted_amplitudes = S11 @ incident_amplitudes), rather than a batched matrix-vector operation.

Installation

FMMAX can be installed via pip:

pip install fmmax

For developers requiring a local installation, you will need to first clone this repository and then perform a local install from within the root directory using:

pip install -e ".[dev]"

The [dev] modifier specifies optional dependencies for developers which are listed in pyproject.toml.

Note: for this to work, it may be necessary to first update your pip installation using e.g. python3 -m pip install --upgrade pip.

Citing FMMAX

If you use FMMAX, please consider citing our paper,

@misc{schubert2023fourier,
      title={Fourier modal method for inverse design of metasurface-enhanced micro-LEDs}, 
      author={Martin F. Schubert and Alec M. Hammond},
      year={2023},
      eprint={2308.08573},
      archivePrefix={arXiv},
      primaryClass={physics.comp-ph}
}

License

FMMAX is licensed under the MIT license.

References

• [2012 Liu] V. Liu and S. Fan, S4: A free electromagnetic solver for layered structures structures, Comput. Phys. Commun. 183, 2233-2244 (2012).

• [1999 Whittaker] D. M. Whittaker and I. S. Culshaw, Scattering-matrix treatment of patterned multilayer photonic structures, Phys. Rev. B 60, 2610 (1999).

• [2020 Jin] W. Jin, W. Li, M. Orenstein, and S. Fan Inverse design of lightweight broadband reflector for relativistic lightsail propulsion, ACS Photonics 7, 9, 2350-2355 (2020).

• [2022 Lopez-Fraguas] E. Lopez-Fraguas, F. Binkowski, S. Burger, B. Garcia-Camara, R. Vergaz, C. Becker and P. Manley Tripling the light extraction efficiency of a deep ultraviolet LED using a nanostructured p-contact, Scientific Reports 12, 11480 (2022).

fmmax's Issues

Use "differentiable optimization trick" for backpropagation through tangent vector field calculation

Currently, we directly backpropagate through the tangent vector field calculation, which involves a Newton solve to find the minimum of a convex quadratic objective. It may be more efficient to define a custom gradient for this operation, in a manner similar to what is done for differentiable optimization.

Include tangent vector field as an attribute of the layer solve result

Improved PML profiles

The current PML profiles are appropriate e.g. for a finite difference grid, and result in sharp features at the edges of the unit cells, likely leading to ringing and perturbations to the permittivity even in the non-PML region. We should investigate this issue and consider better (smoother) spatial profiles for the PML.

Security alerts

The dependabot security tool identifies several issues with fmmax, which seem to be related to docs generation.

https://github.com/facebookresearch/fmmax/security/dependabot

Adding propagation conventions

We discuss a bit what propagation means in the context of scattering parameters, but it might be good to describe what that means in other contexts.

For example, when building a structure, one typically specifies a list of permittivities and a corresponding list of layer thicknesses. With many of the examples in the codebase (if not all), this list is interpreted as a stack assembled from top to bottom. In which case, all quantities (e.g. sources) that describe forward propagation, actually imply downward propagation, and all quantities that describe backward propagation, actually imply upward propagation. Furthermore, the start of layer, in this case, is the top of the layer (and the end is the bottom)

Of course, we can rearrange how we order things, and the forward/backward convention flips (which is why it's good that the codebase arguments are labeled using fwd and bwd, to maximize generalizability). But it would be nice to mention that this is indeed the case somewhere upfront.

Alternate BZ grid generation for better resource utilization

Currently, the brillouin_zone_in_plane_wavevector() function samples a uniform grid of BZ points by building an (N,M) grid, and sampling the centers of each sub region.

This approach requires NxM points. This may not be ideal if you want to have an arbitrary number of points that are still equidistant from each other (for example, 8 points within a square).

In the context of multi-gpu simulations, you typically want to make sure the product of your batch dimensions is a multiple of your available devices. When you're restricted to uniform grids, it can be tricky to make this work (without arbitrarily padding your arrays with zeros...)

One way around this is to do something like a circle-packing algorithm to see how one can fit an arbitrary number of points in your BZ cell. If your BZ cell is not a square (e.g. it's a rectangle parallelogram) you can do a change of basis on on the unit cell, pack your circles, pull out their centers, and then transform the centers back to the original basis.

The advantage of this is you still get a uniform sampling of the BZ with an arbitrary number of points. Of course, one then needs to figure out how to interpret these new points (i.e. what would an equivalent source look like).

Augment microlens array example to include conformal AR coating

It would be interesting (and probably difficult) to modify the microlens array example to include a conformal AR coating. This is particularly important when the index contrast is particularly high, and the injected fields have high spatial frequencies (they aren't described by a simple planewave). In these cases, significant energy is reflected back through the stack.

Additional example to study metal grating convergence

It would be useful to have additional examples which demonstrate the convergence characteristics of the various FMM formulations supported by fmmax. In particular, the crossed metallic grating discussed in #34 could be a good choice. That grating comes from example 3 of Li's paper.

This task could largely be achieved by copying and existing example (e.g. metal grating) and modifying it so that the correct geometry is modeled. The example should also reproduce the figure from Li's paper.

Better type checking in `eigensolve_isotropic_media`

Currently, if the wavelength parameter within eigensolve_isotropic_media is passed as a float, rather than a jax array, then it throws an obscure error (not directly pointing to the wavelength parameter itself).

It would be good to either

• explicitly allow for float values and do the necessary conversions under the hood.
• assert a type up front and throw an error describing why this needs to be the case.

Automated releases with new versions

Document near2far nuances

Normally, near2far transformations within RCWA/FMM codes are really straightforward. All you have to do is propagate your fields until you hit a uniform, homogenous layer. Once you hit that layer, you can look at the relative power in each Fourier order. Since the mode solutions here correspond to plane waves, each Fourier order will have a direction (far field $\theta$ and $\phi$) one can use to map out the far field space. Of course, only planewaves within the lightcone (which aren't evanescent) make it to the far field. If we want to increase your angular far-field resolution, you need to re-run your simulation centered at different points in the Brillouin zone. You'll also want to compute things like solid angle etc. But the idea is that far field maps are easy to generate using RCWA/FMM.

Unless you incorporate PML. In this case, your last homogenous layer propagating infinitely can no longer be decomposed into simple planewaves. So now, in order to compute the far field, you have to use a Green's function convolution. While the Green's function simplifies quite a bit when concerned with the angular far field (it turns into a Fourier transform), you have to include the propagating energy around your emitter -- meaning you'll have to compute the (vectorial) fields through the z stack (depending on your problem).

We should document these nuances.

Specify a stack with a sidewall angle

Currently, if one wants to simulate a structure that exhibits a sidewall angle, then one must slice the geometry into a "wedding cake" manually. It would be great if we could automate this procedure via an API hook.

Similarly, it might be worth investigating how straightforward it would be to cache several different wedding cake layers (even if they aren't currently used), that way one could optimize the global thickness of the stack (or even the angle itself) without recomputing new eigenmodes each iteration.

Multiple source planes

It'd be great to be able to specify multiple source planes within a simulation.

Similarly, it'd be great to have an API for volumetric sources (e.g. imagine if the source had a finite thickness associated with it).

Broken tests due to new numpy version

Some tests depend upon autograd-backed grcwa, which have started failing due to changes in numpy. It seems autograd will not be updated to address the issue. HIPS/autograd#618

We should perhaps set up CI so that autograd-dependent tests run using a pinned version of numpy.

Hanging (or slow) fmmax on GPU with vector formulation and uniform patterned layers

I am finding that fmmax can hang or be unexpectedly slow in certain scenarios, specifically on GPU platform with uniform patterned layers and using a vector formulation.

I put together a colab which shows non-problematic and problematic cases. Right now it uses the invrs.io gym polarization sorter example, so a simpler fmmax-only repro is likely possible.

https://colab.research.google.com/drive/1BdntytzVa8Li66VeY0QCOX8LMoyRL1dG#scrollTo=RT9goZ8pFAXd

Faster vector field generation

Currently, vector field generation takes a significant amount of time, so that the simulation time for FMM calculations with a few hundred Fourier orders could actually be dominated by the vector field calculation. We should implement a faster scheme for obtaining these vector fields.

Validate power loss within lossy materials

We've already done quite a bit of validation. But it would be good to ensure that the computed emitted power within a lossy region really is what we think it is... there are a lot of nuances associated with this.

Docs and examples

Here a some important areas that could benefit from improved documentation and examples:

• Vector field formulations -- discuss the nuances of each method, how the loss function is computed, etc.
• Convergence -- wrt vector field formulation, number of terms, etc.
• Sharp bits example -- (e.g. what happens when the wavelength is a multiple of the unit cell...) all the gotchas
• Explanation of types, calculation basics
• Scattering matrix fundamentals
• Field processing routines
• Errata

Gaussian beam sources

We'd like a function that creates a Gaussian beam source for use with our new amplitudes_for_fields function.

The signature could look something like this:

def gaussian_beam_fields(
    x: jnp.ndarray,  # coordinates within the unit cell
    y: jnp.ndarray,  # coordinates within the unit cell
    polar_angle: jnp.ndarray,  # Defines the incident beam direction
    azimuthal_angle: jnp.ndarray,  # Defines the incident beam direction
    polarization_angle: jnp.ndarray,  # Defines the polarization of the incident beam
    # other beam parameters go here
) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]:  # ex, ey, hx, hy

cc @smartalecH

Improved s-matrix formulation

Occasionally we observe some violation of energy conservation in results computed by FMMAX. We suspect that this is due to a suboptimal treatment of the scattering matrices. There are a few reports in the literature of methods with some advantages, including methods that avoid "risky" operations such as matrix inversion, or methods that preserve symmetry that one might expect from the scattering matrix (e.g. (S21 == S12). Let's collect relevant references here.

• IMPROVED FORMULATION OF SCATTERING MATRICES FOR SEMI-ANALYTICAL METHODS THAT IS CONSISTENT WITH CONVENTION
• Hybrid-matrix algorithm for rigorous coupled-wave analysis of multilayered diffraction gratings
• Enhanced R -matrix algorithms for multilayered diffraction gratings
• Hybridization of electromagnetic numerical methods through the G-matrix algorithm

far field of the metalens

Hi,

I am a graduated student from Taiwan.

Currently, I am going to develop a design flow for large-area metalens.

I plan to decompose the metalens and then optimize each sub-metalens with aperiodic Fourier modal method.

It is mentioned by this jornal:
https://www.nature.com/articles/s41377-019-0159-5

And, how can I get the far field of each sub-metalens?

Thanks for your help.

Vector fields docs

Would be good to add a simple page in the docs that describe how fmmax computes vector fields. The paper describes this to some detail, but even since then, fmmax has changed it's approach.

Some interesting details include...

• FOMs used
• Comparisons of methods (especially compared to their FOURIER counterparts)
• When to use some vs others

We can start will a "brain dump" (it doesn't have to sound like a paper). But documenting this will help us keep things consistent.

Discretization techniques in z

There's a lot of nuance behind how one decides to discretize in z. For example, in the microlens array example, the structure is sliced such that the upper corner of each slice lies right on the ideal curve. Alternatively, one could have sliced along the lower corner, or better yet, performed a "midpoint rule."

It would be interesting to explore the different schemes here, and how one could borrow tricks from other domains (e.g. finite-differences and quadrature) to minimize the error here.

Specifically, for simple structures (like angled sidewalls consisting of linear slants) there may be some analytic solutions that describe the optimal slicing for a given number of slices.

Document batch conventions

In the README.md, we talk about fmmax's ability to batch over operations. But it would be good to discuss exactly how the user should provide the arrays such that they are properly shaped.

In particular, there are three quantities within eigensolve_isotropic_media that must have consistent batch dimensions:

• angular_frequency (which comes from wavelength) with a non-batch dim of [N]
• in_plane_wavevector with a non-batch dim of [2]
• permittivity with a non-batch dim of [N,M]

If the user wants to batch over multiple wavelengths, then they should specify one additional batch dimension in front of these. If the user wants to batch over different k-points, then they should specify two additional batch dimensions. If the user wants to loop over a cartesian product of wavelength vectors and k-points, then they should specify three additional batch dimensions (which obviously must be consistent).

So if we wanted to batch over 5 wavelength points, and a 4x4 set of k-points, and we have a permittivity image with a size of 100x100 we would expect the following shapes:

• angular_frequency - [5, 1, 1]
• in_plane_wavevector - [1, 4, 4, 2]
• permittivity - [1, 1, 1, 100, 100]

@mfschubert does this align with you? Anything else we should add?

Anisotropic structure in 2D simulation

To the contributors of this code，

Appreciate your valuable code for RCWA calculation, especially the exploration in anisotropic materials！

I noticed that there is an example showing how to calculate 1D anisotropic grating. However, the example only considers the condition that the off-diagonal elements of the dielectric constant tensor appear in xy and yx directions. When I try to calculate a two-dimensional periodic structure, such as periodic hole, the off-diagonal elements may appear in xz and zx directions. So, how can I implement the calculation of 2D periodic structures composed with anisotropic materials, especially when the dielectric constant tensor has xz or yz components? I wonder if this will require many modifications, such as the algorithms used in the Transfer Matrix Method (TMM) and the scattering matrix calculations.

This issue has been bothering me for some time, and I am eagerly looking forward to your assistance.

Thank you so much!

Variations to the Lorentzian Broadening

To condition the manual vjp of the eigenvectors, fmmax uses a lorentzian broadening term (like many other FMM/RCWA implementations):

fmmax/src/fmmax/utils.py

Lines 164 to 168 in a509ffa

	eigenvalues_i = eigenvalues[..., jnp.newaxis, :]
	eigenvalues_j = eigenvalues[..., :, jnp.newaxis]
	f_broadened = (eigenvalues_i - eigenvalues_j) / (
	(eigenvalues_i - eigenvalues_j) ** 2 + eps
	)

That being said, the current formulation looks particularly sensitive to phase differences between eigenvalues. The convolution with the lorentzian test function should be complex-valued. One possible change could be:

$$ f_0 = \frac{\bar{Z}}{\bar{Z}Z + \epsilon} $$

where

$$ Z = \lambda_i-\lambda_j $$

Broadband FOMs

It would be great to be able to specify broadband FOMs (with optional weighting functions).

Since fmmax is a frequency-domain method, one would either have to run each frequency solve sequentially, or distribute each solve in parallel (e.g. using MPI). It would be good to have a standard API for both options. The latter option, which may appear difficult, should actually be straightforward thanks to tools like mpi4jax.

One thing to be aware of: If the specified FOM is scalar-valued (e.g. we take the mean across multiple frequencies) then jax will handle the gradients just fine (even if we use MPI). However, if we specify a vector-valued FOM (e.g. for performing a minimax optimization) then jax has some trouble propagating the gradients... there's some ways around this but it will require some tweaking.

Documentation

Would be good to export the API docstrings into some documentation. The examples could be exported too (even better would be to export notebooks with the accompanying text explaining things).

Visualization routines

It might be nice to add some visualization routines. In theory, this should be really straightforward, especially since everything is defined later by layer. But there are many ways one could envision an API for this (especially since FMMAX is fundamentally so low-level).

So here are a few use cases where visualization is really important. Maybe by outlining these, we can sort out a reasonable API:

• We want to ensure the layer stack is in the right order with the right absolute thicknesses, and that each layer has the right material.
• When defining something with curvature (like a uLens) we want to make sure the curvature is properly discretized in the layered direction
• We want to make sure that free-form structures (e.g. for TO) are properly inserted into the system, so full cross sections of arbitrary slices would be useful.
• We don't want to have to simulate anything to visualize things, so we prefer not to work with LayerSolveResult objects (or scattering matrices etc)

Data import/export functionality

It would be great to have robust and standardized tooling for data import/export.

For example, it would be useful to be able to cache to disk LayerSolveResults scattering matrices, and amplitudes corresponding to a particular basis.

Pickle files aren't very robust (or secure). Perhaps Hdf5?

Incorrect magnetic field in crystal example

https://github.com/facebookresearch/fmmax/blob/main/examples/crystal.py#L368

Should this be ex * jnp.sqrt(permittivity_ambient)?

convergence wrt to plasmonics

Hello, dear developers,
First of all, thanks for the code.

I am wondering if you have tested convergence on structures involved metals for plasmonics. I have tried a few RCWA & FMM codes, some do reasonably well (RETICOLO, for example) and some not so (an example is torcwa. Since you have implemented absorbing boundary layer, I guess in some situations it should work.

For reference, I have a small script that uses torcwa.

cheers

Proper vector fields for odd-shaped unit cells

Currently, the vector field calculations are correct only when the basis vectors are orthogonal, and the resolution of the grid along each of the basis directions is equal.

For example, the following would be problematic, for any choice of resolution:

primitive_lattice_vectors = basis.LatticeVectors(
    u=jnp.asarray([1.0, 0.1]), v=jnp.asarray([0.1, 1.0])
)

Another problematic case is,

primitive_lattice_vectors = basis.LatticeVectors(u=basis.X, v=basis.Y)
permittivity = ... # Permittivity with non-square shape, e.g. (200, 100)

Batching with pmap

Currently, FMMAX performs a lot of broadcasting under the hood, such that it's easy to simulate a device over multiple wavelength points and/or multiple k-points (all of which are completely independent simulations and can be executed in an embarrassingly parallel fashion).

In some cases, the full cartesian product of simulations cannot fit on a single accelerator, and it would be nice to distribute this across multiple accelerators.

Jax has some functionality for this using pmap. It might be nice to set up an example that takes an arbitrary combination of wavelengths and k-points and distributes the computation across all available accelerators. There are a lot of things to consider here, of course. For one, the eigendecomposition actually dispatches back to the host. So if all the devices perform this same dispatch, this could quickly become a bottleneck. Also, there are some limitations with pmap semantics when the number of parallel jobs is not an integer number of local accelerators, or if the accelerators live on different nodes.

Absorbing boundary layer

It would be nice to support x/y absorbing boundary conditions, so that non-periodic structures could be modeled.

Unfortunately, this seems to be a bit involved, since u-PML would require magnetic, anisotropic materials. The implementation could be relatively straightforward for the basic "FFT" formulation of the FMM, but it would be good to make this available also for the vector formulations.

If this were implemented, there would also be an opportunity to simplify the layer eigensolve code, which currently has several different functions to handle e.g. uniform, patterned, isotropic, anisotropic cases.

Ditch docusaurus

We need to ditch docusaurus for a better alternative (one less prone to exploits).

@mfschubert are you using Sphinx for all your repos?

Revisit public/private API

In preparation for documentation of the API, to keep the size of the public API reasonable, we should revisit what functions/modules are public and documented, and which are private.

Add Magnetic current sources

Currently, the amplitudes_for_source() function allows the user to specify an arbitrary electric current-source distribution (e.g. $J_x$, $J_y$, $J_z$). It would be useful to also be able to specify accompanying magnetic current sources (e.g. $M_x$, $M_y$, $M_z$). This would enable injecting unidirectional profiles into a simulation (which requires a complete ansatz).

Twisted photonic bilayer crystal slabs

Nice work on FMMAX!

I'm interested in implementing twisted PhCs in FMMAX, following the paper of B. Lou.

In particular I'm interested in coupling this extension with you BZI sources (Lagrange-Gaussian beams).
Are you open to external contributors? I already have a WIP version in a code of mine.

Cheers

"Long" tests

Now that fmmax has some updated vector field formulations, it would be good to rerun some longer, more exhaustive tests just to ensure things are as we expect. Here are some of the axes we should check:

• All of the different vector field formulations (currently the tests includes mostly POL)
• Complicated geometries with metals and metasurfaces (e.g. modify the uLED to include a random metasurface)
• Convergence checks with increased Fourier terms and BZ points

By specifying a separate set of "long" (or expensive) tests, we don't have to run them with CI. Rather, we can just run them locally when a major change is made. If we see regressions, we can manually git bisect to find the root of the issue.

LatticeVector shape checking

There is currently a check within basis.py to ensure that the LatticeVectors are of the proper shape:

fmmax/fmmax/basis.py

Lines 37 to 43 in 1bc9673

	def __post_init__(self) -> None:
	if isinstance(self.u, jnp.ndarray):
	if self.u.shape != (2,) or self.v.shape != (2,):
	raise ValueError(
	f"`u` and `v` must have length 2, but got shapes "
	f"{self.u.shape} and {self.v.shape}."
	)

This works fine for simple problems. But when trying to backpropagate through more complicated structures, jax will expand the shape of the basis accordingly (and this check fails).

Is there a better way to account for this?

(cc @mfschubert )