python/optics.py

# -*- coding: utf-8 -*-
"""
Created on Fri Jun 22 16:22:17 2018
@author: david
"""
import numpy as np
import scipy as sp
import scipy.ndimage
import matplotlib
import math
import matplotlib.pyplot as plt
#adjust the components of E0 so that they are orthogonal to d
    #if d is 2D, the z component is calculated such that ||[dx, dy, dz]||=1
def orthogonalize(E, d):
    if d.shape[-1] == 2:
        dz = 1 - d[..., 0]**2 - d[..., 1]**2
        d = np.append(d, dz)
    s = np.cross(E, d)
    
    dE = np.cross(d, s)
    return dE
    #return dE/np.linalg.norm(dE) * np.linalg.norm(E)
def intensity(E):
    Econj = np.conj(E)
    I = np.sum(E*Econj, axis=-1)
    return np.real(I)
class planewave:
    #implement all features of a plane wave
    #   k, E, frequency (or wavelength in a vacuum)--l
    #   try to enforce that E and k have to be orthogonal
    
    #initialization function that sets these parameters when a plane wave is created
    def __init__ (self, k, E):
        
        self.k = np.array(k)              
        self.E = E
        
        #if ( np.linalg.norm(k) > 1e-15 and np.linalg.norm(E) >1e-15):
        #    s = np.cross(k, E)              #compute an orthogonal side vector
        #    s = s / np.linalg.norm(s)       #normalize it
        #    Edir = np.cross(s, k)              #compute new E vector which is orthogonal
        #    self.k = np.array(k)
        #    self.E = Edir / np.linalg.norm(Edir) * np.linalg.norm(E)
    
    def __str__(self):
        return str(self.k) + "\n" + str(self.E)     #for verify field vectors use print command
    def kmag(self):
        return math.sqrt(self.k[0] ** 2 + self.k[1] ** 2 + self.k[2] ** 2)
    
    #function that renders the plane wave given a set of coordinates
    def evaluate(self, X, Y, Z):
        k_dot_r = self.k[0] * X + self.k[1] * Y + self.k[2] * Z     #phase term k*r
        ex = np.exp(1j * k_dot_r)       #E field equation  = E0 * exp (i * (k * r)) here we simply set amplitude as 1
        Ef = ex[..., None] * self.E
        return Ef
