Repository: spacetelescope/poppy

Issue:
add more efficient rotation for 90 degree cases
pointed out by @ehpor on https://github.com/spacetelescope/hicat-package/pull/556#issuecomment-802155424

> the rotations inside poppy are done with scipy.ndimage. That feels overkill for a 90degree rotation. Is there a better way in poppy to do a 90degree rotation (np.rot90())?

It would be easy to add a check for 90 degree rotation and use a more performant function. This should be added in.



Discussion:

Patch:
diff --git a/poppy/poppy_core.py b/poppy/poppy_core.py
index 99bc19c..b74d66c 100644
--- a/poppy/poppy_core.py
+++ b/poppy/poppy_core.py
@@ -880,8 +880,16 @@ class BaseWavefront(ABC):
         # Huh, the ndimage rotate function does not work for complex numbers. That's weird.
         # so let's treat the real and imaginary parts individually
         # FIXME TODO or would it be better to do this on the amplitude and phase?
-        rot_real = scipy.ndimage.interpolation.rotate(self.wavefront.real, angle, reshape=False)
-        rot_imag = scipy.ndimage.interpolation.rotate(self.wavefront.imag, angle, reshape=False)
+
+        k, remainder = np.divmod(angle, 90)
+        if remainder == 0:
+            # rotation is a multiple of 90
+            rot_real = np.rot90(self.wavefront.real, k=-k)  # negative = CCW
+            rot_imag = np.rot90(self.wavefront.imag, k=-k)
+        else:
+            # arbitrary free rotation with interpolation
+            rot_real = scipy.ndimage.interpolation.rotate(self.wavefront.real, -angle, reshape=False)  # negative = CCW
+            rot_imag = scipy.ndimage.interpolation.rotate(self.wavefront.imag, -angle, reshape=False)
         self.wavefront = rot_real + 1.j * rot_imag
 
         self.history.append('Rotated by {:.2f} degrees, CCW'.format(angle))
@@ -2991,13 +2999,21 @@ class FITSOpticalElement(OpticalElement):
             # ---- transformation: rotation ----
             # If a rotation is specified and we're NOT a null (scalar) optic, then do the rotation:
             if rotation is not None and len(self.amplitude.shape) == 2:
-                # do rotation with interpolation, but try to clean up some of the artifacts afterwards.
-                # this is imperfect at best, of course...
-                self.amplitude = scipy.ndimage.interpolation.rotate(self.amplitude, -rotation,  # negative = CCW
-                                                                    reshape=False).clip(min=0, max=1.0)
-                wnoise = np.where((self.amplitude < 1e-3) & (self.amplitude > 0))
-                self.amplitude[wnoise] = 0
-                self.opd = scipy.ndimage.interpolation.rotate(self.opd, -rotation, reshape=False)  # negative = CCW
+
+                k,remainder = np.divmod(rotation, 90)
+                if remainder==0:
+                    # rotation is a multiple of 90
+                    self.amplitude = np.rot90(self.amplitude, k=-k)   # negative = CCW
+                    self.opd = np.rot90(self.opd, k=-k)
+                else:
+                    # arbitrary free rotation with interpolation
+                    # do rotation with interpolation, but try to clean up some of the artifacts afterwards.
+                    # this is imperfect at best, of course...
+                    self.amplitude = scipy.ndimage.interpolation.rotate(self.amplitude, -rotation,  # negative = CCW
+                                                                        reshape=False).clip(min=0, max=1.0)
+                    wnoise = (self.amplitude < 1e-3) & (self.amplitude > 0)
+                    self.amplitude[wnoise] = 0
+                    self.opd = scipy.ndimage.interpolation.rotate(self.opd, -rotation, reshape=False)  # negative = CCW
                 _log.info("  Rotated optic by %f degrees counter clockwise." % rotation)
                 self._rotation = rotation
 


Test patch:
diff --git a/poppy/tests/test_core.py b/poppy/tests/test_core.py
index 34aa215..8cfb7ae 100644
--- a/poppy/tests/test_core.py
+++ b/poppy/tests/test_core.py
@@ -5,6 +5,8 @@ import numpy as np
 from astropy.io import fits
 import astropy.units as u
 import pytest
+import matplotlib.pyplot as plt
+
 try:
     import scipy
 except ImportError:
@@ -365,7 +367,6 @@ def test_displays():
     # As a result this just tests that the code runs to completion, without any assessment
     # of the correctness of the output displays.
     import poppy
-    import matplotlib.pyplot as plt
 
     osys = poppy.OpticalSystem()
     osys.add_pupil(poppy.CircularAperture())
@@ -403,8 +404,8 @@ def test_displays():
 def test_rotation_in_OpticalSystem(display=False, npix=1024):
     """ Test that rotation planes within an OpticalSystem work as
     expected to rotate the wavefront. We can get equivalent results
-    by rotating an Optic a given amount, or rotating the wavefront
-    in the opposite direction.
+    by rotating an Optic a given amount counterclockwise, or rotating the wavefront
+    in the same direction.
     """
 
     angles_and_tolerances = ((90, 1e-8), (45, 3e-7))
@@ -413,23 +414,24 @@ def test_rotation_in_OpticalSystem(display=False, npix=1024):
         osys = poppy.OpticalSystem(npix=npix)
         osys.add_pupil(poppy.optics.ParityTestAperture(rotation=angle))
         osys.add_detector(fov_pixels=128, pixelscale=0.01)
+        psf1 = osys.calc_psf()
 
