search

facebookresearch / fmmax

Fourier modal method with Jax
MIT License
97 stars 10 forks source link

Support s-matrix generation via scan #103

Closed mfschubert closed 7 months ago

mfschubert commented 7 months ago

Adds a new stack_s_matrix_scan function to scattering module which computes scattering matrix by scan, rather than with python for loop. This could reduce compile time and program size in cases where the number of layers is large.

The "best" way to make use of this functionality is to also compute the layer eigensolves on a batch of permittivities (with leading match dimension corresponding to the layer).

To support this change, we also enabled batch dimensions for LatticeVectors objects throughout. This is needed in this context, because it puts the primitive_lattice_vectors attribute of the LayerSolveResult on the same footing as all other attributes (i.e. it is no longer a special attribute without batch dimensions).

Note that the spatial coordinates now also can have batch dimensions, which are singleton and squeezed out in several tests. This may be relevant when computing fields or sources.

smartalecH commented 7 months ago

(cc @marcoamonteiro )

mfschubert commented 7 months ago

I'm also throwing in an improvement to the detection of permittivities that vary in one direction only (which should lead to uniform vector fields).