class objective:
    
    #initialize an objective lens
    #   sin_alpha is the sine of the angle subtended by the objective (numerical aperture in a refractive index of 1.0)
    #   n is the real refractive index of the imaging media (default air)
    #   N is the resolution of the 2D objective map (field at the input and output aperture)
    def __init__(self, sin_alpha, n, N):
        
        #set member variables
        self.sin_alpha = sin_alpha
        self.n = n
        
        #calculate the transfer function from the back aperture to the front aperture
        self.N = N
        self._CI(N)
        self._CR(N)        
    
       
    #calculates the basis vectors describing the mapping between the back and front aperture
    def _calc_bases(self, N):
        
        #calculate the (sx, sy) coordinates (angular spectrum) corresponding to each position in the entrance/exit pupil
        s = np.linspace(-self.sin_alpha, self.sin_alpha, N)
        [SX, SY] = np.meshgrid(s, s)
        S_SQ = SX**2 + SY**2                                  #calculate sx^2 + sy^2
        
        BAND_M = S_SQ <= self.sin_alpha**2                    #find a mask for pixels within the bandpass of the objective
        self.BAND_M = BAND_M
        
        #pull out pixels that pass through the objective bandpass        
        SX = SX[BAND_M]
        SY = SY[BAND_M]
        S_SQ = S_SQ[BAND_M]
        SZ = np.sqrt(1 - S_SQ)                            #calculate sz = sqrt(1 - sx^2 - sy^2) in normalized (-1, 1) coordinates
        
        #calculate the apodization functions for each polarization
        self.FS = 1.0 / np.sqrt(SZ)
        self.FP = self.FS
        
        #calculate the S and P basis vectors
        THETA = np.arctan2(SY, SX)
        
        SX_SXSY = np.cos(THETA)
        SY_SXSY = np.sin(THETA)
        
        v_size = (len(SX), 3)
        VS_VSP = np.zeros(v_size)
        VP = np.zeros(v_size)
        VPP = np.zeros(v_size)
        
        VS_VSP[:, 0] = -SY_SXSY
        VS_VSP[:, 1] = SX_SXSY
        
        VP[:, 0] = -SX_SXSY
        VP[:, 1] = -SY_SXSY
        
        SXSY = np.sqrt(S_SQ)
        VPPZ = np.zeros(SXSY.shape)
        DC_M = SXSY > 0
        VPPZ[DC_M] = S_SQ[DC_M] / SXSY[DC_M]
        VPP[:, 0] = -SX_SXSY * SZ
        VPP[:, 1] = -SY_SXSY * SZ
        VPP[:, 2] = VPPZ
        
        return [VS_VSP, VP, VPP]        
        
        
    #calculate the transfer function that maps the back aperture field to the front aperture field (Eq. 19 in Davis et al.)
    def _CI(self, N):
        [VS_VSP, VP, VPP] = self._calc_bases(N)
        
        VSP_VST = np.matmul(VS_VSP[:, :, None], VS_VSP[:, None, :])
        VPP_VPT = np.matmul(VPP[:, :, None], VP[:, None, :])
        
        self.CI = self.FS[:, None, None] * VSP_VST + self.FP[:, None, None] * VPP_VPT      
        
    def _CR(self, N):
        [VS_VSP, VP, VPP] = self._calc_bases(N)
        
        VSP_VST = np.matmul(VS_VSP[:, :, None], VS_VSP[:, None, :])
        VP_VPPT = np.matmul(VP[:, :, None], VPP[:, None, :])
        
        self.CR = 1.0 / self.FS[:, None, None] * VSP_VST + 1.0 / self.FP[:, None, None] * VP_VPPT
        
    #condenses a field defined at the back aperture of the objective
    #   E is a 3xNxN image of the field at the back aperture
    def back2front(self, E):
        Ein = E[self.BAND_M, :, None]
        
        Efront = np.zeros((self.N, self.N, 3), dtype=np.complex128)
        Efront[self.BAND_M, :] = np.squeeze(np.matmul(self.CI, Ein))
        
        return Efront
    
        #propagates a field from the front aperture of the objective to the back aperture
    #   E is a 3xNxN image of the field at the front aperture
    def front2back(self, E):
        Ein = E[self.BAND_M, :, None]
        
        Eback = np.zeros((self.N, self.N, 3), dtype=np.complex128)
        Eback[self.BAND_M, :] = np.squeeze(np.matmul(self.CR, Ein))
        
        return Eback
    
    #calculates the image at the focus given the image at the back aperture
    #   E is the field at the back aperture of the objective
    #   k0 is the wavenumber in a vacuum (2pi/lambda_0)
    def focus(self, E, k0, RX, RY, RZ):
        
        #calculate the angular spectrum at the front aperture
        Ef = self.back2front(E)
        
        #generate an array describing the valid components of S
        s = np.linspace(-self.sin_alpha, self.sin_alpha, self.N)
        [SX, SY] = np.meshgrid(s, s)
        S = [SX[self.BAND_M], SY[self.BAND_M]]        
        SXSY_SQ = S[0]**2 + S[1]**2
        S.append(np.sqrt(1 - SXSY_SQ))
        S = np.array(S)
        
        
        #calculate the k vector (accounting for both k0 and the refractive index of the immersion medium)
        K = k0 * self.n * np.transpose(S)[:, None, :]
        
        #calculate the dot product of K and R
        R = np.moveaxis(np.array([RX, RY, RZ]), 0, -1)[..., None, :, None]
        
        K_DOT_R = np.squeeze(np.matmul(K, R))
        EX = np.exp(1j * K_DOT_R)[..., None]
        Ewaves = EX * Ef[self.BAND_M]
        
        Efocus = np.sum(Ewaves, -2)
        return Efocus
