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]

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])``````