facebookresearch / fmmax

Fourier modal method with Jax
MIT License
95 stars 9 forks source link

Alternate BZ grid generation for better resource utilization #81

Open smartalecH opened 7 months ago

smartalecH commented 7 months ago

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).

mfschubert commented 7 months ago

One thing to have in mind is that some "postprocessing" functions (e.g. for computing farfield profiles) rely on sampling the BZ with a regular grid. I suspect that generalizing these would be somewhat involved.

Also, multiple wavelength and incident angles could increase the batch size and would likewise benefit from optimized partitioning. In general, we may want to allow (and provide functionality for) arbitrary shuffling of elements within the batch to best match the hardware, and potentially allow for "dummy" elements to be inserted in cases where no perfect fit is possible.