class gaussian_beam:
    
    #generate a beam
    #   (Ex, Ey) specifies the polarization
    #   p specifies the focal point
    #   l is the free-space wavelength
    #   sigma_p is the standard deviation of the Gaussian beam at the focal point
    #   amp is the amplitude of the DC component of the Gaussian beam
    def __init__(self, x_polar, y_polar, lambda_0=1, amp=1, sigma_p=None, sigma_k=None):
        self.E = [x_polar, y_polar]
        self.l = lambda_0
        
        #calculate the standard deviation of the Gaussian function in k-space 
        if(sigma_p is None and sigma_k is None):
            self.sigma = 2 * math.pi / lambda_0
        elif(sigma_p is not None):
            self.sigma = 1.0 / sigma_p
        else:
            self.sigma = sigma_k
            
        #self.s = self.sigma_p2k(sigma_p)
        self.A = amp
        
    #calculates the maximum k-space frequency supported by this beam    
    def kmax(self):
        return 2 * np.pi / self.l
    
    #calculate the standard deviation of the Gaussian function in k space given the std at the focal point
   # def sigma_p2k(self, sigma):
        
        # this code is based on: http://www.robots.ox.ac.uk/~az/lectures/ia/lect2.pdf
        # f(r) = 1 / (2 pi sigma^2) e^(-r^2 / (2 sigma^2))
        # F(kx, ky) = f( |k_par| ) = e^( -2 pi^2 |k_par|^2 sigma^2 )
        
        # Basically, the Fourier transform of a normally distributed Gaussian is
        #   a non-scaled Gaussian with standard deviation 1/(2 pi sigma)        
        #return 1.0/(2 * math.pi * sigma)
        #return 1.0 / sigma
      
    
    def kspace(self, Kx, Ky):
        #sigma = self.s * (np.pi / self.l)
        #G = 1 / (2 * np.pi * sigma **2) * np.exp(- (0.5/(sigma**2)) * (Kx ** 2 + Ky ** 2))
        G = self.A * np.exp(- (Kx ** 2 + Ky ** 2) / (2 * self.sigma ** 2))
        
        #calculate k-vectors outside of the band-limit
        B = ((Kx ** 2 + Ky **2) > self.kmax() ** 2)
        
        #calculate Kz (required for calculating the Ez component of E)
        #Kz = np.lib.scimath.sqrt(self.kmax() ** 2 - Kx**2 - Ky**2)
        
        Ex = self.E[0] * G
        Ey = self.E[1] * G
        #Ez = -Kx/Kz * Ex
        Ez = np.zeros(Ex.shape)
        
        #set all values outside of the band limit to zero
        Ex[B] = 0
        Ey[B] = 0
        Ez[B] = 0
        
        return [Ex, Ey, Ez]
    
    # The PSF at the focal point is the convolution of a Bessel function and a Gaussian
    # This function calculates the parameters of the Gaussian function in terms of a*e^{-x^2 / (2*c^2)}
    #   a (return) is the scale factor for the Gaussian
    #   c (return) is the standard deviation
    def psf_gaussian(self):
        # this code is based on: http://mathworld.wolfram.com/FourierTransformGaussian.html
        #a = math.sqrt( 2 * math.pi * (self.s ** 2) )
        #c = 1.0 / (2 * math.pi * self.s)
        
        a = math.sqrt( (self.sigma ** 2) )
        c = 1.0 / (self.sigma)
        
        return [a, c]
    
    # The PSF at the focal point is the convolution of a Bessel function and a Gaussian
    # This function calculates the parameters of the bessel function in terms of a J_1(pi * a * x) / x
    def psf_bessel(self):
        return self.kmax() / math.pi  
    
    # calculate the beam point spread function at the focal plane
    #   D is the image extent for the square: (-D, D, -D, D)
    #   N is the number of spatial samples along one dimension of the image (NxN)
    #   order is the mathematical order of the calculation (in terms of the resolution of the FFT)
    def psf(self, D, N):
        
        [a, c] = self.psf_gaussian()
        alpha = self.psf_bessel()
        
        #calculate the sample positions along the x-axis
        x = np.linspace(-D, D, N)
        dx = x[1] - x[0]
        
        [X, Y] = np.meshgrid(x, x)
        R = np.sqrt(X**2 + Y**2)
        
        bessel = alpha * sp.special.j1(math.pi * alpha * R) / R
        
        return a * sp.ndimage.filters.gaussian_filter(bessel, c / dx)    
        
    
    #return the k-space transform of the beam at p as an NxN array
    def kspace_image(self, N):
        #calculate the band limit [-2*pi/lambda, 2*pi/lambda]
        kx = np.linspace(-self.kmax(), self.kmax(), N)
        ky = kx
        
        #generate a meshgrid for the field
        [Kx, Ky] = np.meshgrid(kx, ky)
        #generate a Gaussian with a standard deviation of sigma * pi/lambda
        return self.kspace(Kx, Ky)
    
    # calculate the practical band limit of the beam such that 
    # all frequencies higher than the returned values are less than epsilon
    def practical_band_limit(self, epsilon = 0.001):
        
        return min(self.sigma * math.sqrt(math.log(1.0/epsilon)), self.kmax())
        
        
    #return a plane wave decomposition of the beam using N plane waves
    def decompose(self, N):
        PW = []                     #create an empty list of plane waves
        
        # sample a cylinder and project those samples onto a unit sphere
        theta = np.random.uniform(0, 2*np.pi, N)
        
        # get the practical band limit to minimize sampling
        band_limit = self.practical_band_limit() / self.kmax()
        phi = np.arccos(np.random.uniform(1 - band_limit, 1, N))
        
        kx = -self.kmax() * np.sin(phi) * np.cos(theta)
        ky = -self.kmax() * np.sin(phi) * np.sin(theta)
        
        mc_weight = 2 * math.pi / N
        E0 = self.kspace(kx, ky)
        kz = np.sqrt(self.kmax() ** 2 - kx**2 - ky**2)
        
        for i in range(0, N):
            k = [kx[i], ky[i], kz[i]]
            E = [E0[0][i] * mc_weight, E0[1][i] * mc_weight, E0[2][i] * mc_weight]
            w = planewave(k, E)
            PW.append(w)
            
        return PW
    
