Merge branch 'vector' of https://github.com/kvangorkom/poppy into vector

poppy/optics.py

@@ -2294,15 +2294,15 @@ class LinearPolarizer(PolarizationOpticalElement, AnalyticOpticalElement):
 
     Parameters
     ----------
-    name : string
-        Descriptive name
     angle: float
         Polarization axis angle, in radians.
-    extinction : float
-        NOT IMPLEMENTED: Extinction ratio. Default is infinite (perfect linear polarizer).
+    extinction : float, optional
+        Extinction ratio. Default is infinite (perfect linear polarizer).
+    name : string, optional
+        Descriptive name
     """
 
-    def __init__(self, name=None, angle=0, extinction=xp.inf, **kwargs):
+    def __init__(self, angle, name=None, extinction=xp.inf, **kwargs):
         if name is None:
             name = "Linear polarizer"
         self.angle = angle
@@ -2316,8 +2316,6 @@ def get_jones_matrix(self, wave):
         cth = xp.cos(self.angle)
         sth = xp.sin(self.angle)
         eps = 1/self.extinction
-        #self.jones_matrix = xp.asarray([[cth**2,  sth*cth],
-        #                                [sth*cth, sth**2]])
         self.jones_matrix = xp.asarray([[cth**2 + eps*sth**2,  (1-eps)*sth*cth    ],
                                         [(1-eps)*sth*cth,      eps*cth**2 + sth**2]])
         return self.jones_matrix
@@ -2335,6 +2333,8 @@ class CircularPolarizer(PolarizationOpticalElement, AnalyticOpticalElement):
         Descriptive name
     handedness: str
         Either 'left' or 'right'
+    name : string, optional
+        Descriptive name
     """
 
     def __init__(self, handedness, name=None, **kwargs):
@@ -2362,12 +2362,12 @@ class LinearPhaseRetarder(PolarizationOpticalElement, AnalyticOpticalElement):
 
     Parameters
     ----------
-    name : string
-        Descriptive name
     phase : float
         Phase retardance, in radians
     angle: float
         Fast axis angle, in radians.
+    name : string, optional
+        Descriptive name
     """
 
     def __init__(self, phase, angle, name=None, **kwargs):
@@ -2393,13 +2393,13 @@ class QuarterWavePlate(LinearPhaseRetarder):
 
     Parameters
     ----------
-    name : string
-        Descriptive name
     angle: float
-        Fast axis angle, in radians.
+        Fast axis angle, in radians
+    name : string, optional
+        Descriptive name
     """
 
-    def __init__(self, name=None, angle=0):
+    def __init__(self, angle, name=None):
         if name is None:
             name = "Quarter wave plate"
         super(QuarterWavePlate, self).__init__(np.pi/2, angle, name=name)
@@ -2409,13 +2409,13 @@ class HalfWavePlate(LinearPhaseRetarder):
 
     Parameters
     ----------
-    name : string
-        Descriptive name
     angle: float
         Fast axis angle, in radians.
+    name : string, optional
+        Descriptive name
     """
 
-    def __init__(self, name=None, angle=0):
+    def __init__(self, angle, name=None):
         if name is None:
             name = "Half wave plate"
         super(HalfWavePlate, self).__init__(np.pi, angle, name=name,)
@@ -2448,9 +2448,19 @@ class VectorVortexMask(LinearPhaseRetarder):
     approximate better sampling in the vicinity of the vortex singularity
     and is likely to be limited by numerical artifacts without extreme
     sampling.
+
+    Parameters
+    ----------
+    charge : int, optional
+        The charge of the vector vortex. Default: 6
+    retardance : float, optional
+        Global retardance of the vector vortex. Default is pi, for a
+        half-wave plate.
+    name : string, optional
+        Descriptive name
     """
 
-    def __init__(self, charge=6, retardance=np.pi, name=None, **kwargs):
+    def __init__(self, charge=6, retardance=xp.pi, name=None, **kwargs):
         if name is None:
             name = "VVC"
         self.charge = charge

poppy/polarized_wavefront.py

@@ -1,24 +1,11 @@
-'''
-TO DO:
-* write display code for polarization optics and polarized wavefronts
-    * compatibility with existing display functions
-    * new display functions for Stokes and vector WFs?
-* more tests!
-    * worried about interactions of polarized WFs with various types of optical elements
-    * test cases aren't checking all basic stokes combos yet
-* create polarized versions of some of the standard coronagraph masks
-* put together examples using Fraun+Fresnel prop (vector/stokes), complicated optical system
-* resolve all other TO DOs or NOT IMPLEMENTEDs
-'''
-
 import numpy as np
 import astropy.units as u
 
 from poppy.poppy_core import Wavefront, BaseWavefront
 from poppy.fresnel import FresnelWavefront
 
 from . import accel_math
-from .accel_math import xp, ensure_not_on_gpu
+from .accel_math import xp
 
 if accel_math._NUMEXPR_AVAILABLE:
     import numexpr as ne