-        if display: plt.figure()
-        psf1 = osys.calc_psf(display=display)
-        if display: plt.title("Optic rotated {} deg".format(angle))
-
-
-        osys = poppy.OpticalSystem(npix=npix)
-        osys.add_pupil(poppy.optics.ParityTestAperture())
-        osys.add_rotation(angle=-angle)  # note, opposite sign here.
-        osys.add_detector(fov_pixels=128, pixelscale=0.01)
-        if display: plt.figure()
-        psf2 = osys.calc_psf(display=display)
-        if display: plt.title("Wavefront rotated {} deg".format(angle))
+        osys2 = poppy.OpticalSystem(npix=npix)
+        osys2.add_pupil(poppy.optics.ParityTestAperture())
+        osys2.add_rotation(angle=angle)  # note, same sign here.
+        osys2.add_detector(fov_pixels=128, pixelscale=0.01)
+        psf2 = osys2.calc_psf()
 
+        if display:
+            fig, axes = plt.subplots(figsize=(16, 5), ncols=2)
+            poppy.display_psf(psf1, ax=axes[0])
+            axes[0].set_title("Optic rotated {} deg".format(angle))
+            poppy.display_psf(psf2, ax=axes[1])
+            axes[1].set_title("Wavefront rotated {} deg".format(angle))
 
         assert np.allclose(psf1[0].data, psf2[0].data, atol=atol), ("PSFs did not agree "
-            "within the requested tolerance")
+                                                                    f"within the requested tolerance, for angle={angle}."
+                                                                    f"Max |difference| = {np.max(np.abs(psf1[0].data - psf2[0].data))}")
 
 ### Tests for OpticalElements defined in poppy_core###
 
@@ -516,6 +518,46 @@ def test_OPD_in_waves_for_FITSOpticalElement():
         psf = osys.calc_psf(wavelength=prefactor * u.um)
         assert np.isclose(center_pixel_value, psf[0].data[1,1])
 
+def test_fits_rot90_vs_ndimagerotate_consistency(plot=False):
+    """Test that rotating a FITS HDUList via either of the two
+    methods yields consistent results. This compares an exact
+    90 degree rotation and an interpolating not-quite-90-deg rotation.
+    Both methods should rotate counterclockwise and consistently.
+    """
+    letterf_hdu = poppy.optics.LetterFAperture().to_fits(npix=128)
+    opt1 = poppy.FITSOpticalElement(transmission=letterf_hdu,
+                                   rotation=90)
+    opt2 = poppy.FITSOpticalElement(transmission=letterf_hdu,
+                                   rotation=89.99999)
+    assert np.allclose(opt1.amplitude, opt2.amplitude, atol=1e-5)
+
+    if plot:
+        fig, axes = plt.subplots(figsize=(10, 5), ncols=2)
+        axes[0].imshow(opt1.amplitude)
+        axes[0].set_title("Rot90")
+        axes[1].imshow(opt2.amplitude)
+        axes[1].set_title("ndimage rotate(89.9999)")
+
+def test_analytic_vs_FITS_rotation_consistency(plot=False):
+    """Test that rotating an AnalyticOpticalElement vs
+    rotating a discretized version as a FITSOpticalElement
+    are consistent in rotation direction (counterclockwise)
+    and amount"""
+    opt1 = poppy.optics.LetterFAperture(rotation=90)
+
+    letterf_hdu = poppy.optics.LetterFAperture().to_fits(npix=128)
+    opt2 = poppy.FITSOpticalElement(transmission=letterf_hdu,
+                                    rotation=90)
+
+    if plot:
+        opt1.display()
+        plt.figure()
+        opt2.display()
+
+    array1 = opt1.sample(npix=128)
+    array2 = opt2.amplitude
+    assert np.allclose(array1, array2)
+
 ### OpticalSystem tests and related
 
 def test_source_offsets_in_OpticalSystem(npix=128, fov_size=1):
diff --git a/poppy/tests/test_wavefront.py b/poppy/tests/test_wavefront.py
index d569a81..62bbf4d 100644
--- a/poppy/tests/test_wavefront.py
+++ b/poppy/tests/test_wavefront.py
@@ -1,4 +1,5 @@
 
+import poppy
 from .. import poppy_core
 from .. import optics
 import numpy as np
@@ -62,6 +63,34 @@ def test_wavefront_rotation():
     wave *= rot
     assert np.allclose(wave0.wavefront, wave.wavefront)
 
+def test_wavefront_rot90_vs_ndimagerotate_consistency(plot=False):
+    """Test that rotating a Wavefront via either of the two
+    methods yields consistent results. This compares an exact
+    90 degree rotation and an interpolating not-quite-90-deg rotation.
+    Both methods should rotate counterclockwise and consistently.
+    """
+    letterf = poppy.optics.LetterFAperture()
+    wave = poppy.Wavefront(diam=3 * u.m, npix=128)
+    wave *= letterf
+    wave2 = wave.copy()
+
+    wave.rotate(90)
+    wave2.rotate(89.99999)
+
+    assert np.allclose(wave2.intensity, wave.intensity, atol=1e-5), "Inconsistent results from the two rotation methods"
+
+    from poppy.tests.test_sign_conventions import brighter_top_half, brighter_left_half
+
+    assert brighter_left_half(wave.intensity), "Rotated wavefront orientation not as expected"
+    assert not brighter_top_half(wave.intensity), "Rotated wavefront orientation not as expected"
+
+    if plot:
+        fig, axes = plt.subplots(figsize=(10, 5), ncols=2)
+        wave.display(ax=axes[0])
+        wave2.display(ax=axes[1])
+        axes[0].set_title("Rot90")
+        axes[1].set_title("ndimage rotate(89.9999)")
+
 def test_wavefront_inversion():
     npix=100
     wave = poppy_core.Wavefront(npix=npix, wavelength=1e-6)