class layersample:
    
    #generate a layered sample based on a list of layer refractive indices
    #   layer_ri is a list of complex refractive indices for each layer
    #   z is the list of boundary positions along z
    def __init__(self, layer_ri, z):
        self.n = np.array(layer_ri).astype(np.complex128)
        self.z = np.array(z)
    
    #calculate the index of the field component associated with a layer
    #   l is the layer index [0, L)
    #   c is the component (x=0, y=1, z=2)
    #   d is the direction (0 = transmission, 1 = reflection)
    def i(self, l, c, d):
        i = l * 6 + d * 3 + c - 3
        return  i
    # generate the linear system corresponding to this layered sample and plane wave
    #   s is the direction vector scaled by the refractive index
    #   k0 is the free-space k-vector
    #   E is the field vector for the incident plane wave
    def generate_linsys(self, s, k0, E):
        #allocate space for the matrix
        L = len(self.n)
        M = np.zeros((6*(L-1), 6*(L-1)), dtype=np.complex128)
        
        #allocate space for the RHS vector
        b = np.zeros(6*(L-1), dtype=np.complex128)
        
        #initialize a counter for the equation number
        ei = 0
        
        #calculate the sz component for each layer
        self.sz = np.zeros(L, dtype=np.complex128)
        for l in range(L):
            self.sz[l] = np.sqrt(self.n[l]**2 - s[0]**2 - s[1]**2)
        
        #set constraints based on Gauss' law
        for l in range(0, L):
            #sz = np.sqrt(self.n[l]**2 - s[0]**2 - s[1]**2)
            
            #set the upward components for each layer
            #   note that layer L-1 does not have a downward component
            #   David I, Equation 7
            if l != L-1:
                M[ei, self.i(l, 0, 1)] = s[0]
                M[ei, self.i(l, 1, 1)] = s[1]
                M[ei, self.i(l, 2, 1)] = -self.sz[l]
                ei = ei+1
            
            #set the downward components for each layer
            #   note that layer 0 does not have a downward component
            #   Davis I, Equation 6
            if l != 0:
                M[ei, self.i(l, 0, 0)] = s[0]
                M[ei, self.i(l, 1, 0)] = s[1]
                M[ei, self.i(l, 2, 0)] = self.sz[l]
                ei = ei+1            
            
        #enforce a continuous field across boundaries
        for l in range(1, L):
            sz0 = self.sz[l-1]
            sz1 = self.sz[l]
            A = np.exp(1j * k0 * sz0 * (self.z[l] - self.z[l-1]))
            if l < L - 1:
                dl = self.z[l] - self.z[l+1]
                arg = -1j * k0 * sz1 * dl
                B = np.exp(arg)
            
            
            #if this is the second layer, use the simplified equations that account for the incident field
            if l == 1:
                M[ei, self.i(0, 0, 1)] = 1
                M[ei, self.i(1, 0, 0)] = -1
                if L > 2:
                    #print(-B, M[ei, self.i(1, 0, 1)])
                    M[ei, self.i(1, 0, 1)] = -B
                    
                b[ei] = -A * E[0]
                ei = ei + 1
                
                M[ei, self.i(0, 1, 1)] = 1
                M[ei, self.i(1, 1, 0)] = -1
                if L > 2:
                    M[ei, self.i(1, 1, 1)] = -B
                b[ei] = -A * E[1]
                ei = ei + 1
                
                M[ei, self.i(0, 2, 1)] = s[1]
                M[ei, self.i(0, 1, 1)] = sz0
                M[ei, self.i(1, 2, 0)] = -s[1]
                M[ei, self.i(1, 1, 0)] = sz1
                if L > 2:
                    M[ei, self.i(1, 2, 1)] = -B*s[1]
                    M[ei, self.i(1, 1, 1)] = -B*sz1
                b[ei] = A * sz0 * E[1] - A * s[1]*E[2]
                ei = ei + 1
                
                M[ei, self.i(0, 0, 1)] = -sz0
                M[ei, self.i(0, 2, 1)] = -s[0]
                M[ei, self.i(1, 0, 0)] = -sz1
                M[ei, self.i(1, 2, 0)] = s[0]
                if L > 2:
                    M[ei, self.i(1, 0, 1)] = B*sz1
                    M[ei, self.i(1, 2, 1)] = B*s[0]
                b[ei] = A * s[0] * E[2] - A * sz0 * E[0]
                ei = ei + 1
            
            #if this is the last layer, use the simplified equations that exclude reflections from the last layer
            elif l == L-1:
                M[ei, self.i(l-1, 0, 0)] = A
                M[ei, self.i(l-1, 0, 1)] = 1
                M[ei, self.i(l, 0, 0)] = -1
                ei = ei + 1
                
                M[ei, self.i(l-1, 1, 0)] = A
                M[ei, self.i(l-1, 1, 1)] = 1
                M[ei, self.i(l, 1, 0)] = -1
                ei = ei + 1
                
                M[ei, self.i(l-1, 2, 0)] = A*s[1]
                M[ei, self.i(l-1, 1, 0)] = -A*sz0
                M[ei, self.i(l-1, 2, 1)] = s[1]
                M[ei, self.i(l-1, 1, 1)] = sz0
                M[ei, self.i(l, 2, 0)] = -s[1]
                M[ei, self.i(l, 1, 0)] = sz1
                ei = ei + 1
                
                M[ei, self.i(l-1, 0, 0)] = A*sz0
                M[ei, self.i(l-1, 2, 0)] = -A*s[0]
                M[ei, self.i(l-1, 0, 1)] = -sz0
                M[ei, self.i(l-1, 2, 1)] = -s[0]
                M[ei, self.i(l, 0, 0)] = -sz1
                M[ei, self.i(l, 2, 0)] = s[0]
                ei = ei + 1
            #otherwise use the full set of boundary conditions
            else:
                M[ei, self.i(l-1, 0, 0)] = A
                M[ei, self.i(l-1, 0, 1)] = 1
                M[ei, self.i(l, 0, 0)] = -1
                M[ei, self.i(l, 0, 1)] = -B
                ei = ei + 1
                
                M[ei, self.i(l-1, 1, 0)] = A
                M[ei, self.i(l-1, 1, 1)] = 1
                M[ei, self.i(l, 1, 0)] = -1
                M[ei, self.i(l, 1, 1)] = -B
                ei = ei + 1
                
                M[ei, self.i(l-1, 2, 0)] = A*s[1]
                M[ei, self.i(l-1, 1, 0)] = -A*sz0
                M[ei, self.i(l-1, 2, 1)] = s[1]
                M[ei, self.i(l-1, 1, 1)] = sz0
                M[ei, self.i(l, 2, 0)] = -s[1]
                M[ei, self.i(l, 1, 0)] = sz1
                M[ei, self.i(l, 2, 1)] = -B*s[1]
                M[ei, self.i(l, 1, 1)] = -B*sz1
                ei = ei + 1
                
                M[ei, self.i(l-1, 0, 0)] = A*sz0
                M[ei, self.i(l-1, 2, 0)] = -A*s[0]
                M[ei, self.i(l-1, 0, 1)] = -sz0
                M[ei, self.i(l-1, 2, 1)] = -s[0]
                M[ei, self.i(l, 0, 0)] = -sz1
                M[ei, self.i(l, 2, 0)] = s[0]
                M[ei, self.i(l, 0, 1)] = B*sz1
                M[ei, self.i(l, 2, 1)] = B*s[0]
                ei = ei + 1
        
        return [M, b]
    
    #create a matrix for a single plane wave specified by k and E
    #   d = [dx, dy] are the x and y coordinates of the normalized direction of propagation
    #   k0 is the free space wave number (2 pi / lambda0)
    #   E is the electric field vector
    
    def solve1(self, d, k0, E):
                
        #s is the plane wave direction scaled by the refractive index
        s = np.array(d) * self.n[0]
               
        #store the matrix and RHS vector (for debugging)   
        [self.M, self.b] = self.generate_linsys(s, k0, E)
        #self.M = M
        #self.b = b
        
        #evaluate the linear system
        P = np.linalg.solve(self.M, self.b)
        
        #save the results (also for debugging)
        self.P = P
        
        #store the coefficients for each layer
        L = len(self.n)                             # calculate the number of layers
        self.Pt = np.zeros((3, L), np.complex128)
        self.Pr = np.zeros((3, L), np.complex128)
        for l in range(L):
            if l == 0:
                self.Pt[:, 0] = [E[0], E[1], E[2]]
            else:
                px = P[self.i(l, 0, 0)]
                py = P[self.i(l, 1, 0)]
                pz = P[self.i(l, 2, 0)]
                self.Pt[:, l] = [px, py, pz]
                
            if l == L-1:
                self.Pr[:, L-1] = [0, 0, 0]
            else:
                px = P[self.i(l, 0, 1)]
                py = P[self.i(l, 1, 1)]
                pz = P[self.i(l, 2, 1)]
                self.Pr[:, l] = [px, py, pz]
        
        #store values required for evaluation
        #store k
        self.k = k0
        
        #store sx and sy
        self.s = np.array([s[0], s[1]])
        
        self.solved = True
    
    #evaluate a solved homogeneous substrate
    def evaluate(self, X, Y, Z):
        
        if not self.solved:
            print("ERROR: the layered substrate hasn't been solved")
            return
        
        #this code is a bit cumbersome and could probably be optimized
        #   Basically, it vectorizes everything by creating an image
        #   such that the value at each pixel is the corresponding layer
        #   that the pixel resides in. That index is then used to calculate
        #   the field within the layer
        
        #allocate space for layer indices
        LI = np.zeros(Z.shape, dtype=np.int)
                
        #find the layer index for each sample point
        L = len(self.z)
        LI[Z < self.z[0]] = 0
        for l in range(L-1):
            idx = np.logical_and(Z > self.z[l], Z <= self.z[l+1])
            LI[idx] = l
            LI[Z > self.z[-1]] = L - 1
        
        #calculate the appropriate phase shift for the wave transmitted through the layer
        Ph_t = np.exp(1j * self.k * self.sz[LI] * (Z - self.z[LI]))
        
        #calculate the appropriate phase shift for the wave reflected off of the layer boundary
        LIp = LI + 1
        LIp[LIp >= L] = 0
        Ph_r = np.exp(-1j * self.k * self.sz[LI] * (Z - self.z[LIp]))
