Reflections on prysm, part 1: pupil planes

Brandon Dube
Last updated on Jan 16, 2023

The subject of this blog post is prysm, a python library I’ve created which is used for numerous tasks in optics, which are enumerated in its documentation. Users reading this should not fear that prysm is going away or will be abandoned. However, after years of development there is not much more to add to prysm without expanding into new domains or applications. Those are not really things I have much expertise in, and would require community development. Prysm also recently turned five, and I want to take the time to reflect on its design for the domains I know how to solve with it.

The library has evolved to have a “fluent” and expressive API, where very few keystrokes accomplish most tasks. Curious readers are encouraged to look at the examples, which include, e.g., parametric phase retrieval in less than 30 lines of user code. The library is very ‘magic’ and some elements of the design are internally very convoluted.

I would like to actively solicit whether the few hundred users would like the API to remain stable as it has been the last few releases and declare 1.0, or make one final sweeping set of design changes before releasing 1.0. This series of posts is not a sign prysm will be abandoned.

This post concerns itself with the propagation engine or, more aptly, with the creation of numerical models of pupils for propagation.

Propagation is the action which moves you from one plane to another in an optical system via diffraction. These are two types of propagation supported by prysm:

  • from pupil to PSF (or vice versa)
  • from plane to plane

Propagation to a region that is neither near a focus nor near a pupil is not a feature, since neither myself nor the community of users have needed it.

A pupil plane is typically thought of as a complex field, whose amplitude abs(f) defines the transmission and whose phase arg(f) describes the optical path error of the field from a perfect plane or spherical wave.

The role of the classes prysm provides for setting up pupils is to automate with few keystrokes the creation of 2D phase and amplitude profiles and the transformation to a complex number. There are many relevant parameters, which has led to growth in the complexity of those code over time, although the outward-facing interface has remained stable and is, in my opinion, wonderfully terse and expressive.

As an example of the growth of inner complexity, we can investigate the method resolution order of one of the most common classes:

Help on class FringeZernike in module prysm.zernike:

class FringeZernike(BaseZernike)
 |  FringeZernike(*args, **kwargs)
 |
 |  Fringe Zernike description of an optical pupil.
 |
 |  Method resolution order:
 |      FringeZernike
 |      BaseZernike
 |      prysm.pupil.Pupil
 |      prysm._phase.OpticalPhase
 |      prysm._richdata.RichData
 |      builtins.object
 |

There are five inheritances which give this class a total of 39 fields, attributes, or methods. The constructor takes **kwargs, which you only find documented for the other classes it shares an inheritance tree with. Different portions of the kwargs propagate up different extents of the __init__ chain, which has lead to subtle bugs over time. FringeZernike, NollZernike and the counterparts are basically front-ends to BaseZernike, whose constructor looks like:

def __init__(self, *args, **kwargs):
    """Initialize a new Zernike instance."""
    self.coefs = {}

    self.normalize = False
    pass_args = {}

    if args is not None:
        if len(args) == 1:
            enumerator = args[0]
        else:
            enumerator = args

        for idx, coef in enumerate(enumerator):
            self.coefs[idx+1] = coef

    if kwargs is not None:
        for key, value in kwargs.items():
            if key[0].lower() == 'z' and key[1].isnumeric():
                idx = int(key[1:])  # strip 'Z' from index
                self.coefs[idx] = value
            elif key.lower() == 'norm':
                self.normalize = value
            else:
                pass_args[key] = value

    super().__init__(**pass_args)

While this provides beautiful polymorphism that “just works the way you expect”, what it is doing may be less than obvious. It enables each of these calling cases:

coefs = np.random.rand(36)
FringeZernike(coefs) # terms 1..36
FringeZernike(Z123=123.456) # sparse terms
FringeZernike(coefs, dia=123.456) # terms w/ kwargs

The code would be clearer and users would more quickly discover these ways of computing sums of polynomials if the code were not polymorphic. Three functions or methods instead of polymorphism would be clearer.

The code itself is not maximally clear; it passes the kwargs through a sieve to filter which ones it propagates up the chain with multiple “stops” along the way (some are for FringeZernike, some are for Pupil, and so forth).

The Pupil constructor is next in the chain and is also complicated, with two if blocks that set a conditional that controls a third, and a nested ‘if’ inside in that third that further switches conditional behavior. Bugs resulting from this have been fixed three times, and each time I was sure would be the last.

We can break down what is done here into a few functions:

