Poke: an open-source ray-based physical optics platform

From the Rayfront we have immediate access to data that describes the constituent rays. These attributes are shown in Table 2.

These data can be accessed directly from the Rayfront object. They are numpy arrays whose axes correspond to rayset, surface, and then ray (see the example in Listing 4). With the basic attributes, we can perform simple analyses such as spot diagrams or examining the OPD of the ray bundle at the exit pupil (shown in Figure 2.

We can access the coordinates at any surface in the optical system, which means that the totality of ray-based analyses can be built in Poke. However, Poke’s primary utility is to serve as a platform for ray-based physical optics. In the next section we illustrate two key physics modules built into Poke that can simulate diffraction and polarization from purely ray data.

Traditionally, physical optics propagators employ diffraction integrals derived from the Huygens-Fresnel principle. Plane-to-plane propagation is done with the Fresnel diffraction integral typically by using the angular spectrum method to model the transfer function of free space using a Fast Fourier Transformation. However, using these diffraction integrals imposes the scalar and paraxial assumption on the optical system of interest. This limits the scope of what the diffraction model can accurately determine.

Gaussian Beamlet Decomposition (GBD) is an alternative method of physical optics propagation that has seen substantive implementation in commercial software. The operating principle of GBD is to decompose an arbitrary field into a finite sum of coherent Gaussian beamlets. The analytical propagation laws for these beamlets are known via ray transfer matrices, so they can be propagated along geometric ray paths. GBD doesn’t have an apparent presence in the open source. This is likely because a clear algorithm for its implementation is not present in the literature. We derived a method for its implementation and built it into Poke to study its suitability for high-contrast imaging[6]. The operating principle of GBD is exhaustively covered in the literature [7, 8, 9]. Because we know the analytic propagation laws of Gaussian beams, we can decompose an arbitrary wavefront into a finite sum of Gaussian beams and propagate them anywhere in a 3D non-paraxial optical system. In Figure 3 we show the ability of our poke.gbd module to reconstruct the analytical airy function. GBD requires that each beamlet is propagated to each point that we wish to know the field at. If we desire to propagate a field decomposed into $N_{beamlets}$ to a plane that is $N_{pix}\times N_{pix}$ in size, then the total number of propagations is equal to $N_{beamlets}\times N_{pix}\times N_{pix}$. This is a daunting challenge computationally, so careful consideration was built into the structure of our algorithm. All computation is vectorized, with the option to break it the simulation up into loops over the beamlet index should the simulation require more memory than a given computer has. We have also implemented accelerated computing options for computation on GPUs.

Optimizing the GBD technique is an active area of research for which Poke is a prime platform. It is ray tracer agnostic and can be studied by anyone with access to Python 3.8+. Furthermore, we believe this to be the first open-source implementation of the transfer-matrix method of GBD pioneered by Worku and Gross. Harvey et al [7] report on the complex ray tracing method implemented in FRED, which is the first known iteration of the technique. The transfer-matrix method generalizes the problem, enabling the propagation of any analytical solution to the general Collins’ integral[10, 6]. Consequently, we have the ability to alter the electric field profile of the Gaussian beamlets for higher-fidelity diffraction simulation.

High-contrast imaging instruments in the next decade will be forced to contend with polarization aberrations[11]. The theory of polarization aberrations is described by McGuire and Chipman[12, 13], and Sasian[14]. These effects introduce low-order aberrations in even perfectly aligned optical systems, resulting from the polarized interaction of light with material interfaces. Yun[15, 16] outlined the basics of Polarization Ray Tracing (PRT), an algorithm by which polarization aberrations can be diagnosed. An exhaustive review of the technique is covered in Yun and Chipman, and will not be explained in detail here, but we refer the readers to Polarized Light and Optical Systems[17] for a tutorial. PRT propagates the effects of Fresnel Reflections through a 3D optical system to understand the total vector field behavior. Most notably, PRT gives us access to the Jones pupil. This is a critical data product for conducting diffraction modeling with consideration for polarization aberrations. PRT has been implemented in commercial design codes (CODE V, Zemax), and POLARIS-M⁵⁵5https://airyoptics.com/polaris-m-software/, which operates in Mathematica. However, there isn’t an apparent package with support in the open source community. To address this need, we built PRT into Poke to bring the technique into Python 3.8+. Users can retain their optical system ray traces in their preferred software, but load the PRT data directly into a Python environment. Using Poke, one can generate a Jones pupil using the following code. In this example, we illustrate the polarization aberrations of a Cassegrain telescope with a fold mirror[18].

Poke can decompose the elements of the Jones pupil into Zernike polynomials to re-sample them onto a regularly spaced grid. We can then leverage poke.interfaces to export the data to a physical optics propagator to perform diffraction analysis.

Poke provides us with the platform we need to democratize PRT, while still being useful for investigators with access to commercial ray tracing licenses.

Corollary to PRT is the need to simulate thin film multilayer stacks. Polarization in optical systems is frequently caused by the thin films deposited on the optics. To accurately model these effects in PRT, we built poke.thinfilms to be able to encompass the influence of arbitrary stacks of homogeneous, linear, and isotropic thin films. We do so via the characteristic matrix[19], but shown below in Equations 1 and 2

Equations 1 and 2 show the characteristic matrix $A$ from which the complex amplitude coefficients are readily derived. While this utility is useful in PRT simulations, it can be called independently to design thin film multilayer stacks. Below we show the design of a long pass filter using poke.thinfilms and scipy.minimize. This is an excellent example of the utility available to us by bringing these physical optics techniques into Python. In Figure 6 we illustrate a short-pass filter that was designed using the thin films module, and the runtime to compute a Jones pupil from this film data.

The thin film interface with optical surfaces is also compatible with spatially-varying thin films. This utility is critical in understanding how changes in film thickness can influence the Jones pupil of an optical system. In Figure 7 we apply an $45^{o}$ astigmatic coating to the primary mirror of the Extremely Large Telescope (ELT) and evaluate the Jones pupil. From these data, we are able to understand such a film’s influence on the amplitude and phase of our polarization state.

Poke has already been used to model the nominal polarization aberrations present in the GSMT’s. In Anche et al [5] we discovered that the polarization aberrations due to the extremely large angles of incidence on the observatory mirrors lead to polarization aberration residuals that were greater than the desired contrast at the focal plane of ideal coronagraphic instrumentation. To achieve the baseline contrast proposed for these instruments, we must develop a method to mitigate the effect of polarization aberrations. In Ashcraft et al. in prep, we are developing an optimization-based method to design thin-film stacks to mitigate polarization aberrations.

In Figures 8-9 we illustrate the results of using PRT with optimization routines in Python to design thin film stacks that can reduce the presence of polarization aberrations. This technique has potential implications for maximizing the achievable contrast near the inner working angle of coronagraphs for next-generation astronomical observatories.

Poke defaults to numpy for the vast majority of its numerical calculations. However, for high-fidelity simulations it is advantageous to employ Python packages with numpy-like API’s that grant the user more utility. Poke has adopted the interchangeable backend system of prysm (see section 4A of Dube et al[20]). Poke’s backend can be swapped from numpy to cupy for accelerated GPU-based computation, and jax to enable automatic differentiation.

cupy is an incredible Python package that allows users to scale their computations to GPUs without making significant changes to their code. It’s numpy-like API makes it trivial to implement in any codebase. GPUs are especially suited to large matrix multiplications, which are the fundamental operation of the physical optics modules discussed in Section 3.

jax is a Python package developed by Google that supports automatic differentiation (or autodiff). Autodiff enables the computation of analytic gradients without requiring the user to derive the analytic expression for gradient back-propagation. This is particularly useful for complex expressions, where the mathematics would grow unwieldy. jax makes Poke fully differentiable, enabling accurate computation of gradients for accelerated optimization.