Add support for polarized wavefronts

[kvangorkom]

This PR extends the Wavefront classes to support Fraunhofer and Fresnel propagation of fully- and partially-polarized fields.

Wavefronts and propagation

Fully polarized fields are treated as 2xYxX vector fields as, e.g,. wf = PolarizedWavefront(input_polarization=(1,1j))

Following the hcipy treatment, partial polarization is handled by propagating 2x2xYxX tensor fields through a wavefront initialized as, e.g., wf = PolarizedWavefront(input_stokes_vector=(1,0,0,0)) and evaluating the stokes parameters wf.stokes_parameters.

(and the same for Fresnel diffraction with the PolarizedFresnelWavefront class)

FFTs and MFTs are generalized to handle the extra dimensions by performing the transforms and fftshifts over the final two dimensions.

Polarization optics

This also adds a handful of polarization optics, including linear and circular polarizers, quarter- and half-wave plates, and general Jones matrix elements.

I've updated the BaseWavefront.__imul__ method to special case the interaction with polarization optics such that the multiplication happens over the correct axes, and check that polarization optics can't be multiplied with scalar Wavefront classes.

Notebooks and tests

I've added a notebook with a couple very simple polarization demos here

A bunch of tests of the new Wavefront classes and interactions with polarization optics are added in test_polarization.py. Some initial multi-dim propagation tests have been added in test_matrixDFT.py

Open items

• Still needs a lot more tests (more testing of Stokes calculations, interactions with polarization optics in image/pupil planes, propagation, etc.)
• Current display code is pretty broken and needs to be extended to handle vector/tensor/stokes visualization. OpticalSystems can display polarization optics, but it ignores the Jones matrices. Displaying wavefronts is just broken.
• Add more polarization optics (probably extend a couple of the existing coronagraph masks to their polarization equivalents)
• Lots of general code clean up
• Timing tests
• Check/adapt cupy-specific code paths

---

There's still a fair amount of work to do before this is ready to merge, but I thought I'd get this out there before it has a chance to stagnate on my fork!

[mperrin]

Let me just say -- while awaiting further work, I see you've marked this as WIP -- but already, bravo and wow! This is not a small amount of work. It's going to one a much appreciated significant enhancement to have this; thank you!

[kvangorkom]

This PR is at a point now where I'm generally happy with it. I think the test suite has pretty good coverage of the changes/additions I've made, and all tests are currently passing on both CPU and GPU.

There are definitely a few cases where there's room for healthy debate on the best way to implement (e.g., BaseWavefront.shape only includes spatial dimensions, which hides a bit of what's happening under the hood with PolarizedWavefronts, and the display/plotting code is a bit rough), and I'm guessing there are still some edge cases lurking where PolarizedWavefronts either break or do something unexpected.

In any case, I think it's probably ready for some more eyes!

[marcio-melendez]

@kvangorkom this is an impressive amount of work and well documented with the notebook and the tests. I ran the notebook without any issues, apart from some overlapping labels in some of the figures, but that's ok.
I did not try the GPU part, as I don't have an easy access to test it.
Since there has been two months since this PR was ready for a review I wanted to ask if there is something else that you may have worked/improved during this time and want to include before I accept.
Thanks!

[kvangorkom]

I haven't touched this since the commits above, but I am about to start using this code to do some analysis. I fully expect to find some things that need tweaking/expanding in the next few weeks, so I'm okay waiting a bit longer before merging.

That said, it would be super useful to have a couple people take look through this to make sure I haven't done anything egregious! Maybe @ramya-anche and/or @Jashcraf could take a glance? (😇)