phase = phase_fn(*args)
amplitude = amplitude_fn(*args)
meta = {...} # parameters
cmplx = cmplx_transformation(amplitude, phase, *some_of_meta)

meta would include things like $\lambda$ or $\Delta\xi$, the sample spacing of the pupil plane. The task of replacing four function calls with one produced a few dozen lines of meta machinery, excluding the implementations of those phase and amplitude functions. That meta machinery is complicated, too, as we’ve seen. Further, there is some “ball of mud” effect going on, where we basically have all the parameters of the system available all the time. This is not necessary.

If we switch to Go syntax for a moment (fear not, prysm is not and will not be re-written in Go), we can really describe a pupil plane as a relatively small interface. That interface is the same for pupil or PSF planes, so we’ll use Plane as a label:

type Plane interface {
  Field() ComplexArray2D
  Dx() float64 // == Delta xi
  Lambda() float64 // wavelength, not lambda expr
}

There are some minor limitations to this due to assumptions it makes that are invalid for highly oblique propagations, but they are out of scope. The class fields with these purposes are called fcn (wavefunction), sample_spacing, and wavelength in prysm. The changed names reflect gained experience with questions on what the previous names meant, and Go’s lack of lambda as a predeclared identifier.

It’s strongly desireable that the propagation engine only ever know about complex numbers, which is why the Field method would return a complex matrix. The extent in x and y are already available through Field, so they are not repeated in this interface.

Because this is an interface, any type can satisfy it. This means that we are free to define any struct/class we like for the various phase and amplitude combinations we’d like to synthesize.

Those structs do not require any inheritance, though they may all end up wrapping a simple function to convert phase and amplitude into a complex field.

The described interface allows any type to provide 2D fields which are fully described. It would be much clearer than what exists today. A little composition goes a long way in allowing amplitude and phase to be combined:

type AmpAndPhase struct {
    Amp mat.Matrix
    Phase mat.Matrix
    lambda float64
    dx float64
}

func (a AmpAndPhase) k() float64 {
    return 2 * math.Pi / a.Lambda
}
func (a AmpAndPhase) Field() mat.CMatrix {
    // this is pseudocode, or valid python.  Take your pick.
    return a.Amp + math.Exp(-i*a.k() * a.Phase)
}
// a.Dx and a.Lambda can be imagined by the reader

There remains a design issue over what to do with lambda and dx, since they are private and users could not initialize such a struct without a constructor. They overlap with the capitals for the methods. This is a problem with Go; you can only express interfaces (~= explicit dependencies) on functions, not data.

For computing Zernike polynomials, or other phases, the call would look like this:

pln := AmpAndPhase{
    Phase: NollSequence(128, []float64{1,2,3}),
    Amp: Circle(128),
    lambda: .5,
    dx: DxDiam(100, 128)
}

ps := prop.Focus(pln, params...)

The 128 appearing in three places would be the number of samples in the square arrays. DxDiam is a poor name from a to-be-made utility for computing the sample spacing from the diameter and sample count. The name is not very good and needs work.

The comparable python today, with the same explicitness, would be:

pln = NollZernike([1,2,3], samples=128, transmission='circle', wavelength=.5)

ps = PSF.from_pupil(pln, params...)

The verbosity of the former is a bit higher, but you probably understand everything that’s happening, modulo DxDiam without any complicated machinery. This would basically remove the need for the Pupil class in prysm, which is 216 lines (granted, much fewer SLOC). I likely would not provide anything to recreate the default args that let that be written without the three kwargs given. I get too many questions about how to change the parameters the defaults “hide.” The OpticalPhase class would go, too, and all the garbage in the various Zernike class constructors used to deal with the inheritance.

That’s probably 500 SLOC that are not needed. Removing them is not just less code, but a lot less complexity. Most of the needless complexity comes from the use of inheritance. I don’t think I will use much of it in future projects; after the first level, it is not worth it. And the evolution of prysm towards all classes having a data attribute gives, even if implicitly, the behavior inheritance was used to give in the first place. In other words,

class Zernike(Pupil):
    def __init__(self, ...):
        phase = do(...)
        super().init(phase=phase)


class Pupil:
    def __init__(self, phase):
        self.phase = phase

says absolutely nothing more than:

class Zernike:
    def __init__(self, ...):
        self.phase = do(...)

It’s true that users of these types depend on the existence of the phase field. I do not know how to express that in dynamic languages, but I know inheritance isn’t it.