#        print(Ph_r)
        Ph_r[LI >= L-1] = 0
        
        #calculate the phase shift based on the X and Y positions
        Ph_xy = np.exp(1j * self.k * (self.s[0] * X + self.s[1] * Y))
        
        #apply the phase shifts
        Et = self.Pt[:, LI] * Ph_t[:, :]
        Er = self.Pr[:, LI] * Ph_r[:, :]
        
        #add everything together coherently
        E = (Et + Er) * Ph_xy[:, :]
        
        #return the electric field
        return np.moveaxis(E, 0, -1)
  
def example_psf():
    
    #specify the angle of accepted waves (NA if n = 1)
    sin_angle = 0.8
    
    #refractive index of the imaging medium
    n = 1.0
    
    #specify the size of the spectral domain (resolution of the front/back aperture)
    N = 100
    
    #create a new objective, specifying the NA (angle * refractive index) and resolution of the transform    
    O = objective(sin_angle, n, N)
    
    #specify the size of the spatial domain
    D = 10
    
    
    #specify the resolution of the spatial domain
    M = 100
    
    #free-space wavelength
    l0 = 1
    
    #calculate the free-space wavenumber
    k0 = 2 * np.pi * l0
    
    #specify the field incident at the aperture
    pap = np.array([1, 0, 0])                     #polarization
    sap = np.array([0, 0, 1])                  #incident angle of the plane wave (will be normalized)
    sap = sap / np.linalg.norm(sap)             #calculate the normalized direction vector for the plane wave
    
    k0ap = sap * k0                                 #calculate the k-vector of the plane wave incident on the aperture
    
    #create a plane wave
    pw = planewave(k0ap, pap)
    
    #generate the domain for the simulation (where the focus will be simulated)
    x = np.linspace(-D, D, M)
    z = np.linspace(-D, D, M)
    [RX, RZ] = np.meshgrid(x, z)
    RY = np.zeros(RX.shape)
    
    #create an image of the back aperture (just an incident plane wave)
    xap = np.linspace(-D, D, N)
    [XAP, YAP] = np.meshgrid(xap, xap)
    ZAP = np.zeros(XAP.shape)
    
    #evaluate the plane wave at the back aperture to produce a 2D image of the field
    EAP = pw.evaluate(XAP, YAP, ZAP)
    
    #calculate the field at the focus (coordinates are specified in RX, RY, and RZ)
    Efocus = O.focus(EAP, k0, RX, RY, RZ)
    
    #display an image of the calculated field
    plt.figure()
    plt.imshow(intensity(Efocus))
    plt.colorbar()
    
    return Efocus
    
