codebase / stimlib

Commit fdc7f32c10427a43e2cbbe6a7807455f4ca6c007

Authored by David Mayerich
1 parent da132b67

Added a coupled wave library containing code to implement Bryn's homogeneous layered model

Showing 2 changed files with 384 additions and 17 deletions   Show diff stats
Inline Side-by-side
python/coupledwave.py 0 → 100644
  Show/Hide comments   View file @fdc7f32
  1 +# -*- coding: utf-8 -*-
  2 +"""
  3 +Created on Thu Sep 6 22:14:36 2018
  4 +
  5 +@author: david
  6 +"""
  7 +
  8 +import numpy as np
  9 +import matplotlib.pyplot as plt
  10 +
  11 +#evaluate a plane wave for all points in R = [Rx, Ry, Rz]
  12 +# E0 [Ex, Ey, Ez] is the electric field vector
  13 +# k is the k vector
  14 +# n is the refractive index of the material
  15 +# R is a list of coordinates in x, y, and z
  16 +def planewave2field(E0, k, n, R):
  17 + knr = n * (k[0] * R[0] + k[1] * R[1] + k[2] + R[2])
  18 + exp_2pknr = np.exp(1j * 2 * np.pi * knr)
  19 + Ex = E0[0] * exp_2pknr
  20 + Ey = E0[1] * exp_2pknr
  21 + Ez = E0[2] * exp_2pknr
  22 +
  23 + return [Ex, Ey, Ez]
  24 +
  25 +def orthogonalize(E, k):
  26 + s = np.cross(E, k)
  27 + return np.cross(k, s)
  28 +
  29 +class layersample:
  30 +
  31 + #generate a layered sample based on a list of layer refractive indices
  32 + # layer_ri is a list of complex refractive indices for each layer
  33 + # z is the list of boundary positions along z
  34 + def __init__(self, layer_ri, z):
  35 + self.n = np.array(layer_ri).astype(np.complex128)
  36 + self.z = np.array(z)
  37 +
  38 + #calculate the index of the field component associated with a layer
  39 + # l is the layer index [0, L)
  40 + # c is the component (x=0, y=1, z=2)
  41 + # d is the direction (0 = transmission, 1 = reflection)
  42 + def idx(self, l, c, d):
  43 + i = l * 6 + d * 3 + c - 3
  44 + print(i)
  45 + return i
  46 +
  47 + #create a matrix for a single plane wave specified by k and E
  48 + def solve(self, k, E):
  49 +
  50 + #calculate the wavenumber
  51 + kn = np.linalg.norm(k)
  52 +
  53 + #s is the plane wave direction scaled by the refractive index
  54 + s = k / kn * self.n[0]
  55 +
  56 + #allocate space for the matrix
  57 + L = len(self.n)
  58 + M = np.zeros((6*(L-1), 6*(L-1)), dtype=np.complex128)
  59 +
  60 + #allocate space for the RHS vector
  61 + b = np.zeros(6*(L-1), dtype=np.complex128)
  62 +
  63 + #initialize a counter for the equation number
  64 + ei = 0
  65 +
  66 + #calculate the sz component for each layer
  67 + self.sz = np.zeros(L, dtype=np.complex128)
  68 + for l in range(L):
  69 + self.sz[l] = np.sqrt(self.n[l]**2 - s[0]**2 - s[1]**2)
  70 +
  71 + #set constraints based on Gauss' law
  72 + for l in range(0, L):
  73 + #sz = np.sqrt(self.n[l]**2 - s[0]**2 - s[1]**2)
  74 +
  75 + #set the upward components for each layer
  76 + # note that layer L-1 does not have a downward component
  77 + # David I, Equation 7
  78 + if l != L-1:
  79 + M[ei, self.idx(l, 0, 1)] = s[0]
  80 + M[ei, self.idx(l, 1, 1)] = s[1]
  81 + M[ei, self.idx(l, 2, 1)] = -self.sz[l]
  82 + ei = ei+1
  83 +
  84 + #set the downward components for each layer
  85 + # note that layer 0 does not have a downward component
  86 + # Davis I, Equation 6
  87 + if l != 0:
  88 + M[ei, self.idx(l, 0, 0)] = s[0]
  89 + M[ei, self.idx(l, 1, 0)] = s[1]
  90 + M[ei, self.idx(l, 2, 0)] = self.sz[l]
  91 + ei = ei+1
  92 +
  93 + #enforce a continuous field across boundaries
  94 + for l in range(1, L):
  95 + sz0 = self.sz[l-1]
  96 + sz1 = self.sz[l]
  97 + A = np.exp(1j * kn * sz0 * (self.z[l] - self.z[l-1]))
  98 + if l < L - 1:
  99 + B = np.exp(-1j * kn * sz1 * (self.z[l] - self.z[l+1]))
  100 +
  101 +
  102 + #if this is the second layer, use the simplified equations that account for the incident field
  103 + if l == 1:
  104 + M[ei, self.idx(0, 0, 1)] = 1
  105 + M[ei, self.idx(1, 0, 0)] = -1
  106 + if L > 2:
  107 + M[ei, self.idx(1, 0, 1)] = -B
  108 +
  109 + b[ei] = -A * E[0]
  110 + ei = ei + 1
  111 +
  112 + M[ei, self.idx(0, 1, 1)] = 1
  113 + M[ei, self.idx(1, 1, 0)] = -1
  114 + if L > 2:
  115 + M[ei, self.idx(1, 1, 1)] = -B
  116 + b[ei] = -A * E[1]
  117 + ei = ei + 1
  118 +
  119 + M[ei, self.idx(0, 2, 1)] = s[1]
  120 + M[ei, self.idx(0, 1, 1)] = sz0
  121 + M[ei, self.idx(1, 2, 0)] = -s[1]
  122 + M[ei, self.idx(1, 1, 0)] = sz1
  123 + if L > 2:
  124 + M[ei, self.idx(1, 2, 1)] = -B*s[1]
  125 + M[ei, self.idx(1, 1, 1)] = -B*sz1
  126 + b[ei] = A * sz0 * E[1] - A * s[1]*E[2]
  127 + ei = ei + 1
  128 +
  129 + M[ei, self.idx(0, 0, 1)] = -sz0
  130 + M[ei, self.idx(0, 2, 1)] = -s[0]
  131 + M[ei, self.idx(1, 0, 0)] = -sz1
  132 + M[ei, self.idx(1, 2, 0)] = s[0]
  133 + if L > 2:
  134 + M[ei, self.idx(1, 0, 1)] = B*sz1
  135 + M[ei, self.idx(1, 2, 1)] = B*s[0]
  136 + b[ei] = A * s[0] * E[2] - A * sz0 * E[0]
  137 + ei = ei + 1
  138 +
  139 + #if this is the last layer, use the simplified equations that exclude reflections from the last layer
  140 + elif l == L-1:
  141 + M[ei, self.idx(l-1, 0, 0)] = A
  142 + M[ei, self.idx(l-1, 0, 1)] = 1
  143 + M[ei, self.idx(l, 0, 0)] = -1
  144 + ei = ei + 1
  145 +
  146 + M[ei, self.idx(l-1, 1, 0)] = A
  147 + M[ei, self.idx(l-1, 1, 1)] = 1
  148 + M[ei, self.idx(l, 1, 0)] = -1
  149 + ei = ei + 1
  150 +
  151 + M[ei, self.idx(l-1, 2, 0)] = A*s[1]
  152 + M[ei, self.idx(l-1, 1, 0)] = -A*sz0
  153 + M[ei, self.idx(l-1, 2, 1)] = s[1]
  154 + M[ei, self.idx(l-1, 1, 1)] = sz0
  155 + M[ei, self.idx(l, 2, 0)] = -s[1]
  156 + M[ei, self.idx(l, 1, 0)] = sz1
  157 + ei = ei + 1
  158 +
  159 + M[ei, self.idx(l-1, 0, 0)] = A*sz0
  160 + M[ei, self.idx(l-1, 2, 0)] = -A*s[0]
  161 + M[ei, self.idx(l-1, 0, 1)] = -sz0
  162 + M[ei, self.idx(l-1, 2, 1)] = -s[0]
  163 + M[ei, self.idx(l, 0, 0)] = -sz1
  164 + M[ei, self.idx(l, 2, 0)] = s[0]
  165 + ei = ei + 1
  166 + #otherwise use the full set of boundary conditions
  167 + else:
  168 + M[ei, self.idx(l-1, 0, 0)] = A
  169 + M[ei, self.idx(l-1, 0, 1)] = 1
  170 + M[ei, self.idx(l, 0, 0)] = -1
  171 + M[ei, self.idx(l, 0, 1)] = -B
  172 + ei = ei + 1
  173 +
  174 + M[ei, self.idx(l-1, 1, 0)] = A
  175 + M[ei, self.idx(l-1, 1, 1)] = 1
  176 + M[ei, self.idx(l, 1, 0)] = -1
  177 + M[ei, self.idx(l, 1, 1)] = -B
  178 + ei = ei + 1
  179 +
  180 + M[ei, self.idx(l-1, 2, 0)] = A*s[1]
  181 + M[ei, self.idx(l-1, 1, 0)] = -A*sz0
  182 + M[ei, self.idx(l-1, 2, 1)] = s[1]
  183 + M[ei, self.idx(l-1, 1, 1)] = sz0
  184 + M[ei, self.idx(l, 2, 0)] = -s[1]
  185 + M[ei, self.idx(l, 1, 0)] = sz1
  186 + M[ei, self.idx(l, 2, 1)] = -B*s[1]
  187 + M[ei, self.idx(l, 1, 1)] = -B*sz1
  188 + ei = ei + 1
  189 +
  190 + M[ei, self.idx(l-1, 0, 0)] = A*sz0
  191 + M[ei, self.idx(l-1, 2, 0)] = -A*s[0]
  192 + M[ei, self.idx(l-1, 0, 1)] = -sz0
  193 + M[ei, self.idx(l-1, 2, 1)] = -s[0]
  194 + M[ei, self.idx(l, 0, 0)] = -sz1
  195 + M[ei, self.idx(l, 2, 0)] = s[0]
  196 + M[ei, self.idx(l, 0, 1)] = B*sz1
  197 + M[ei, self.idx(l, 2, 1)] = B*s[0]
  198 + ei = ei + 1
  199 +
  200 + #store the matrix and RHS vector (for debugging)
  201 + self.M = M
  202 + self.b = b
  203 +
  204 + #evaluate the linear system
  205 + P = np.linalg.solve(M, b)
  206 +
  207 + #save the results (also for debugging)
  208 + self.P = P
  209 +
  210 + #store the coefficients for each layer
  211 + self.Ptrans = []
  212 + self.Prefl = []
  213 + for l in range(L):
  214 + if l == 0:
  215 + self.Ptrans.append((E[0], E[1], E[2]))
  216 + else:
  217 + px = P[self.idx(l, 0, 0)]
  218 + py = P[self.idx(l, 1, 0)]
  219 + pz = P[self.idx(l, 2, 0)]
  220 + self.Ptrans.append((px, py, pz))
  221 +
  222 + if l == L-1:
  223 + self.Prefl.append((0, 0, 0))
  224 + else:
  225 + px = P[self.idx(l, 0, 1)]
  226 + py = P[self.idx(l, 1, 1)]
  227 + pz = P[self.idx(l, 2, 1)]
  228 + self.Prefl.append((px, py, pz))
  229 +
  230 + #this converts the coefficients into a NumPy array for convenience later on
  231 + self.Ptrans = np.array(self.Ptrans)
  232 + self.Prefl = np.array(self.Prefl)
  233 +
  234 + #store values required for evaluation
  235 + #store k
  236 + self.k = kn
  237 +
  238 + #store sx and sy
  239 + self.s = np.array([s[0], s[1]])
  240 +
  241 + self.solved = True
  242 +
  243 + #evaluate a solved homogeneous substrate
  244 + def evaluate(self, X, Y, Z):
  245 +
  246 + if not self.solved:
  247 + print("ERROR: the layered substrate hasn't been solved")
  248 + return
  249 +
  250 + #this code is a bit cumbersome and could probably be optimized
  251 + # Basically, it vectorizes everything by creating an image
  252 + # such that the value at each pixel is the corresponding layer
  253 + # that the pixel resides in. That index is then used to calculate
  254 + # the field within the layer
  255 +
  256 + #allocate space for layer indices
  257 + LI = np.zeros(Z.shape, dtype=np.int)
  258 +
  259 + #find the layer index for each sample point
  260 + L = len(self.z)
  261 + LI[Z < self.z[0]] = 0
  262 + for l in range(L-1):
  263 + idx = np.logical_and(Z > self.z[l], Z <= self.z[l+1])
  264 + LI[idx] = l
  265 + LI[Z > self.z[-1]] = L - 1
  266 +
  267 + #calculate the appropriate phase shift for the wave transmitted through the layer
  268 + Ph_t = np.exp(1j * self.k * self.sz[LI] * (Z - self.z[LI]))
  269 +
  270 + #calculate the appropriate phase shift for the wave reflected off of the layer boundary
  271 + LIp = LI + 1
  272 + LIp[LIp >= L] = 0
  273 + Ph_r = np.exp(-1j * self.k * self.sz[LI] * (Z - self.z[LIp]))
  274 + Ph_r[LI >= L-1] = 0
  275 +
  276 + #calculate the phase shift based on the X and Y positions
  277 + Ph_xy = np.exp(1j * self.k * (self.s[0] * X + self.s[1] * Y))
  278 +
  279 + #apply the phase shifts
  280 + Et = self.Ptrans[LI] * Ph_t[:, :, None]
  281 + Er = self.Prefl[LI] * Ph_r[:, :, None]
  282 +
  283 + #add everything together coherently
  284 + E = (Et + Er) * Ph_xy[:, :, None]
  285 +
  286 + #return the electric field
  287 + return E
  288 +
  289 +
  290 +#This sample code produces a field similar to Figure 2 in Davis et al., 2010
  291 +#set the material properties
  292 +depths = [-50, -15, 15, 40]
  293 +n = [1.0, 1.4+1j*0.05, 1.4, 1.0]
  294 +
  295 +#create a layered sample
  296 +layers = layersample(n, depths)
  297 +
  298 +#set the input light parameters
  299 +d = np.array([0.2, 0, 1])
  300 +d = d / np.linalg.norm(d)
  301 +l = 5.0
  302 +k = 2 * np.pi / l * d
  303 +E0 = [0, 1, 0]
  304 +E0 = orthogonalize(E0, d)
  305 +
  306 +#solve for the substrate field
  307 +layers.solve(k, E0)
  308 +
  309 +#set the simulation domain
  310 +N = 512
  311 +M = 1024
  312 +D = [-50, 80]
  313 +
  314 +x = np.linspace(D[0], D[1], N)
  315 +z = np.linspace(D[0], D[1], M)
  316 +[X, Z] = np.meshgrid(x, z)
  317 +Y = np.zeros(X.shape)
  318 +E = layers.evaluate(X, Y, Z)
  319 +
  320 +Er = np.real(E)
  321 +I = Er[..., 0] ** 2 + Er[..., 1] **2 + Er[..., 2] ** 2
  322 +plt.figure()
  323 +plt.set_cmap("afmhot")
  324 +plt.imshow(Er[..., 1])
