always fftshift+iffshift for Fraunhofer propagation

[kvangorkom]

This PR addresses what I think may be a bug in the accel_math.fft_2d implementation, where fftshifts weren't being performed before and after every fft2, leading to checkboard patterns in the complex fields.

You can reproduce this with a simple code block:

npix = 128
D = 10 * u.mm
wavelen = 1e-6 * u.m
oversample = 2

circ = poppy.CircularAperture(radius=D/2)
img = poppy.ScalarTransmission()
osys = poppy.OpticalSystem(npix=npix, oversample=oversample)
osys.add_pupil(circ)
osys.add_image(img)

wf = poppy.Wavefront(diam=2*D, wavelength=wavelen, npix=npix, oversample=oversample)
psf, wfs = osys.calc_psf(inwave=wf, return_intermediates=True)
plt.imshow(wfs[-1].wavefront.real)

the output of which is

vs the expected output of

My proposed fix is to always have accel_math.fft_2d perform an ifftshift and fftshift when fftshift=True, regardless of whether forward or inverse FFTs are being performed.

I added test coverage by modifying test_matrixDFT.test_MFT_FFT_equivalence_in_OpticalSystem to explicitly check agreement between the complex outputs of the MFT and FFT (instead of just intensities).

All tests appear to be passing locally, but it's not 100% obvious to me that this couldn't have an unintended consequence that's not getting test coverage!

[mperrin]

Thanks much @kvangorkom! Good catch. I agree this is pretty clearly a bug.

This fix definitely improves it, but I'm wondering if it's 100%, or if there are other edge cases to think about. In particular, I'm this case you've made it do ifftshift then fftshift regardless of the overall transform direction. The other potential fix could instead be directionally-dependent, like:

if fftshift: 
    if forward: 
        wavefront = _ifftshift(wavefront)
    else: 
        wavefront = _fftshift(wavefront

And then the reverse order later on.

I'm not asserting the above is correct or better. Rather I'm realizing it would be good to check this with a sufficient suite of test cases including odd- and even-sized wavefronts, in both forward and backward direction transforms.

I'd like to say I would work on that at some point, but alas I've been out all month dealing with family medical issues, so am not sure on what timescale I could get to checking on this in any more detail. You might be able to double check some of this sooner than I could, perhaps. In any case thanks very much for this!

[mperrin]

Very very minor: It would be nice to also update the comment text here to "shift before all propagations" instead of "shift before backwards propagations"

[kvangorkom]

Hey, @mperrin! I hope you're doing okay—feel free to ignore this until things have settled back down for you.

I've done some more digging, and I agree that the full solution is a bit more involved than the code changes I've made above. I have a suspicion that there are some other subtle issues with the handling of fftshift/ifftshift elsewhere that were hiding behind this bug, and untangling that is going to take a bit of work.

Stay tuned!