#This sample code produces a field similar to Figure 2 in Davis et al., 2010
def example_layer():
    
    #set the material properties
    depths = [-50, -15, 15, 40]                         #specify the position (first boundary) of each layer (first boundary is where the field is defined)
    n = [1.0, 1.4+1j*0.05, 1.4, 1.0]                    #specify the refractive indices of each layer
    #create a layered sample
    layers = layersample(n, depths)
    
    #set the input light parameters
    d = np.array([0.5, 0])                             #direction of propagation of the plane wave
    
    #d = d / np.linalg.norm(d)                           #normalize this direction vector
    l0 = 5                                              #specify the wavelength in free-space         
    k0 = 2 * np.pi / l0                              #calculate the wavenumber in free-space
    E0 = [1, 1, 0]                                      #specify the E vector
    E0 = orthogonalize(E0, d)                           #make sure that both vectors are orthogonal
    #solve for the substrate field
    
    layers.solve1(d, k0, E0)
    #set the simulation domain
    N = 512
    M = 1024
    D = [-40, 80, 0, 60]
    x = np.linspace(D[2], D[3], N)
    z = np.linspace(D[0], D[1], M)
    [X, Z] = np.meshgrid(x, z)
    Y = np.zeros(X.shape)
    E = layers.evaluate(X, Y, Z)
    Er = np.real(E)
    I = intensity(E)
    plt.figure()
    plt.set_cmap("afmhot")
    matplotlib.rcParams.update({'font.size': 32})
    plt.subplot(1, 4, 1)
    plt.imshow(Er[..., 0], extent=(D[3], D[2], D[1], D[0]))
    plt.title("Ex")
    plt.subplot(1, 4, 2)
    plt.imshow(Er[..., 1], extent=(D[3], D[2], D[1], D[0]))
    plt.title("Ey")
    plt.subplot(1, 4, 3)
    plt.imshow(Er[..., 2], extent=(D[3], D[2], D[1], D[0]))
    plt.title("Ez")
    plt.subplot(1, 4, 4)
    plt.imshow(I, extent=(D[3], D[2], D[1], D[0]))
    plt.title("I")
