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]

These need some way to set:

• Wavefront.ranges
• Wavefront._pad.pad

These setters are needed for smart propagation (when you resize the grid as the beam grows)

[ChristopherMayes]

We also need a method to write the Genesis4 field data from Wavefront (and corresponding test)

[ken-lauer]

@ChristopherMayes noted below

These need some way to set: Wavefront.ranges Wavefront._pad.pad
These setters are needed for smart propagation (when you resize the grid as the beam grows)

ranges is not settable as it is calculated based on the grid spacing and the shape of rmesh.

With the latest code I pushed, you can provide a new rmesh by way of:

new_wf = wf.with_rmesh(rmesh=expanded_version)

and the previous metadata will be retained. Padding can optionally be manually specified (or automatically adjusted based on scipy's suggested size)

[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?