Addition of CuPy as an Accelerated Computing Option

[kian1377]

This is a copy of PR spacetelescope#499 to the spacetelescope/poppy repository.

This is a bit more of an extensive PR as it seeks to add CuPy as an accelerated math option for many computations including FFTs, MFTs, and exponentials for propagation with the OpticalSystem or FresnelOpticalSystem classes. To start with, CuPy is a package for GPU accelerated computing. While POPPY has GPU computing options available with PyOpenCL and with pyculib (which has been deprecated by Numba in favor of CuPy), this implementation seeks to perform all calculations on the GPU until the end of propagation. This reduces time for calculations as arrays no longer need to be transferred between GPU memory and standard memory when performing different calculations.

CuPy has been designed to be very similar to Numpy and has many of the same features requiring the same syntax to use. An example is numpy.fft.fft2, which with CuPy is cupy.fft.fft2. The catch is that CuPy functions can not be used on Numpy arrays since Numpy arrays are stored on standard memory. As such, the implementation strategy was to use the statement “import cupy as np” when POPPY’s Config has been set to use CuPy. This makes performing all calculations on GPU more seamless as wavefront arrays for Wavefront or FresnelWavefront objects along with optical element phasors are automatically calculated as CuPy arrays.

In addition to using CuPy for basic computations of FFTs and exponentials, CuPy also enables the use of many SciPy functions through its cupyx functions. The cupyx functions can be imported much like many scipy functions with a statement such as “import cupyx.scipy.ndimage as ndimage”. So many functions for computing array rotations or interpolations are also imported through cupyx or scipy based on POPPY’s Config setup. The same method is also used to compute Bessel functions.

Currently, many optics have been tested for use with CuPy. A table with these optics is provided below with some comments regarding the functionality of some optics.

Optical Element Class	Compatible with CuPy	Comments
CircularAperture	Yes
MultiCircularAperture	Yes
SquareAperture	Yes
RectangularAperture	Yes
HexagonAperture	Yes
MultiHexagonAperture	Yes
NgonAperture	Yes
SecondaryObscuration	Yes
AsymmetricSecondaryObscuration	Yes
CompoundAnalyticOptic	Yes
ThinLens	Yes
GaussianAperture	Yes
KnifeEdge	Yes
SquareFieldStop	Yes	These image plane elements have only been tested in Fraunhofer systems since it is difficult to use these elements in Fresnel systems anyway. This is because of the issue in converting to angular units in a Fresnel system.
RectagularFieldStop	Yes
AnnularFieldStop	Yes
HexagonFieldStop	Yes
CircularOcculter	Yes
BarOcculter	Yes
BandLimitedCoronagraph	Yes
IdealFQPM	Yes
ScalarTransmission	Yes
InverseTransmission	Yes
ScalarOpticalPathDifference	Yes
ZernikeWFE	Yes
KolmogorovWFE	Yes	Does not pass test when using Tatarski power spectrum.
SineWaveWFE	Yes
StatisticalPSDWFE	Yes	Had to implement hack found in Issue spacetelescope#452, pre-existing bug not related to CuPyNote: slight differences in numpy vs cupy random generators so does not produce the exact same result with the same seed value.
PowerSpectrumWFE	Yes	Had to assume opd is calculated in nanometers and then rescale it to meters, there should be a better solution to this. Note: slight differences in numpy vs cupy random generators so does not produce exact same result
ContinuousDeformableMirror	Yes	Influence function must be provided. When using .set_surface(), you should still provide a numpy.ndarray instead of a cupy.ndarray.
HexSegmentedDeformableMirror	Yes	These segmented optics are functional but when converted to an ArrayOpticalElement with fixed_sampling_optic(), the OPD is only 0. Still not 100% sure why but it does not seem to be an issue with CuPy because get_opd() functions as expected.
CircularSegmentedDeformableMirror	Yes
ArrayOpticalElement	Yes
FITSOpticalElement	Yes
FixedSamplingImagePlaneElement	Yes

All tests in the test suite for POPPY were also run. Currently, there is only an issue with the KolmogorovWFE test not passing due to a units issue in an exponential. This only comes up when using a Tatarski power spectrum. The tests have only been run with standard CPU computations to make sure POPPY is not running into critical errors because of the addition of CuPy even if a user isn’t using the CuPy feature.

Computation comparisons have been performed to illustrate the benefit of this accelerated computing feature. Below are comparisons of the times required for a PSF to be calculated for varying array sizes using the MKL FFT option versus the CuPy calculations. The optical systems tested had 5 different surfaces/optics. The system used for these comparisons was the University of Arizona’s HPC Puma nodes. The node utilized 32 AMD EPYC 7642 CPUs and the NVIDIA Tesla V100S GPU.

Propagation Type	Array Size	MKL Method Times [s]	CuPy Method Times [s]	Speed Up Factor
Fraunhofer	1024	0.218	0.0261	8.35
Fraunhofer	2048	0.755	0.0294	25.7
Fraunhofer	4096	3.36	0.0423	79.4
Fresnel	1024	0.714	0.0438	16.3
Fresnel	2048	4.16	0.0845	49.2
Fresnel	4096	17.5	0.225	77.8

One catch with this is none of POPPY’s display functionality is compatible with CuPy. This is because matplotlib cannot plot CuPy arrays since they are only on GPU memory. In order to obtain a Numpy array from a CuPy array, the “cupy.ndarray.get()” method can be used. So users can obtain intensity and phase arrays by adding .get().

It should be noted that the following POPPY features have not been tested for functionality with CuPy:

• PhysicalFresnelWavefront
• Active optics such as the TipTiltStage (found in active_optics.py)
• The Instrument class
• Special propagation features such as SemiAnalyticCoronagraph (found in special_prop.py)
• Sub-sampled optics such as ShackHartmannWavefrontSensor
• Floating-window centroid calculations