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,7 +52,7 @@ namespace stim{ | ||
52 | // compute capability | 52 | // compute capability |
53 | int testDevices(int* dlist, unsigned n_devices, int major, int minor){ | 53 | int testDevices(int* dlist, unsigned n_devices, int major, int minor){ |
54 | int valid = 0; | 54 | int valid = 0; |
55 | - for(int d = 0; d < n_devices; d++){ | 55 | + for(unsigned d = 0; d < n_devices; d++){ |
56 | if(testDevice(dlist[d], major, minor)) | 56 | if(testDevice(dlist[d], major, minor)) |
57 | valid++; | 57 | valid++; |
58 | } | 58 | } |