... ...
stim/image/image.h
  Show/Hide comments   View file @fdc7f32
... ... @@ -33,6 +33,7 @@ class image{
33 33  
34 34 T* img; //pointer to the image data (interleaved RGB for color)
35 35 size_t R[3];
  36 + bool interleaved = true; //by default the data is interleaved
36 37  
37 38 inline size_t X() const { return R[1]; }
38 39 inline size_t Y() const { return R[2]; }
... ... @@ -73,13 +74,24 @@ class image{
73 74  
74 75 #ifdef USING_OPENCV
75 76 int cv_type(){
76   - if(typeid(T) == typeid(unsigned char)) return CV_MAKETYPE(CV_8U, (int)C());
77   - if(typeid(T) == typeid(char)) return CV_MAKETYPE(CV_8S, (int)C());
78   - if(typeid(T) == typeid(unsigned short)) return CV_MAKETYPE(CV_16U, (int)C());
79   - if(typeid(T) == typeid(short)) return CV_MAKETYPE(CV_16S, (int)C());
80   - if(typeid(T) == typeid(int)) return CV_MAKETYPE(CV_32S, (int)C());
81   - if(typeid(T) == typeid(float)) return CV_MAKETYPE(CV_32F, (int)C());
82   - if(typeid(T) == typeid(double)) return CV_MAKETYPE(CV_64F, (int)C());
  77 + if (C() == 1 || C() == 3) { //if the image has 1 or 3 channels, this will work
  78 + if (typeid(T) == typeid(unsigned char)) return CV_MAKETYPE(CV_8U, (int)C());
  79 + if (typeid(T) == typeid(char)) return CV_MAKETYPE(CV_8S, (int)C());
  80 + if (typeid(T) == typeid(unsigned short)) return CV_MAKETYPE(CV_16U, (int)C());
  81 + if (typeid(T) == typeid(short)) return CV_MAKETYPE(CV_16S, (int)C());
  82 + if (typeid(T) == typeid(int)) return CV_MAKETYPE(CV_32S, (int)C());
  83 + if (typeid(T) == typeid(float)) return CV_MAKETYPE(CV_32F, (int)C());
  84 + if (typeid(T) == typeid(double)) return CV_MAKETYPE(CV_64F, (int)C());
  85 + }
  86 + else if (C() == 4) { //OpenCV only supports saving BGR images - 4-channel isn't supported
  87 + if (typeid(T) == typeid(unsigned char)) return CV_MAKETYPE(CV_8U, 3);
  88 + if (typeid(T) == typeid(char)) return CV_MAKETYPE(CV_8S, 3);
  89 + if (typeid(T) == typeid(unsigned short)) return CV_MAKETYPE(CV_16U, 3);
  90 + if (typeid(T) == typeid(short)) return CV_MAKETYPE(CV_16S, 3);
  91 + if (typeid(T) == typeid(int)) return CV_MAKETYPE(CV_32S, 3);
  92 + if (typeid(T) == typeid(float)) return CV_MAKETYPE(CV_32F, 3);
  93 + if (typeid(T) == typeid(double)) return CV_MAKETYPE(CV_64F, 3);
  94 + }
83 95  
84 96 std::cout<<"ERROR in stim::image::cv_type - no valid data type found"<<std::endl;
85 97 exit(1);
... ... @@ -351,15 +363,37 @@ public:
351 363 }
352 364 #ifdef USING_OPENCV
353 365 //OpenCV uses an interleaved format, so convert first and then output
354   - T* buffer = (T*) malloc(bytes());
355   -
356   - if(C() == 1)
  366 +
  367 + //single-channel image
  368 + if (C() == 1) {
  369 + T* buffer = (T*)malloc(bytes());
357 370 memcpy(buffer, img, bytes());
358   - else if(C() == 3)
  371 + cv::Mat cvImage((int)Y(), (int)X(), cv_type(), buffer);
  372 + cv::imwrite(filename, cvImage);
  373 + free(buffer);
  374 + }
  375 + //3-channel image
  376 + else if (C() == 3) {
  377 + T* buffer = (T*)malloc(bytes());
359 378 get_interleaved_bgr(buffer);
360   - cv::Mat cvImage((int)Y(), (int)X(), cv_type(), buffer);
361   - cv::imwrite(filename, cvImage);
362   - free(buffer);
  379 + cv::Mat cvImage((int)Y(), (int)X(), cv_type(), buffer);
  380 + cv::imwrite(filename, cvImage);
  381 + free(buffer);
  382 + }
  383 + //ARGB image
  384 + else if (C() == 4) {
  385 + size_t channelsize = X() * Y() * sizeof(T);
  386 + T* buffer = (T*)malloc(bytes() - channelsize);
  387 + get_interleaved_bgr(buffer);
  388 + cv::Mat cvImage((int)Y(), (int)X(), cv_type(), buffer);
  389 + cv::imwrite(filename, cvImage);
  390 + free(buffer);
  391 + }
  392 + else {
  393 + throw std::runtime_error("ERROR: number of channels not supported for image saving.");
  394 + }
  395 +
  396 +
363 397 #else
364 398 if (file.extension() == "ppm")
365 399 save_netpbm(filename);
... ... @@ -419,12 +453,21 @@ public:
419 453 }
420 454  
421 455 void get_interleaved_bgr(T* data){
422   -
423 456 //for each channel
  457 + T* source;
  458 + if (C() == 3) {
  459 + source = img; //if the image has 3 channels, interleave all three
  460 + }
  461 + else if (C() == 4) {
  462 + source = img + X() * Y(); //if the image has 4 channels, skip the alpha channel
  463 + }
  464 + else {
  465 + throw std::runtime_error("ERROR: a BGR image must be 3 or 4 channels"); //throw an error if any other number of channels is provided
  466 + }
424 467 for(size_t y = 0; y < Y(); y++){
425 468 for(size_t x = 0; x < X(); x++){
426   - for(size_t c = 0; c < C(); c++){
427   - data[y * X() * C() + x * C() + (2-c)] = img[idx(x, y, c)];
  469 + for(size_t c = 0; c < 3; c++){
  470 + data[y * X() * 3 + x * 3 + (2-c)] = source[y*3*R[1] + x*3 + c];
428 471 }
429 472 }
430 473 }
... ...