Commit 057bed5a7cb679805baca6d87ba49230e854abfe
1 parent
c685653d
added optics python library
Showing
2 changed files
with
844 additions
and
1 deletions
Show diff stats
1 | +# -*- coding: utf-8 -*- | |
2 | +""" | |
3 | +Created on Fri Jun 22 16:22:17 2018 | |
4 | + | |
5 | +@author: david | |
6 | +""" | |
7 | + | |
8 | +import numpy as np | |
9 | +import scipy as sp | |
10 | +import scipy.ndimage | |
11 | +import matplotlib | |
12 | +import math | |
13 | +import matplotlib.pyplot as plt | |
14 | + | |
15 | +#adjust the components of E0 so that they are orthogonal to d | |
16 | + #if d is 2D, the z component is calculated such that ||[dx, dy, dz]||=1 | |
17 | +def orthogonalize(E, d): | |
18 | + if d.shape[-1] == 2: | |
19 | + dz = 1 - d[..., 0]**2 - d[..., 1]**2 | |
20 | + d = np.append(d, dz) | |
21 | + s = np.cross(E, d) | |
22 | + | |
23 | + dE = np.cross(d, s) | |
24 | + return dE | |
25 | + #return dE/np.linalg.norm(dE) * np.linalg.norm(E) | |
26 | + | |
27 | +def intensity(E): | |
28 | + Econj = np.conj(E) | |
29 | + I = np.sum(E*Econj, axis=-1) | |
30 | + return np.real(I) | |
31 | + | |
32 | +class planewave: | |
33 | + #implement all features of a plane wave | |
34 | + # k, E, frequency (or wavelength in a vacuum)--l | |
35 | + # try to enforce that E and k have to be orthogonal | |
36 | + | |
37 | + #initialization function that sets these parameters when a plane wave is created | |
38 | + def __init__ (self, k, E): | |
39 | + | |
40 | + self.k = np.array(k) | |
41 | + self.E = E | |
42 | + | |
43 | + if ( np.linalg.norm(k) > 1e-15 and np.linalg.norm(E) >1e-15): | |
44 | + s = np.cross(k, E) #compute an orthogonal side vector | |
45 | + s = s / np.linalg.norm(s) #normalize it | |
46 | + Edir = np.cross(s, k) #compute new E vector which is orthogonal | |
47 | + self.k = np.array(k) | |
48 | + self.E = Edir / np.linalg.norm(Edir) * np.linalg.norm(E) | |
49 | + | |
50 | + def __str__(self): | |
51 | + return str(self.k) + "\n" + str(self.E) #for verify field vectors use print command | |
52 | + | |
53 | + def kmag(self): | |
54 | + return math.sqrt(self.k[0] ** 2 + self.k[1] ** 2 + self.k[2] ** 2) | |
55 | + | |
56 | + #function that renders the plane wave given a set of coordinates | |
57 | + def evaluate(self, X, Y, Z): | |
58 | + k_dot_r = self.k[0] * X + self.k[1] * Y + self.k[2] * Z #phase term k*r | |
59 | + ex = np.exp(1j * k_dot_r) #E field equation = E0 * exp (i * (k * r)) here we simply set amplitude as 1 | |
60 | + Ef = ex[..., None] * self.E | |
61 | + return Ef | |
62 | + | |
63 | + | |
64 | +class objective: | |
65 | + | |
66 | + #initialize an objective lens | |
67 | + # sin_alpha is the sine of the angle subtended by the objective (numerical aperture in a refractive index of 1.0) | |
68 | + # n is the real refractive index of the imaging media (default air) | |
69 | + # N is the resolution of the 2D objective map (field at the input and output aperture) | |
70 | + def __init__(self, sin_alpha, n, N): | |
71 | + | |
72 | + #set member variables | |
73 | + self.sin_alpha = sin_alpha | |
74 | + self.n = n | |
75 | + | |
76 | + #calculate the transfer function from the back aperture to the front aperture | |
77 | + self.N = N | |
78 | + self._CI(N) | |
79 | + self._CR(N) | |
80 | + | |
81 | + | |
82 | + #calculates the basis vectors describing the mapping between the back and front aperture | |
83 | + def _calc_bases(self, N): | |
84 | + | |
85 | + #calculate the (sx, sy) coordinates (angular spectrum) corresponding to each position in the entrance/exit pupil | |
86 | + s = np.linspace(-self.sin_alpha, self.sin_alpha, N) | |
87 | + [SX, SY] = np.meshgrid(s, s) | |
88 | + S_SQ = SX**2 + SY**2 #calculate sx^2 + sy^2 | |
89 | + | |
90 | + BAND_M = S_SQ <= self.sin_alpha**2 #find a mask for pixels within the bandpass of the objective | |
91 | + self.BAND_M = BAND_M | |
92 | + | |
93 | + #pull out pixels that pass through the objective bandpass | |
94 | + SX = SX[BAND_M] | |
95 | + SY = SY[BAND_M] | |
96 | + S_SQ = S_SQ[BAND_M] | |
97 | + SZ = np.sqrt(1 - S_SQ) #calculate sz = sqrt(1 - sx^2 - sy^2) in normalized (-1, 1) coordinates | |
98 | + | |
99 | + #calculate the apodization functions for each polarization | |
100 | + self.FS = 1.0 / np.sqrt(SZ) | |
101 | + self.FP = self.FS | |
102 | + | |
103 | + #calculate the S and P basis vectors | |
104 | + THETA = np.arctan2(SY, SX) | |
105 | + | |
106 | + SX_SXSY = np.cos(THETA) | |
107 | + SY_SXSY = np.sin(THETA) | |
108 | + | |
109 | + v_size = (len(SX), 3) | |
110 | + VS_VSP = np.zeros(v_size) | |
111 | + VP = np.zeros(v_size) | |
112 | + VPP = np.zeros(v_size) | |
113 | + | |
114 | + VS_VSP[:, 0] = -SY_SXSY | |
115 | + VS_VSP[:, 1] = SX_SXSY | |
116 | + | |
117 | + VP[:, 0] = -SX_SXSY | |
118 | + VP[:, 1] = -SY_SXSY | |
119 | + | |
120 | + SXSY = np.sqrt(S_SQ) | |
121 | + VPPZ = np.zeros(SXSY.shape) | |
122 | + DC_M = SXSY > 0 | |
123 | + VPPZ[DC_M] = S_SQ[DC_M] / SXSY[DC_M] | |
124 | + VPP[:, 0] = -SX_SXSY * SZ | |
125 | + VPP[:, 1] = -SY_SXSY * SZ | |
126 | + VPP[:, 2] = VPPZ | |
127 | + | |
128 | + return [VS_VSP, VP, VPP] | |
129 | + | |
130 | + | |
131 | + #calculate the transfer function that maps the back aperture field to the front aperture field (Eq. 19 in Davis et al.) | |
132 | + def _CI(self, N): | |
133 | + [VS_VSP, VP, VPP] = self._calc_bases(N) | |
134 | + | |
135 | + VSP_VST = np.matmul(VS_VSP[:, :, None], VS_VSP[:, None, :]) | |
136 | + VPP_VPT = np.matmul(VPP[:, :, None], VP[:, None, :]) | |
137 | + | |
138 | + self.CI = self.FS[:, None, None] * VSP_VST + self.FP[:, None, None] * VPP_VPT | |
139 | + | |
140 | + def _CR(self, N): | |
141 | + [VS_VSP, VP, VPP] = self._calc_bases(N) | |
142 | + | |
143 | + VSP_VST = np.matmul(VS_VSP[:, :, None], VS_VSP[:, None, :]) | |
144 | + VP_VPPT = np.matmul(VP[:, :, None], VPP[:, None, :]) | |
145 | + | |
146 | + self.CR = 1.0 / self.FS[:, None, None] * VSP_VST + 1.0 / self.FP[:, None, None] * VP_VPPT | |
147 | + | |
148 | + #condenses a field defined at the back aperture of the objective | |
149 | + # E is a 3xNxN image of the field at the back aperture | |
150 | + def back2front(self, E): | |
151 | + Ein = E[self.BAND_M, :, None] | |
152 | + | |
153 | + Efront = np.zeros((self.N, self.N, 3), dtype=np.complex128) | |
154 | + Efront[self.BAND_M, :] = np.squeeze(np.matmul(self.CI, Ein)) | |
155 | + | |
156 | + return Efront | |
157 | + | |
158 | + #propagates a field from the front aperture of the objective to the back aperture | |
159 | + # E is a 3xNxN image of the field at the front aperture | |
160 | + def front2back(self, E): | |
161 | + Ein = E[self.BAND_M, :, None] | |
162 | + | |
163 | + Eback = np.zeros((self.N, self.N, 3), dtype=np.complex128) | |
164 | + Eback[self.BAND_M, :] = np.squeeze(np.matmul(self.CR, Ein)) | |
165 | + | |
166 | + return Eback | |
167 | + | |
168 | + #calculates the image at the focus given the image at the back aperture | |
169 | + # E is the field at the back aperture of the objective | |
170 | + # k0 is the wavenumber in a vacuum (2pi/lambda_0) | |
171 | + def focus(self, E, k0, RX, RY, RZ): | |
172 | + | |
173 | + #calculate the angular spectrum at the front aperture | |
174 | + Ef = self.back2front(E) | |
175 | + | |
176 | + #generate an array describing the valid components of S | |
177 | + s = np.linspace(-self.sin_alpha, self.sin_alpha, self.N) | |
178 | + [SX, SY] = np.meshgrid(s, s) | |
179 | + S = [SX[self.BAND_M], SY[self.BAND_M]] | |
180 | + SXSY_SQ = S[0]**2 + S[1]**2 | |
181 | + S.append(np.sqrt(1 - SXSY_SQ)) | |
182 | + S = np.array(S) | |
183 | + | |
184 | + | |
185 | + #calculate the k vector (accounting for both k0 and the refractive index of the immersion medium) | |
186 | + K = k0 * self.n * np.transpose(S)[:, None, :] | |
187 | + | |
188 | + #calculate the dot product of K and R | |
189 | + R = np.moveaxis(np.array([RX, RY, RZ]), 0, -1)[..., None, :, None] | |
190 | + | |
191 | + K_DOT_R = np.squeeze(np.matmul(K, R)) | |
192 | + EX = np.exp(1j * K_DOT_R)[..., None] | |
193 | + Ewaves = EX * Ef[self.BAND_M] | |
194 | + | |
195 | + Efocus = np.sum(Ewaves, -2) | |
196 | + return Efocus | |
197 | + | |
198 | +class gaussian_beam: | |
199 | + | |
200 | + #generate a beam | |
201 | + # (Ex, Ey) specifies the polarization | |
202 | + # p specifies the focal point | |
203 | + # l is the free-space wavelength | |
204 | + # sigma_p is the standard deviation of the Gaussian beam at the focal point | |
205 | + # amp is the amplitude of the DC component of the Gaussian beam | |
206 | + def __init__(self, x_polar, y_polar, lambda_0=1, amp=1, sigma_p=None, sigma_k=None): | |
207 | + self.E = [x_polar, y_polar] | |
208 | + self.l = lambda_0 | |
209 | + | |
210 | + #calculate the standard deviation of the Gaussian function in k-space | |
211 | + if(sigma_p is None and sigma_k is None): | |
212 | + self.sigma = 2 * math.pi / lambda_0 | |
213 | + elif(sigma_p is not None): | |
214 | + self.sigma = 1.0 / sigma_p | |
215 | + else: | |
216 | + self.sigma = sigma_k | |
217 | + | |
218 | + #self.s = self.sigma_p2k(sigma_p) | |
219 | + self.A = amp | |
220 | + | |
221 | + #calculates the maximum k-space frequency supported by this beam | |
222 | + def kmax(self): | |
223 | + return 2 * np.pi / self.l | |
224 | + | |
225 | + #calculate the standard deviation of the Gaussian function in k space given the std at the focal point | |
226 | + # def sigma_p2k(self, sigma): | |
227 | + | |
228 | + # this code is based on: http://www.robots.ox.ac.uk/~az/lectures/ia/lect2.pdf | |
229 | + # f(r) = 1 / (2 pi sigma^2) e^(-r^2 / (2 sigma^2)) | |
230 | + # F(kx, ky) = f( |k_par| ) = e^( -2 pi^2 |k_par|^2 sigma^2 ) | |
231 | + | |
232 | + # Basically, the Fourier transform of a normally distributed Gaussian is | |
233 | + # a non-scaled Gaussian with standard deviation 1/(2 pi sigma) | |
234 | + #return 1.0/(2 * math.pi * sigma) | |
235 | + #return 1.0 / sigma | |
236 | + | |
237 | + | |
238 | + def kspace(self, Kx, Ky): | |
239 | + #sigma = self.s * (np.pi / self.l) | |
240 | + #G = 1 / (2 * np.pi * sigma **2) * np.exp(- (0.5/(sigma**2)) * (Kx ** 2 + Ky ** 2)) | |
241 | + G = self.A * np.exp(- (Kx ** 2 + Ky ** 2) / (2 * self.sigma ** 2)) | |
242 | + | |
243 | + #calculate k-vectors outside of the band-limit | |
244 | + B = ((Kx ** 2 + Ky **2) > self.kmax() ** 2) | |
245 | + | |
246 | + #calculate Kz (required for calculating the Ez component of E) | |
247 | + #Kz = np.lib.scimath.sqrt(self.kmax() ** 2 - Kx**2 - Ky**2) | |
248 | + | |
249 | + Ex = self.E[0] * G | |
250 | + Ey = self.E[1] * G | |
251 | + #Ez = -Kx/Kz * Ex | |
252 | + Ez = np.zeros(Ex.shape) | |
253 | + | |
254 | + #set all values outside of the band limit to zero | |
255 | + Ex[B] = 0 | |
256 | + Ey[B] = 0 | |
257 | + Ez[B] = 0 | |
258 | + | |
259 | + return [Ex, Ey, Ez] | |
260 | + | |
261 | + # The PSF at the focal point is the convolution of a Bessel function and a Gaussian | |
262 | + # This function calculates the parameters of the Gaussian function in terms of a*e^{-x^2 / (2*c^2)} | |
263 | + # a (return) is the scale factor for the Gaussian | |
264 | + # c (return) is the standard deviation | |
265 | + def psf_gaussian(self): | |
266 | + # this code is based on: http://mathworld.wolfram.com/FourierTransformGaussian.html | |
267 | + #a = math.sqrt( 2 * math.pi * (self.s ** 2) ) | |
268 | + #c = 1.0 / (2 * math.pi * self.s) | |
269 | + | |
270 | + a = math.sqrt( (self.sigma ** 2) ) | |
271 | + c = 1.0 / (self.sigma) | |
272 | + | |
273 | + return [a, c] | |
274 | + | |
275 | + # The PSF at the focal point is the convolution of a Bessel function and a Gaussian | |
276 | + # This function calculates the parameters of the bessel function in terms of a J_1(pi * a * x) / x | |
277 | + def psf_bessel(self): | |
278 | + return self.kmax() / math.pi | |
279 | + | |
280 | + # calculate the beam point spread function at the focal plane | |
281 | + # D is the image extent for the square: (-D, D, -D, D) | |
282 | + # N is the number of spatial samples along one dimension of the image (NxN) | |
283 | + # order is the mathematical order of the calculation (in terms of the resolution of the FFT) | |
284 | + def psf(self, D, N): | |
285 | + | |
286 | + [a, c] = self.psf_gaussian() | |
287 | + alpha = self.psf_bessel() | |
288 | + | |
289 | + #calculate the sample positions along the x-axis | |
290 | + x = np.linspace(-D, D, N) | |
291 | + dx = x[1] - x[0] | |
292 | + | |
293 | + [X, Y] = np.meshgrid(x, x) | |
294 | + R = np.sqrt(X**2 + Y**2) | |
295 | + | |
296 | + bessel = alpha * sp.special.j1(math.pi * alpha * R) / R | |
297 | + | |
298 | + return a * sp.ndimage.filters.gaussian_filter(bessel, c / dx) | |
299 | + | |
300 | + | |
301 | + #return the k-space transform of the beam at p as an NxN array | |
302 | + def kspace_image(self, N): | |
303 | + #calculate the band limit [-2*pi/lambda, 2*pi/lambda] | |
304 | + kx = np.linspace(-self.kmax(), self.kmax(), N) | |
305 | + ky = kx | |
306 | + | |
307 | + #generate a meshgrid for the field | |
308 | + [Kx, Ky] = np.meshgrid(kx, ky) | |
309 | + | |
310 | + #generate a Gaussian with a standard deviation of sigma * pi/lambda | |
311 | + return self.kspace(Kx, Ky) | |
312 | + | |
313 | + # calculate the practical band limit of the beam such that | |
314 | + # all frequencies higher than the returned values are less than epsilon | |
315 | + def practical_band_limit(self, epsilon = 0.001): | |
316 | + | |
317 | + return min(self.sigma * math.sqrt(math.log(1.0/epsilon)), self.kmax()) | |
318 | + | |
319 | + | |
320 | + #return a plane wave decomposition of the beam using N plane waves | |
321 | + def decompose(self, N): | |
322 | + PW = [] #create an empty list of plane waves | |
323 | + | |
324 | + # sample a cylinder and project those samples onto a unit sphere | |
325 | + theta = np.random.uniform(0, 2*np.pi, N) | |
326 | + | |
327 | + # get the practical band limit to minimize sampling | |
328 | + band_limit = self.practical_band_limit() / self.kmax() | |
329 | + phi = np.arccos(np.random.uniform(1 - band_limit, 1, N)) | |
330 | + | |
331 | + kx = -self.kmax() * np.sin(phi) * np.cos(theta) | |
332 | + ky = -self.kmax() * np.sin(phi) * np.sin(theta) | |
333 | + | |
334 | + mc_weight = 2 * math.pi / N | |
335 | + E0 = self.kspace(kx, ky) | |
336 | + kz = np.sqrt(self.kmax() ** 2 - kx**2 - ky**2) | |
337 | + | |
338 | + for i in range(0, N): | |
339 | + k = [kx[i], ky[i], kz[i]] | |
340 | + E = [E0[0][i] * mc_weight, E0[1][i] * mc_weight, E0[2][i] * mc_weight] | |
341 | + w = planewave(k, E) | |
342 | + PW.append(w) | |
343 | + | |
344 | + return PW | |
345 | + | |
346 | +class layersample: | |
347 | + | |
348 | + #generate a layered sample based on a list of layer refractive indices | |
349 | + # layer_ri is a list of complex refractive indices for each layer | |
350 | + # z is the list of boundary positions along z | |
351 | + def __init__(self, layer_ri, z): | |
352 | + self.n = np.array(layer_ri).astype(np.complex128) | |
353 | + self.z = np.array(z) | |
354 | + | |
355 | + #calculate the index of the field component associated with a layer | |
356 | + # l is the layer index [0, L) | |
357 | + # c is the component (x=0, y=1, z=2) | |
358 | + # d is the direction (0 = transmission, 1 = reflection) | |
359 | + def i(self, l, c, d): | |
360 | + i = l * 6 + d * 3 + c - 3 | |
361 | + return i | |
362 | + | |
363 | + #create a matrix for a single plane wave specified by k and E | |
364 | + # d = [dx, dy] are the x and y coordinates of the normalized direction of propagation | |
365 | + # k0 is the free space wave number (2 pi / lambda0) | |
366 | + # E is the electric field vector | |
367 | + | |
368 | + def solve1(self, d, k0, E): | |
369 | + | |
370 | + #s is the plane wave direction scaled by the refractive index | |
371 | + s = np.array(d) * self.n[0] | |
372 | + | |
373 | + #allocate space for the matrix | |
374 | + L = len(self.n) | |
375 | + M = np.zeros((6*(L-1), 6*(L-1)), dtype=np.complex128) | |
376 | + | |
377 | + #allocate space for the RHS vector | |
378 | + b = np.zeros(6*(L-1), dtype=np.complex128) | |
379 | + | |
380 | + #initialize a counter for the equation number | |
381 | + ei = 0 | |
382 | + | |
383 | + #calculate the sz component for each layer | |
384 | + self.sz = np.zeros(L, dtype=np.complex128) | |
385 | + for l in range(L): | |
386 | + self.sz[l] = np.sqrt(self.n[l]**2 - s[0]**2 - s[1]**2) | |
387 | + | |
388 | + #set constraints based on Gauss' law | |
389 | + for l in range(0, L): | |
390 | + #sz = np.sqrt(self.n[l]**2 - s[0]**2 - s[1]**2) | |
391 | + | |
392 | + #set the upward components for each layer | |
393 | + # note that layer L-1 does not have a downward component | |
394 | + # David I, Equation 7 | |
395 | + if l != L-1: | |
396 | + M[ei, self.i(l, 0, 1)] = s[0] | |
397 | + M[ei, self.i(l, 1, 1)] = s[1] | |
398 | + M[ei, self.i(l, 2, 1)] = -self.sz[l] | |
399 | + ei = ei+1 | |
400 | + | |
401 | + #set the downward components for each layer | |
402 | + # note that layer 0 does not have a downward component | |
403 | + # Davis I, Equation 6 | |
404 | + if l != 0: | |
405 | + M[ei, self.i(l, 0, 0)] = s[0] | |
406 | + M[ei, self.i(l, 1, 0)] = s[1] | |
407 | + M[ei, self.i(l, 2, 0)] = self.sz[l] | |
408 | + ei = ei+1 | |
409 | + | |
410 | + #enforce a continuous field across boundaries | |
411 | + for l in range(1, L): | |
412 | + sz0 = self.sz[l-1] | |
413 | + sz1 = self.sz[l] | |
414 | + A = np.exp(1j * k0 * sz0 * (self.z[l] - self.z[l-1])) | |
415 | + if l < L - 1: | |
416 | + dl = self.z[l] - self.z[l+1] | |
417 | + arg = -1j * k0 * sz1 * dl | |
418 | + B = np.exp(arg) | |
419 | + | |
420 | + | |
421 | + #if this is the second layer, use the simplified equations that account for the incident field | |
422 | + if l == 1: | |
423 | + M[ei, self.i(0, 0, 1)] = 1 | |
424 | + M[ei, self.i(1, 0, 0)] = -1 | |
425 | + if L > 2: | |
426 | + #print(-B, M[ei, self.i(1, 0, 1)]) | |
427 | + M[ei, self.i(1, 0, 1)] = -B | |
428 | + | |
429 | + b[ei] = -A * E[0] | |
430 | + ei = ei + 1 | |
431 | + | |
432 | + M[ei, self.i(0, 1, 1)] = 1 | |
433 | + M[ei, self.i(1, 1, 0)] = -1 | |
434 | + if L > 2: | |
435 | + M[ei, self.i(1, 1, 1)] = -B | |
436 | + b[ei] = -A * E[1] | |
437 | + ei = ei + 1 | |
438 | + | |
439 | + M[ei, self.i(0, 2, 1)] = s[1] | |
440 | + M[ei, self.i(0, 1, 1)] = sz0 | |
441 | + M[ei, self.i(1, 2, 0)] = -s[1] | |
442 | + M[ei, self.i(1, 1, 0)] = sz1 | |
443 | + if L > 2: | |
444 | + M[ei, self.i(1, 2, 1)] = -B*s[1] | |
445 | + M[ei, self.i(1, 1, 1)] = -B*sz1 | |
446 | + b[ei] = A * sz0 * E[1] - A * s[1]*E[2] | |
447 | + ei = ei + 1 | |
448 | + | |
449 | + M[ei, self.i(0, 0, 1)] = -sz0 | |
450 | + M[ei, self.i(0, 2, 1)] = -s[0] | |
451 | + M[ei, self.i(1, 0, 0)] = -sz1 | |
452 | + M[ei, self.i(1, 2, 0)] = s[0] | |
453 | + if L > 2: | |
454 | + M[ei, self.i(1, 0, 1)] = B*sz1 | |
455 | + M[ei, self.i(1, 2, 1)] = B*s[0] | |
456 | + b[ei] = A * s[0] * E[2] - A * sz0 * E[0] | |
457 | + ei = ei + 1 | |
458 | + | |
459 | + #if this is the last layer, use the simplified equations that exclude reflections from the last layer | |
460 | + elif l == L-1: | |
461 | + M[ei, self.i(l-1, 0, 0)] = A | |
462 | + M[ei, self.i(l-1, 0, 1)] = 1 | |
463 | + M[ei, self.i(l, 0, 0)] = -1 | |
464 | + ei = ei + 1 | |
465 | + | |
466 | + M[ei, self.i(l-1, 1, 0)] = A | |
467 | + M[ei, self.i(l-1, 1, 1)] = 1 | |
468 | + M[ei, self.i(l, 1, 0)] = -1 | |
469 | + ei = ei + 1 | |
470 | + | |
471 | + M[ei, self.i(l-1, 2, 0)] = A*s[1] | |
472 | + M[ei, self.i(l-1, 1, 0)] = -A*sz0 | |
473 | + M[ei, self.i(l-1, 2, 1)] = s[1] | |
474 | + M[ei, self.i(l-1, 1, 1)] = sz0 | |
475 | + M[ei, self.i(l, 2, 0)] = -s[1] | |
476 | + M[ei, self.i(l, 1, 0)] = sz1 | |
477 | + ei = ei + 1 | |
478 | + | |
479 | + M[ei, self.i(l-1, 0, 0)] = A*sz0 | |
480 | + M[ei, self.i(l-1, 2, 0)] = -A*s[0] | |
481 | + M[ei, self.i(l-1, 0, 1)] = -sz0 | |
482 | + M[ei, self.i(l-1, 2, 1)] = -s[0] | |
483 | + M[ei, self.i(l, 0, 0)] = -sz1 | |
484 | + M[ei, self.i(l, 2, 0)] = s[0] | |
485 | + ei = ei + 1 | |
486 | + #otherwise use the full set of boundary conditions | |
487 | + else: | |
488 | + M[ei, self.i(l-1, 0, 0)] = A | |
489 | + M[ei, self.i(l-1, 0, 1)] = 1 | |
490 | + M[ei, self.i(l, 0, 0)] = -1 | |
491 | + M[ei, self.i(l, 0, 1)] = -B | |
492 | + ei = ei + 1 | |
493 | + | |
494 | + M[ei, self.i(l-1, 1, 0)] = A | |
495 | + M[ei, self.i(l-1, 1, 1)] = 1 | |
496 | + M[ei, self.i(l, 1, 0)] = -1 | |
497 | + M[ei, self.i(l, 1, 1)] = -B | |
498 | + ei = ei + 1 | |
499 | + | |
500 | + M[ei, self.i(l-1, 2, 0)] = A*s[1] | |
501 | + M[ei, self.i(l-1, 1, 0)] = -A*sz0 | |
502 | + M[ei, self.i(l-1, 2, 1)] = s[1] | |
503 | + M[ei, self.i(l-1, 1, 1)] = sz0 | |
504 | + M[ei, self.i(l, 2, 0)] = -s[1] | |
505 | + M[ei, self.i(l, 1, 0)] = sz1 | |
506 | + M[ei, self.i(l, 2, 1)] = -B*s[1] | |
507 | + M[ei, self.i(l, 1, 1)] = -B*sz1 | |
508 | + ei = ei + 1 | |
509 | + | |
510 | + M[ei, self.i(l-1, 0, 0)] = A*sz0 | |
511 | + M[ei, self.i(l-1, 2, 0)] = -A*s[0] | |
512 | + M[ei, self.i(l-1, 0, 1)] = -sz0 | |
513 | + M[ei, self.i(l-1, 2, 1)] = -s[0] | |
514 | + M[ei, self.i(l, 0, 0)] = -sz1 | |
515 | + M[ei, self.i(l, 2, 0)] = s[0] | |
516 | + M[ei, self.i(l, 0, 1)] = B*sz1 | |
517 | + M[ei, self.i(l, 2, 1)] = B*s[0] | |
518 | + ei = ei + 1 | |
519 | + | |
520 | + #store the matrix and RHS vector (for debugging) | |
521 | + self.M = M | |
522 | + self.b = b | |
523 | + | |
524 | + #evaluate the linear system | |
525 | + P = np.linalg.solve(M, b) | |
526 | + | |
527 | + #save the results (also for debugging) | |
528 | + self.P = P | |
529 | + | |
530 | + #store the coefficients for each layer | |
531 | + self.Pt = np.zeros((3, L), np.complex128) | |
532 | + self.Pr = np.zeros((3, L), np.complex128) | |
533 | + for l in range(L): | |
534 | + if l == 0: | |
535 | + self.Pt[:, 0] = [E[0], E[1], E[2]] | |
536 | + else: | |
537 | + px = P[self.i(l, 0, 0)] | |
538 | + py = P[self.i(l, 1, 0)] | |
539 | + pz = P[self.i(l, 2, 0)] | |
540 | + self.Pt[:, l] = [px, py, pz] | |
541 | + | |
542 | + if l == L-1: | |
543 | + self.Pr[:, L-1] = [0, 0, 0] | |
544 | + else: | |
545 | + px = P[self.i(l, 0, 1)] | |
546 | + py = P[self.i(l, 1, 1)] | |
547 | + pz = P[self.i(l, 2, 1)] | |
548 | + self.Pr[:, l] = [px, py, pz] | |
549 | + | |
550 | + #store values required for evaluation | |
551 | + #store k | |
552 | + self.k = k0 | |
553 | + | |
554 | + #store sx and sy | |
555 | + self.s = np.array([s[0], s[1]]) | |
556 | + | |
557 | + self.solved = True | |
558 | + | |
559 | + #evaluate a solved homogeneous substrate | |
560 | + def evaluate(self, X, Y, Z): | |
561 | + | |
562 | + if not self.solved: | |
563 | + print("ERROR: the layered substrate hasn't been solved") | |
564 | + return | |
565 | + | |
566 | + #this code is a bit cumbersome and could probably be optimized | |
567 | + # Basically, it vectorizes everything by creating an image | |
568 | + # such that the value at each pixel is the corresponding layer | |
569 | + # that the pixel resides in. That index is then used to calculate | |
570 | + # the field within the layer | |
571 | + | |
572 | + #allocate space for layer indices | |
573 | + LI = np.zeros(Z.shape, dtype=np.int) | |
574 | + | |
575 | + #find the layer index for each sample point | |
576 | + L = len(self.z) | |
577 | + LI[Z < self.z[0]] = 0 | |
578 | + for l in range(L-1): | |
579 | + idx = np.logical_and(Z > self.z[l], Z <= self.z[l+1]) | |
580 | + LI[idx] = l | |
581 | + LI[Z > self.z[-1]] = L - 1 | |
582 | + | |
583 | + #calculate the appropriate phase shift for the wave transmitted through the layer | |
584 | + Ph_t = np.exp(1j * self.k * self.sz[LI] * (Z - self.z[LI])) | |
585 | + | |
586 | + #calculate the appropriate phase shift for the wave reflected off of the layer boundary | |
587 | + LIp = LI + 1 | |
588 | + LIp[LIp >= L] = 0 | |
589 | + Ph_r = np.exp(-1j * self.k * self.sz[LI] * (Z - self.z[LIp])) | |
590 | +# print(Ph_r) | |
591 | + Ph_r[LI >= L-1] = 0 | |
592 | + | |
593 | + #calculate the phase shift based on the X and Y positions | |
594 | + Ph_xy = np.exp(1j * self.k * (self.s[0] * X + self.s[1] * Y)) | |
595 | + | |
596 | + #apply the phase shifts | |
597 | + Et = self.Pt[:, LI] * Ph_t[:, :] | |
598 | + Er = self.Pr[:, LI] * Ph_r[:, :] | |
599 | + | |
600 | + #add everything together coherently | |
601 | + E = (Et + Er) * Ph_xy[:, :] | |
602 | + | |
603 | + #return the electric field | |
604 | + return np.moveaxis(E, 0, -1) | |
605 | + | |
606 | + | |
607 | +def example_psf(): | |
608 | + | |
609 | + #specify the angle of accepted waves (NA if n = 1) | |
610 | + sin_angle = 0.8 | |
611 | + | |
612 | + #refractive index of the imaging medium | |
613 | + n = 1.0 | |
614 | + | |
615 | + #specify the size of the spectral domain (resolution of the front/back aperture) | |
616 | + N = 100 | |
617 | + | |
618 | + #create a new objective, specifying the NA (angle * refractive index) and resolution of the transform | |
619 | + O = objective(sin_angle, n, N) | |
620 | + | |
621 | + #specify the size of the spatial domain | |
622 | + D = 10 | |
623 | + | |
624 | + | |
625 | + #specify the resolution of the spatial domain | |
626 | + M = 100 | |
627 | + | |
628 | + #free-space wavelength | |
629 | + l0 = 1 | |
630 | + | |
631 | + #calculate the free-space wavenumber | |
632 | + k0 = 2 * np.pi * l0 | |
633 | + | |
634 | + #specify the field incident at the aperture | |
635 | + pap = np.array([1, 0, 0]) #polarization | |
636 | + sap = np.array([0, 0, 1]) #incident angle of the plane wave (will be normalized) | |
637 | + sap = sap / np.linalg.norm(sap) #calculate the normalized direction vector for the plane wave | |
638 | + | |
639 | + k0ap = sap * k0 #calculate the k-vector of the plane wave incident on the aperture | |
640 | + | |
641 | + #create a plane wave | |
642 | + pw = planewave(k0ap, pap) | |
643 | + | |
644 | + #generate the domain for the simulation (where the focus will be simulated) | |
645 | + x = np.linspace(-D, D, M) | |
646 | + z = np.linspace(-D, D, M) | |
647 | + [RX, RZ] = np.meshgrid(x, z) | |
648 | + RY = np.zeros(RX.shape) | |
649 | + | |
650 | + #create an image of the back aperture (just an incident plane wave) | |
651 | + xap = np.linspace(-D, D, N) | |
652 | + [XAP, YAP] = np.meshgrid(xap, xap) | |
653 | + ZAP = np.zeros(XAP.shape) | |
654 | + | |
655 | + #evaluate the plane wave at the back aperture to produce a 2D image of the field | |
656 | + EAP = pw.evaluate(XAP, YAP, ZAP) | |
657 | + | |
658 | + #calculate the field at the focus (coordinates are specified in RX, RY, and RZ) | |
659 | + Efocus = O.focus(EAP, k0, RX, RY, RZ) | |
660 | + | |
661 | + #display an image of the calculated field | |
662 | + plt.figure() | |
663 | + plt.imshow(intensity(Efocus)) | |
664 | + plt.colorbar() | |
665 | + | |
666 | + return Efocus | |
667 | + | |
668 | +#This sample code produces a field similar to Figure 2 in Davis et al., 2010 | |
669 | +def example_layer(): | |
670 | + | |
671 | + #set the material properties | |
672 | + depths = [-50, -15, 15, 40] #specify the position (first boundary) of each layer (first boundary is where the field is defined) | |
673 | + n = [1.0, 1.4+1j*0.05, 1.4, 1.0] #specify the refractive indices of each layer | |
674 | + | |
675 | + #create a layered sample | |
676 | + layers = layersample(n, depths) | |
677 | + | |
678 | + #set the input light parameters | |
679 | + d = np.array([0.5, 0]) #direction of propagation of the plane wave | |
680 | + | |
681 | + #d = d / np.linalg.norm(d) #normalize this direction vector | |
682 | + l0 = 5 #specify the wavelength in free-space | |
683 | + k0 = 2 * np.pi / l0 #calculate the wavenumber in free-space | |
684 | + E0 = [1, 1, 0] #specify the E vector | |
685 | + E0 = orthogonalize(E0, d) #make sure that both vectors are orthogonal | |
686 | + #solve for the substrate field | |
687 | + | |
688 | + layers.solve1(d, k0, E0) | |
689 | + #set the simulation domain | |
690 | + N = 512 | |
691 | + M = 1024 | |
692 | + D = [-40, 80, 0, 60] | |
693 | + x = np.linspace(D[2], D[3], N) | |
694 | + z = np.linspace(D[0], D[1], M) | |
695 | + [X, Z] = np.meshgrid(x, z) | |
696 | + Y = np.zeros(X.shape) | |
697 | + E = layers.evaluate(X, Y, Z) | |
698 | + Er = np.real(E) | |
699 | + I = intensity(E) | |
700 | + plt.figure() | |
701 | + plt.set_cmap("afmhot") | |
702 | + matplotlib.rcParams.update({'font.size': 32}) | |
703 | + plt.subplot(1, 4, 1) | |
704 | + plt.imshow(Er[..., 0], extent=(D[3], D[2], D[1], D[0])) | |
705 | + plt.title("Ex") | |
706 | + plt.subplot(1, 4, 2) | |
707 | + plt.imshow(Er[..., 1], extent=(D[3], D[2], D[1], D[0])) | |
708 | + plt.title("Ey") | |
709 | + plt.subplot(1, 4, 3) | |
710 | + plt.imshow(Er[..., 2], extent=(D[3], D[2], D[1], D[0])) | |
711 | + plt.title("Ez") | |
712 | + plt.subplot(1, 4, 4) | |
713 | + plt.imshow(I, extent=(D[3], D[2], D[1], D[0])) | |
714 | + plt.title("I") | |
715 | + | |
716 | +def example_psflayer(): | |
717 | + #specify the angle of accepted waves (NA if n = 1) | |
718 | + sin_angle = 0.7 | |
719 | + | |
720 | + #refractive index of the imaging medium | |
721 | + n = 1.0 | |
722 | + | |
723 | + #specify the size of the spectral domain (resolution of the front/back aperture) | |
724 | + N = 31 | |
725 | + | |
726 | + #create a new objective, specifying the NA (angle * refractive index) and resolution of the transform | |
727 | + O = objective(sin_angle, n, N) | |
728 | + | |
729 | + #specify the size of the spatial domain | |
730 | + D = 5 | |
731 | + | |
732 | + | |
733 | + #specify the resolution of the spatial domain | |
734 | + M = 200 | |
735 | + | |
736 | + #free-space wavelength | |
737 | + l0 = 1 | |
738 | + | |
739 | + #calculate the free-space wavenumber | |
740 | + k0 = 2 * np.pi /l0 | |
741 | + | |
742 | + #specify the field incident at the aperture | |
743 | + pap = np.array([1, 0, 0]) #polarization | |
744 | + sap = np.array([0, 0, 1]) #incident angle of the plane wave (will be normalized) | |
745 | + sap = sap / np.linalg.norm(sap) #calculate the normalized direction vector for the plane wave | |
746 | + | |
747 | + k0ap = sap * k0 #calculate the k-vector of the plane wave incident on the aperture | |
748 | + | |
749 | + #create a plane wave | |
750 | + pw = planewave(k0ap, pap) | |
751 | + | |
752 | + #create an image of the back aperture (just an incident plane wave) | |
753 | + xap = np.linspace(-D, D, N) | |
754 | + [XAP, YAP] = np.meshgrid(xap, xap) | |
755 | + ZAP = np.zeros(XAP.shape) | |
756 | + | |
757 | + #evaluate the plane wave at the back aperture to produce a 2D image of the field | |
758 | + Eback = pw.evaluate(XAP, YAP, ZAP) | |
759 | + Efront = O.back2front(Eback) | |
760 | + | |
761 | + #set the material properties | |
762 | + depths = [0, -1, 0, 2] #specify the position (first boundary) of each layer (first boundary is where the field is defined) | |
763 | + n = [1.0, 1.4+1j*0.05, 1.4, 1.0] #specify the refractive indices of each layer | |
764 | + | |
765 | + #create a layered sample | |
766 | + L = layersample(n, depths) | |
767 | + | |
768 | + #generate the domain for the simulation (where the focus will be simulated) | |
769 | + x = np.linspace(-D, D, M) | |
770 | + z = np.linspace(-D, D, M) | |
771 | + [RX, RZ] = np.meshgrid(x, z) | |
772 | + RY = np.zeros(RX.shape) | |
773 | + | |
774 | + i = 0 | |
775 | + ang = np.linspace(-sin_angle, sin_angle, N) | |
776 | + | |
777 | + for yi in range(N): | |
778 | + dy = ang[yi] | |
779 | + for xi in range(N): | |
780 | + dx = ang[xi] | |
781 | + dxdy_sq = dx**2 + dy**2 | |
782 | + | |
783 | + if dxdy_sq <= sin_angle**2: | |
784 | + | |
785 | + d = np.array([dx, dy]) | |
786 | + e = Efront[xi, yi, :] | |
787 | + L.solve1(d, k0, e) | |
788 | + | |
789 | + if i == 0: | |
790 | + E = L.evaluate(RX, RY, RZ) | |
791 | + i = i + 1 | |
792 | + else: | |
793 | + E = E + L.evaluate(RX, RY, RZ) | |
794 | + | |
795 | + plt.imshow(intensity(E)) | |
796 | + | |
797 | +def example_spp(): | |
798 | + l = 0.647 | |
799 | + k0 = 2 * np.pi / l | |
800 | + | |
801 | + n_oil = 1.51 | |
802 | + n_mica = 1.56 | |
803 | + n_gold = 0.16049 + 1j * 3.5726 | |
804 | + n_water = 1.33 | |
805 | + | |
806 | + #all units are in um | |
807 | + d_oil = 100 | |
808 | + d_mica = 80 | |
809 | + d_gold = 0.05 | |
810 | + | |
811 | + #specify the parameters for the layered sample | |
812 | + zp = [0, d_oil, d_oil + d_mica, d_oil + d_mica + d_gold] | |
813 | + n = [n_oil, n_mica, n_gold, n_water] | |
814 | + | |
815 | + #generate a layered sample | |
816 | + L = layersample(n, zp) | |
817 | + | |
818 | + N = 10000 | |
819 | + min_theta = -90 | |
820 | + max_theta = 90 | |
821 | + DEG = np.linspace(min_theta, max_theta, N) | |
822 | + | |
823 | + I = np.zeros(N) | |
824 | + | |
825 | + for i in range(N): | |
826 | + theta = DEG[i] * np.pi / 180 | |
827 | + | |
828 | + #create a plane wave | |
829 | + x = np.sin(theta) | |
830 | + z = np.cos(theta) | |
831 | + d = [x, 0] | |
832 | + | |
833 | + #calculate the E vector by finding a value orthogonal to d | |
834 | + E0 = [-z, 0, x] | |
835 | + | |
836 | + #solve for the layered sample | |
837 | + L.solve1(d, k0, E0) | |
838 | + | |
839 | + Er = L.Pr[:, 0] | |
840 | + I[i] = intensity(Er) | |
841 | + plt.plot(I) | |
842 | + | |
843 | + | ... | ... |
stim/cuda/cudatools/devices.h
... | ... | @@ -52,7 +52,7 @@ namespace stim{ |
52 | 52 | // compute capability |
53 | 53 | int testDevices(int* dlist, unsigned n_devices, int major, int minor){ |
54 | 54 | int valid = 0; |
55 | - for(int d = 0; d < n_devices; d++){ | |
55 | + for(unsigned d = 0; d < n_devices; d++){ | |
56 | 56 | if(testDevice(dlist[d], major, minor)) |
57 | 57 | valid++; |
58 | 58 | } | ... | ... |