@@ -49,10 +36,8 @@ def __init__(self,
         super(BasePolarizedWavefront, self).__init__(
             **kwargs
         )
-        # TO DO: clean up the logic of checking which is specified and handling appropriately
-        #self.input_stokes_vector = input_stokes_vector
-        #self.input_polarization = input_polarization
-        self.pol_type = None
+        self.input_stokes_vector = None
+        self.input_polarization = None
 
         if input_stokes_vector is not None: # wavefront tensor
             self.input_polarization = None
@@ -228,7 +213,7 @@ def jones_to_stokes(jones_matrix, input_stokes_vector):
     return xp.einsum('i...,i...', M, input_stokes_vector).real
 
 def jones_to_mueller(jones_matrix):
-    """ Convert 2x2 Jones matrix to Stokes parameters
+    """ Convert 2x2 Jones matrix to Mueller matrix
     
     Based on Eqn A4.13, Spectroscopic Ellipsometry: Principles and Applications H. Fujiwara
     """

poppy/poppy_core.py

@@ -3361,12 +3361,7 @@ def get_opd(self, wave):
         return xp.asarray(self.opd)
 
 class PolarizationOpticalElement(OpticalElement):
-    """ Abstract class for defining polarization optics.
-
-    TO DO:
-    * Add a check that PolarizationOpticalElements are interacting with PolarizedWavefronts
-    (i.e., reject interactions with Wavefront or FresnelWavefront objects)
-    """
+    """ Abstract class for defining polarization optics """
 
     def __init__(self, **kwargs):
         OpticalElement.__init__(self, **kwargs)

poppy/tests/test_fft.py

@@ -197,9 +197,9 @@ def test_pyfftw_vs_numpyfft(verbose=False):
     """ Create an optical system with 2 parity test apertures,
     propagate light through it, and compare that we get the same results from both numpy and pyfftw"""
 
-    defaults = conf.use_fftw, conf.use_cuda, conf.use_opencl, conf.use_cupy
+    defaults = conf.use_fftw, conf.use_cupy, conf.use_opencl, conf.use_cupy
 
-    conf.use_cuda = False
+    conf.use_cupy = False
     conf.use_opencl = False
     conf.use_cupy = False
 
@@ -231,18 +231,18 @@ def test_pyfftw_vs_numpyfft(verbose=False):
             print(" Int. WF {} intensity diff: {}".format(i, np.abs(intermediates[i].total_intensity-total_int_input)) )
         print(" Final PSF intensity diff:", np.abs(intermediates[3].total_intensity-total_int_input) - expected)
 
-    conf.use_fftw, conf.use_cuda, conf.use_opencl, conf.use_cupy = defaults
+    conf.use_fftw, conf.use_cupy, conf.use_opencl, conf.use_cupy = defaults
 
 
 @pytest.mark.skipif(accel_math._OPENCL_AVAILABLE is False, reason="OPENCL not available")
 def test_opencl_vs_numpyfft(verbose=False):
     """ Create an optical system with 2 parity test apertures,
     propagate light through it, and compare that we get the same results from both numpy and CUDA"""
 
-    defaults = conf.use_fftw, conf.use_cuda, conf.use_opencl
+    defaults = conf.use_fftw, conf.use_cupy, conf.use_opencl
 
     conf.use_fftw = False
-    conf.use_cudal = False
+    conf.use_cupy = False
 
     sys = setup_test_osys()
 
@@ -272,7 +272,7 @@ def test_opencl_vs_numpyfft(verbose=False):
             print(" Int. WF {} intensity diff: {}".format(i, np.abs(intermediates[i].total_intensity-total_int_input)) )
         print(" Final PSF intensity diff:", np.abs(intermediates[3].total_intensity-total_int_input) - expected)
 
-    conf.use_fftw, conf.use_cuda, conf.use_opencl = defaults
+    conf.use_fftw, conf.use_cupy, conf.use_opencl = defaults
 
 
 # TODO: Add a function that uses both the DFT and MFT for the exact same calc, and compare the results

poppy/tests/test_polarization.py

@@ -124,11 +124,11 @@ def test_qwp():
     """
     # linear to circular
     wf = poppy.PolarizedWavefront(diam=1, npix=1, input_stokes_vector=(1,1,0,0))
-    qwp = poppy.QuarterWavePlate(angle=xp.pi/4)
+    qwp = poppy.QuarterWavePlate(xp.pi/4)
     assert xp.allclose( xp.squeeze((wf * qwp).stokes_parameters), [1,0,0,1] )
     # circular to linear
     wf = poppy.PolarizedWavefront(diam=1, npix=1, input_stokes_vector=(1,0,0,-1))
-    qwp = poppy.QuarterWavePlate(angle=xp.pi/4)
+    qwp = poppy.QuarterWavePlate(xp.pi/4)
     assert xp.allclose( xp.squeeze((wf * qwp).stokes_parameters), [1,1,0,0] )
 
 def test_hwp():
@@ -137,13 +137,31 @@ def test_hwp():
     """
     # linear 2 theta rotation
     wf = poppy.PolarizedWavefront(diam=1, npix=1, input_stokes_vector=(1,1,0,0))
-    hwp = poppy.HalfWavePlate(angle=xp.pi/4)
+    hwp = poppy.HalfWavePlate(xp.pi/4)
     assert xp.allclose( xp.squeeze((wf * hwp).stokes_parameters), [1,-1,0,0] )
     # circular handedness flip
     wf = poppy.PolarizedWavefront(diam=1, npix=1, input_stokes_vector=(1,0,0,-1))
-    hwp = poppy.HalfWavePlate(angle=0)
+    hwp = poppy.HalfWavePlate(0)
     assert xp.allclose( xp.squeeze((wf * hwp).stokes_parameters), [1,0,0,1] )
 
+
+def test_jones_matrix():
+    """
+    Test tensor fields interacting with JonesMatrixOpticalElements.
+
+    Note that this doesn't test any values -- it just exercises the code
+    """
+    npix = 256
+    wf = poppy.PolarizedWavefront(diam=1, npix=npix, input_stokes_vector=(1,0,0,0))
+
+    zbasis = poppy.zernike.zernike_basis(nterms=3, npix=npix, outside=0)
+    tilt = 0.5 # waves
+    jones_matrix = xp.asarray([[xp.exp(1j*tilt*2*xp.pi*zbasis[1]),    xp.zeros((npix,npix), dtype=complex)],
+                            [xp.zeros((npix,npix), dtype=complex), xp.exp(-1j*tilt*2*xp.pi*zbasis[1]) ]])
+    jm = poppy.JonesMatrixOpticalElement(jones_matrix)
+
+    wf_out = wf * jm
+
 # ---- Fresnel + Stokes (Partial Polarization) ---
 
 def test_fresnel_stokes_linearpolarizer():
@@ -195,7 +213,7 @@ def test_fresnel_stokes_qwp():
     for i in range(len(qwp_angles)):
         osys = fresnel.FresnelOpticalSystem(npix=npix)
         circ = poppy.CircularAperture(radius=D)
-        qwp = poppy.QuarterWavePlate(angle=qwp_angles[i])
+        qwp = poppy.QuarterWavePlate(qwp_angles[i])
         osys.add_optic(circ)
         osys.add_optic(qwp, distance=500*u.mm)
 
@@ -225,7 +243,7 @@ def test_fresnel_stokes_hwp():
     for i in range(len(stokes_true)):
         osys = fresnel.FresnelOpticalSystem(npix=npix)
         circ = poppy.CircularAperture(radius=D)
-        hwp = poppy.HalfWavePlate()
+        hwp = poppy.HalfWavePlate(0)
         osys.add_optic(circ)
         osys.add_optic(hwp, distance=500*u.mm)
 
@@ -289,15 +307,15 @@ def test_fresnel_vector_hwp():
     input_vector = [(1, 0), (f, 1j*f), (f, -1j*f)]
     #output_vector = [(0, 1), (f, -1j*f), (f, 1j*f)]
     filter_out = [
-        poppy.LinearPolarizer(angle=xp.pi/2.),
+        poppy.LinearPolarizer(xp.pi/2.),
         poppy.CircularPolarizer(handedness='right'),
         poppy.CircularPolarizer(handedness='left')
         ]
 
     for i in range(len(input_vector)):
         osys = fresnel.FresnelOpticalSystem(npix=npix)
         circ = poppy.CircularAperture(radius=D)
-        hwp = poppy.HalfWavePlate(angle=hwp_angle)
+        hwp = poppy.HalfWavePlate(hwp_angle)
         osys.add_optic(circ)
         osys.add_optic(hwp, distance=500*u.mm)
         osys.add_optic(filter_out[i])
@@ -392,7 +410,7 @@ def test_fraunhofer_stokes_hwp():
     for i in range(len(stokes_true)):
         osys = poppy.OpticalSystem(npix=npix)
         circ = poppy.CircularAperture(radius=D)
-        hwp = poppy.HalfWavePlate()
+        hwp = poppy.HalfWavePlate(0)
         osys.add_pupil(circ)
         osys.add_pupil(hwp)
         osys.add_image()
@@ -459,15 +477,15 @@ def test_fraunhofer_vector_hwp():
     input_vector = [(1, 0), (f, 1j*f), (f, -1j*f)]
     #output_vector = [(0, 1), (f, -1j*f), (f, 1j*f)]
     filter_out = [
-        poppy.LinearPolarizer(angle=xp.pi/2.),
+        poppy.LinearPolarizer(xp.pi/2.),
         poppy.CircularPolarizer(handedness='right'),
         poppy.CircularPolarizer(handedness='left')
         ]
 
     for i in range(len(input_vector)):
         osys = poppy.OpticalSystem(npix=npix)
         circ = poppy.CircularAperture(radius=D)
-        hwp = poppy.HalfWavePlate(angle=hwp_angle)
+        hwp = poppy.HalfWavePlate(hwp_angle)
         osys.add_pupil(circ)
         osys.add_pupil(hwp)
         osys.add_image()