# simulate a PSF incident on a 2-layered sample    
def example_psflayer():
    #specify the angle of accepted waves (NA if n = 1)
    sin_angle = 0.7
    
    #refractive index of the imaging medium
    n = 1.0
    
    #specify the size of the spectral domain (resolution of the front/back aperture)
    N = 31
    
    #create a new objective, specifying the NA (angle * refractive index) and resolution of the transform    
    O = objective(sin_angle, n, N)
    
    #specify the size of the spatial domain
    D = 5
    
    
    #specify the resolution of the spatial domain
    M = 200
    
    #free-space wavelength
    l0 = 1
    
    #calculate the free-space wavenumber
    k0 = 2 * np.pi /l0
    
    #specify the field incident at the aperture
    pap = np.array([1, 0, 0])                     #polarization
    sap = np.array([0, 0, 1])                  #incident angle of the plane wave (will be normalized)
    sap = sap / np.linalg.norm(sap)             #calculate the normalized direction vector for the plane wave
    
    k0ap = sap * k0                             #calculate the k-vector of the plane wave incident on the aperture
    
    #create a plane wave
    pw = planewave(k0ap, pap)
    
    #create an image of the back aperture (just an incident plane wave)
    xap = np.linspace(-D, D, N)
    [XAP, YAP] = np.meshgrid(xap, xap)
    ZAP = np.zeros(XAP.shape)
    
    #evaluate the plane wave at the back aperture to produce a 2D image of the field
    Eback = pw.evaluate(XAP, YAP, ZAP)
    Efront = O.back2front(Eback)
    
    #set the material properties
    depths = [0, -1, 0, 2]                         #specify the position (first boundary) of each layer (first boundary is where the field is defined)
    n = [1.0, 1.4+1j*0.05, 1.4, 1.0]                    #specify the refractive indices of each layer
    #create a layered sample
    L = layersample(n, depths)
    
    #generate the domain for the simulation (where the focus will be simulated)
    x = np.linspace(-D, D, M)
    z = np.linspace(-D, D, M)
    [RX, RZ] = np.meshgrid(x, z)
    RY = np.zeros(RX.shape)
    
    i = 0
    ang = np.linspace(-sin_angle, sin_angle, N)
    
    for yi in range(N):
        dy = ang[yi]
        for xi in range(N):
            dx = ang[xi]
            dxdy_sq = dx**2 + dy**2
            
            if dxdy_sq <= sin_angle**2:
                
                d = np.array([dx, dy])
                e = Efront[xi, yi, :]
                L.solve1(d, k0, e)
                
                if i == 0:
                    E = L.evaluate(RX, RY, RZ)
                    i = i + 1
                else:
                    E = E + L.evaluate(RX, RY, RZ)
            
    plt.imshow(intensity(E))
    
def example_spp():
    l = 0.647
    k0 = 2 * np.pi / l
    
    n_oil = 1.51
    n_mica = 1.56
    n_gold = 0.16049 + 1j * 3.5726
    n_water = 1.33
    
    #all units are in um
    d_oil = 100
    d_mica = 80
    d_gold = 0.05
    
    #specify the parameters for the layered sample
    zp = [0, d_oil, d_oil + d_mica, d_oil + d_mica + d_gold]
    n = [n_oil, n_mica, n_gold, n_water]
    
    #generate a layered sample
    L = layersample(n, zp)
    
    N = 10000
    min_theta = -90
    max_theta = 90
    DEG = np.linspace(min_theta, max_theta, N)
    
    I = np.zeros(N)
    
    for i in range(N):
        theta = DEG[i] * np.pi / 180
        
        #create a plane wave
        x = np.sin(theta)
        z = np.cos(theta)
        d = [x, 0]
        
        #calculate the E vector by finding a value orthogonal to d
        E0 = [-z, 0, x]
        
        #solve for the layered sample
        L.solve1(d, k0, E0)    
        
        Er = L.Pr[:, 0]
        I[i] = intensity(Er)
    plt.plot(I)