Document batch conventions

[smartalecH]

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?

[mfschubert]

Hi @smartalecH, we have some logic which automatically does some broadcasting, here: https://github.com/facebookresearch/fmmax/blob/main/fmmax/fmm.py#L942

I'm not sure your summary above is exactly right, so let me just repeat some of it.

• angular frequency has no non-batch dimensions
• in_plane_wavevector has a trailing non-batch dimension (it has length 2 and contains kx and ky)
• permittivity has two trailing non-batch dimensions, corresponding to the u and v directions (i.e. the basis lattice vectors)

To carry out the calculation you describe, these shapes would indeed work. But, the automatic broadcasting would also allow e.g. the following:

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

These will then automatically be broadcast up to the shapes you've specified.

[smartalecH]

Got it, thanks! Is there a reason the angular_frequency needs to account for additional k-points, but in_plane_wavevector is agnostic to the number of wavelengths, and permittivity is agnostic to both?

Would it be worth it to augment some of the automatic broadcasting logic to allow the user to specify something like:

• angular_frequency with shape [5,]
• in_plane_wavevector with shape [4, 4, 2]
• permittivity with shape [100, 100]

that way each quantity only "cares" about shapes related to that quantity?

[mfschubert]

Ah, I could have been more clear. There is actually nothing special about angular_frequency here -- we equivalently could provide:

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

All that matters is that the desired combinations exist when everything is broadcast together.

Although simpler logic might be possible, I think it would not be as general. Imagine a case where a user wants to simulate certain wavelength/angle combinations.

We could conceivably create some helper functions to reduce the boilerplate, but I don't think we'd want a less general solution at the lowest level.

[smartalecH]

Although simpler logic might be possible, I think it would not be as general. Imagine a case where a user wants to simulate certain wavelength/angle combinations.

Ah yep this put things into context for me, thanks!

Adding a blurb about this in the (future) documentation would be useful!