coupledwave.py 11.1 KB
# -*- coding: utf-8 -*-
"""
Created on Thu Sep  6 22:14:36 2018

@author: david
"""

import numpy as np
import matplotlib.pyplot as plt

#evaluate a plane wave for all points in R = [Rx, Ry, Rz]
#   E0 [Ex, Ey, Ez] is the electric field vector
#   k is the k vector
#   n is the refractive index of the material
#   R is a list of coordinates in x, y, and z
def planewave2field(E0, k, n, R):
    knr = n * (k[0] * R[0] + k[1] * R[1] + k[2] + R[2])
    exp_2pknr = np.exp(1j * 2 * np.pi * knr)
    Ex = E0[0] * exp_2pknr
    Ey = E0[1] * exp_2pknr
    Ez = E0[2] * exp_2pknr
    
    return [Ex, Ey, Ez]

def orthogonalize(E, k):
    s = np.cross(E, k)
    return np.cross(k, s)
    
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 idx(self, l, c, d):
        i = l * 6 + d * 3 + c - 3
        print(i)
        return  i
    
    #create a matrix for a single plane wave specified by k and E
    def solve(self, k, E):
        
        #calculate the wavenumber
        kn = np.linalg.norm(k)
        
        #s is the plane wave direction scaled by the refractive index
        s = k / kn * self.n[0]

        #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.idx(l, 0, 1)] = s[0]
                M[ei, self.idx(l, 1, 1)] = s[1]
                M[ei, self.idx(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.idx(l, 0, 0)] = s[0]
                M[ei, self.idx(l, 1, 0)] = s[1]
                M[ei, self.idx(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 * kn * sz0 * (self.z[l] - self.z[l-1]))
            if l < L - 1:
                B = np.exp(-1j * kn * sz1 * (self.z[l] - self.z[l+1]))
            
            
            #if this is the second layer, use the simplified equations that account for the incident field
            if l == 1:
                M[ei, self.idx(0, 0, 1)] = 1
                M[ei, self.idx(1, 0, 0)] = -1
                if L > 2:
                    M[ei, self.idx(1, 0, 1)] = -B
                    
                b[ei] = -A * E[0]
                ei = ei + 1
                
                M[ei, self.idx(0, 1, 1)] = 1
                M[ei, self.idx(1, 1, 0)] = -1
                if L > 2:
                    M[ei, self.idx(1, 1, 1)] = -B
                b[ei] = -A * E[1]
                ei = ei + 1
                
                M[ei, self.idx(0, 2, 1)] = s[1]
                M[ei, self.idx(0, 1, 1)] = sz0
                M[ei, self.idx(1, 2, 0)] = -s[1]
                M[ei, self.idx(1, 1, 0)] = sz1
                if L > 2:
                    M[ei, self.idx(1, 2, 1)] = -B*s[1]
                    M[ei, self.idx(1, 1, 1)] = -B*sz1
                b[ei] = A * sz0 * E[1] - A * s[1]*E[2]
                ei = ei + 1
                
                M[ei, self.idx(0, 0, 1)] = -sz0
                M[ei, self.idx(0, 2, 1)] = -s[0]
                M[ei, self.idx(1, 0, 0)] = -sz1
                M[ei, self.idx(1, 2, 0)] = s[0]
                if L > 2:
                    M[ei, self.idx(1, 0, 1)] = B*sz1
                    M[ei, self.idx(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.idx(l-1, 0, 0)] = A
                M[ei, self.idx(l-1, 0, 1)] = 1
                M[ei, self.idx(l, 0, 0)] = -1
                ei = ei + 1
                
                M[ei, self.idx(l-1, 1, 0)] = A
                M[ei, self.idx(l-1, 1, 1)] = 1
                M[ei, self.idx(l, 1, 0)] = -1
                ei = ei + 1
                
                M[ei, self.idx(l-1, 2, 0)] = A*s[1]
                M[ei, self.idx(l-1, 1, 0)] = -A*sz0
                M[ei, self.idx(l-1, 2, 1)] = s[1]
                M[ei, self.idx(l-1, 1, 1)] = sz0
                M[ei, self.idx(l, 2, 0)] = -s[1]
                M[ei, self.idx(l, 1, 0)] = sz1
                ei = ei + 1
                
                M[ei, self.idx(l-1, 0, 0)] = A*sz0
                M[ei, self.idx(l-1, 2, 0)] = -A*s[0]
                M[ei, self.idx(l-1, 0, 1)] = -sz0
                M[ei, self.idx(l-1, 2, 1)] = -s[0]
                M[ei, self.idx(l, 0, 0)] = -sz1
                M[ei, self.idx(l, 2, 0)] = s[0]
                ei = ei + 1
            #otherwise use the full set of boundary conditions
            else:
                M[ei, self.idx(l-1, 0, 0)] = A
                M[ei, self.idx(l-1, 0, 1)] = 1
                M[ei, self.idx(l, 0, 0)] = -1
                M[ei, self.idx(l, 0, 1)] = -B
                ei = ei + 1
                
                M[ei, self.idx(l-1, 1, 0)] = A
                M[ei, self.idx(l-1, 1, 1)] = 1
                M[ei, self.idx(l, 1, 0)] = -1
                M[ei, self.idx(l, 1, 1)] = -B
                ei = ei + 1
                
                M[ei, self.idx(l-1, 2, 0)] = A*s[1]
                M[ei, self.idx(l-1, 1, 0)] = -A*sz0
                M[ei, self.idx(l-1, 2, 1)] = s[1]
                M[ei, self.idx(l-1, 1, 1)] = sz0
                M[ei, self.idx(l, 2, 0)] = -s[1]
                M[ei, self.idx(l, 1, 0)] = sz1
                M[ei, self.idx(l, 2, 1)] = -B*s[1]
                M[ei, self.idx(l, 1, 1)] = -B*sz1
                ei = ei + 1
                
                M[ei, self.idx(l-1, 0, 0)] = A*sz0
                M[ei, self.idx(l-1, 2, 0)] = -A*s[0]
                M[ei, self.idx(l-1, 0, 1)] = -sz0
                M[ei, self.idx(l-1, 2, 1)] = -s[0]
                M[ei, self.idx(l, 0, 0)] = -sz1
                M[ei, self.idx(l, 2, 0)] = s[0]
                M[ei, self.idx(l, 0, 1)] = B*sz1
                M[ei, self.idx(l, 2, 1)] = B*s[0]
                ei = ei + 1
                
        #store the matrix and RHS vector (for debugging)        
        self.M = M
        self.b = b
        
        #evaluate the linear system
        P = np.linalg.solve(M, b)
        
        #save the results (also for debugging)
        self.P = P
        
        #store the coefficients for each layer
        self.Ptrans = []
        self.Prefl = []
        for l in range(L):
            if l == 0:
                self.Ptrans.append((E[0], E[1], E[2]))
            else:
                px = P[self.idx(l, 0, 0)]
                py = P[self.idx(l, 1, 0)]
                pz = P[self.idx(l, 2, 0)]
                self.Ptrans.append((px, py, pz))
                
            if l == L-1:
                self.Prefl.append((0, 0, 0))
            else:
                px = P[self.idx(l, 0, 1)]
                py = P[self.idx(l, 1, 1)]
                pz = P[self.idx(l, 2, 1)]
                self.Prefl.append((px, py, pz))
        
        #this converts the coefficients into a NumPy array for convenience later on
        self.Ptrans = np.array(self.Ptrans)
        self.Prefl = np.array(self.Prefl)

        #store values required for evaluation
        #store k
        self.k = kn
        
        #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]))
        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.Ptrans[LI] * Ph_t[:, :, None]
        Er = self.Prefl[LI] * Ph_r[:, :, None]
        
        #add everything together coherently
        E = (Et + Er) * Ph_xy[:, :, None]
        
        #return the electric field
        return E


#This sample code produces a field similar to Figure 2 in Davis et al., 2010
#set the material properties
depths = [-50, -15, 15, 40]
n = [1.0, 1.4+1j*0.05, 1.4, 1.0]

#create a layered sample
layers = layersample(n, depths)

#set the input light parameters
d = np.array([0.2, 0, 1])
d = d / np.linalg.norm(d)
l = 5.0
k = 2 * np.pi / l * d
E0 = [0, 1, 0]
E0 = orthogonalize(E0, d)

#solve for the substrate field
layers.solve(k, E0)

#set the simulation domain
N = 512
M = 1024
D = [-50, 80]

x = np.linspace(D[0], D[1], 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 = Er[..., 0] ** 2 + Er[..., 1] **2 + Er[..., 2] ** 2
plt.figure()
plt.set_cmap("afmhot")
plt.imshow(Er[..., 1])