WIP/ENH: wavefront class

[ken-lauer]

Add a Wavefront class

Entirely based on code from @balticfish and many discussions with @ChristopherMayes

• Automatically handles real space domain padding
• Auto-determines padding values which result in padded array dimensions which:
  1. Are odd integers (per @balticfish, this is required to get balanced/accurate output)
  2. Are efficient for the scipy FFT (https://docs.scipy.org/doc/scipy/reference/generated/scipy.fft.next_fast_len.html#scipy.fft.next_fast_len; this offers significant speed-ups)
• Supports
  • Propagation in the Z direction
  • Focusing element in real space
  • Automatic recalculation of real/k-space data (by fft/ifft) on an as-needed basis
• Simple example notebook, modified from @balticfish as well
• Genesis4 notes
  • Propagation is exactly z axis in genesis, propagation can work only with small angle (paraxial approximation for drift propagation)

Remaining work

• Refine the API with feedback from @balticfish and @ChristopherMayes
• Consider avoiding property for .rspace and .kspace as their calculation can take a long time, and IPython auto-completion can automatically evaluate these when you don't want it to (double-check where this was coming from in Jedi; is it still an issue in 2024?)
• Add docstrings to everything
• Add units to properties that are missing them
• Remove / rename additional properties, depending on utility
• Remove non-MKS units
  • fs -> s
• Standard axis layout of xyz is desirable. Is there any benefit to z,x,y? Benchmark some operations to verify?
• Support only meters in the Wavefront class
  • Change initializer to reflect this
  • @balticfish will modify functions with t->m (z=c*t)
• Define import/export functionality
  • Can we validate with https://github.com/openPMD/openPMD-validator/tree/2.0.X ?
  • Can we use https://github.com/openPMD/openPMD-api to read/write the data? Possibly an optional dependency?
  • Support reading/writing openpmd-wavefront data
  • See how this integrates with lume-genesis
• Add some tests
• More after discussions
• Rename propagate to drift
• Plotting:
  • XY/ZX/ZY projection: plot("x", "z")
  • Plot Re, Im, Abs and Phase
  • Spectrum plot_spectrum() (this is from G.S.; Chris will supply it)
  • Fancy "Wigner" plot (spectrogram-like) (@balticfish has custom code for this)
• Consider moving some methods out to independent functions
  • @ChristopherMayes prefers to have utility functions that aren't baked into the parameter/wavefront classes in this PR
• Add minimal support for polarization - a direction (X stored as Ex with openpmd, for example)
  • Make an issue for future support of generalized polarization, noting that we do not currently have a use for more complicated polarization schemes
• Ask for feedback from G.S.
• Rework for clarity - refactored code is now unusable to @balticfish

[ChristopherMayes]

sum_axis implies a sum, but here np.mean is called.

For power, it should probably be a sum. Others are questionable.

[ken-lauer]

@balticfish?