Blame view

python/coupledwave.py 11.1 KB
fdc7f32c   David Mayerich   Added a coupled w...
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
  # -*- 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])