Commit 057bed5a7cb679805baca6d87ba49230e854abfe

Authored by David Mayerich
1 parent c685653d

added optics python library

Showing 2 changed files with 844 additions and 1 deletions   Show diff stats
python/optics.py 0 → 100644
  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 }
... ...