Move MIN,MAX,SQ macros to a common header
[qpms.git] / qpms / legacy.py
blob3fb2d712d5f556214c1420ff316b1588cceb7c25
1 import math
2 import numpy as np
3 nx = None
5 _s3 = math.sqrt(3)
7 from qpms_c import get_mn_y, trans_calculator
8 from .qpms_p import cart2sph
13 # Functions copied from scattering_xu, additionaly normalized
14 from py_gmm.gmm import vec_trans as vc
16 #@ujit
17 def q_max(m,n,μ,ν):
18 return min(n,ν,(n+ν-abs(m+μ))/2)
20 q_max_v = np.vectorize(q_max)
22 # returns array with indices corresponding to q
23 # argument q does nothing for now
24 #@ujit
25 def a_q(m,n,μ,ν,q = None):
26 qm=q_max(m,n,μ,ν)
27 res, err= vc.gaunt_xu(m,n,μ,ν,qm)
28 if(err):
29 print("m,n,μ,ν,qm = ",m,n,μ,ν,qm)
30 raise ValueError('Something bad in the fortran subroutine gaunt_xu happened')
31 return res
33 a_q_v = np.vectorize(a_q)
36 # All arguments are single numbers (for now)
37 # ZDE VYCHÁZEJÍ DIVNÁ ZNAMÉNKA
38 #@ujit
39 def(m,n,μ,ν,kdlj,θlj,φlj,r_ge_d,J):
40 """
41 The à translation coefficient for spherical vector waves.
43 Parameters
44 ----------
45 m, n: int
46 The indices (degree and order) of the destination basis.
47 μ, ν: int
48 The indices of the source basis wave.
49 kdlj, θlj, φlj: float
50 The spherical coordinates of the relative position of
51 the new center vs. the old one (R_new - R_old);
52 the distance has to be already multiplied by the wavenumber!
53 r_ge_d: TODO
54 J: 1, 2, 3 or 4
55 Type of the wave in the old center.
57 Returns
58 -------
59 TODO
61 Bugs
62 ----
63 gevero's gaunt coefficient implementation fails for large m, n (the unsafe territory
64 is somewhere around -72, 80)
66 """
67 exponent=(math.lgamma(2*n+1)-math.lgamma(n+2)+math.lgamma(2*ν+3)-math.lgamma(ν+2)
68 +math.lgamma(n+ν+m-μ+1)-math.lgamma(n-m+1)-math.lgamma(ν+μ+1)
69 +math.lgamma(n+ν+1) - math.lgamma(2*(n+ν)+1))
70 presum = math.exp(exponent)
71 presum = presum * np.exp(1j*(μ-m)*φlj) * (-1)**m * 1j**(ν+n) / (4*n)
72 qmax = math.floor(q_max(-m,n,μ,ν)) #nemá tu být +m?
73 q = np.arange(qmax+1, dtype=int)
74 # N.B. -m !!!!!!
75 a1q = a_q(-m,n,μ,ν) # there is redundant calc. of qmax
76 ã1q = a1q / a1q[0]
77 p = n+ν-2*q
78 if(r_ge_d):
79 J = 1
80 zp = zJn(n+ν,kdlj,J)[0][p]
81 Pp = lpmv(μ-m,p,math.cos(θlj))
82 summandq = (n*(n+1) + ν*(ν+1) - p*(p+1)) * (-1)**q * ã1q * zp * Pp
84 # Taylor normalisation v2, proven to be equivalent (NS which is better)
85 prenormratio = 1j**(ν-n) * math.sqrt(((2*ν+1)/(2*n+1))* math.exp(
86 math.lgamma(n+m+1)-math.lgamma(n-m+1)+math.lgamma(ν-μ+1)-math.lgamma(ν+μ+1)))
87 presum = presum / prenormratio
89 # Taylor normalisation
90 #prenormmn = math.sqrt((2*n + 1)*math.factorial(n-m)/(4*π*factorial(n+m)))
91 #prenormμν = math.sqrt((2*ν + 1)*math.factorial(ν-μ)/(4*π*factorial(ν+μ)))
92 #presum = presum * prenormμν / prenormmn
94 return presum * np.sum(summandq)
96 # ZDE OPĚT JINAK ZNAMÉNKA než v Xu (J. comp. phys 127, 285)
97 #@ujit
98 def(m,n,μ,ν,kdlj,θlj,φlj,r_ge_d,J):
99 """
100 The B̃ translation coefficient for spherical vector waves.
102 Parameters
103 ----------
104 m, n: int
105 The indices (degree and order) of the destination basis.
106 μ, ν: int
107 The indices of the source basis wave.
108 kdlj, θlj, φlj: float
109 The spherical coordinates of the relative position of
110 the new center vs. the old one (R_new - R_old);
111 the distance has to be already multiplied by the wavenumber!
112 r_ge_d: TODO
113 J: 1, 2, 3 or 4
114 Type of the wave in the old center.
116 Returns:
117 --------
118 TODO
120 exponent=(math.lgamma(2*n+3)-math.lgamma(n+2)+math.lgamma(2*ν+3)-math.lgamma(ν+2)
121 +math.lgamma(n+ν+m-μ+2)-math.lgamma(n-m+1)-math.lgamma(ν+μ+1)
122 +math.lgamma(n+ν+2) - math.lgamma(2*(n+ν)+3))
123 presum = math.exp(exponent)
124 presum = presum * np.exp(1j*(μ-m)*φlj) * (-1)**m * 1j**(ν+n+1) / (
125 (4*n)*(n+1)*(n+m+1))
126 Qmax = math.floor(q_max(-m,n+1,μ,ν))
127 q = np.arange(Qmax+1, dtype=int)
128 if (μ == ν): # it would disappear in the sum because of the factor (ν-μ) anyway
129 ã2q = 0
130 else:
131 a2q = a_q(-m-1,n+1,μ+1,ν)
132 ã2q = a2q / a2q[0]
133 a3q = a_q(-m,n+1,μ,ν)
134 ã3q = a3q / a3q[0]
135 #print(len(a2q),len(a3q))
136 p = n+ν-2*q
137 if(r_ge_d):
138 J = 1
139 zp_ = zJn(n+1+ν,kdlj,J)[0][p+1] # je ta +1 správně?
140 Pp_ = lpmv(μ-m,p+1,math.cos(θlj))
141 summandq = ((2*(n+1)*(ν-μ)*ã2q
142 -(-ν*(ν+1) - n*(n+3) - 2*μ*(n+1)+p*(p+3))* ã3q)
143 *(-1)**q * zp_ * Pp_)
145 # Taylor normalisation v2, proven to be equivalent
146 prenormratio = 1j**(ν-n) * math.sqrt(((2*ν+1)/(2*n+1))* math.exp(
147 math.lgamma(n+m+1)-math.lgamma(n-m+1)+math.lgamma(ν-μ+1)-math.lgamma(ν+μ+1)))
148 presum = presum / prenormratio
150 ## Taylor normalisation
151 #prenormmn = math.sqrt((2*n + 1)*math.factorial(n-m)/(4*π*factorial(n+m)))
152 #prenormμν = math.sqrt((2*ν + 1)*math.factorial(ν-μ)/(4*π*factorial(ν+μ)))
153 #presum = presum * prenormμν / prenormmn
155 return presum * np.sum(summandq)
164 #@ujit
165 def G_Mie_scat_precalc_cart(source_cart, dest_cart, RH, RV, a, nmax, k_i, k_e, μ_i=1, μ_e=1, J_ext=1, J_scat=3):
167 r1_cart (destination), r2_cart (source) and the result are in cartesian coordinates
168 the result indices are in the source-destination order
169 TODO
171 my, ny = get_mn_y(nmax)
172 nelem = len(my)
173 #source to origin
174 so_sph = cart2sph(-source_cart)
175 kd_so = k_e * so_sph[0]
176 θ_so = so_sph[1]
177 φ_so = so_sph[2]
178 # Decomposition of the source N_0,1, N_-1,1, and N_1,1 in the nanoparticle center
179 p_0 = np.empty((nelem), dtype=np.complex_)
180 q_0 = np.empty((nelem), dtype=np.complex_)
181 p_minus = np.empty((nelem), dtype=np.complex_)
182 q_minus = np.empty((nelem), dtype=np.complex_)
183 p_plus = np.empty((nelem), dtype=np.complex_)
184 q_plus = np.empty((nelem), dtype=np.complex_)
185 for y in range(nelem):
186 m = my[y]
187 n = ny[y]
188 p_0[y] =(m,n, 0,1,kd_so,θ_so,φ_so,False,J=J_scat)
189 q_0[y] =(m,n, 0,1,kd_so,θ_so,φ_so,False,J=J_scat)
190 p_minus[y] =(m,n,-1,1,kd_so,θ_so,φ_so,False,J=J_scat)
191 q_minus[y] =(m,n,-1,1,kd_so,θ_so,φ_so,False,J=J_scat)
192 p_plus[y] =(m,n, 1,1,kd_so,θ_so,φ_so,False,J=J_scat)
193 q_plus[y] =(m,n, 1,1,kd_so,θ_so,φ_so,False,J=J_scat)
194 a_0 = RV[ny] * p_0
195 b_0 = RH[ny] * q_0
196 a_plus = RV[ny] * p_plus
197 b_plus = RH[ny] * q_plus
198 a_minus = RV[ny] * p_minus
199 b_minus = RH[ny] * q_minus
200 orig2dest_sph = cart2sph(dest_cart)
201 orig2dest_sph[0] = k_e*orig2dest_sph[0]
202 M_dest_y, N_dest_y = vswf_yr1(orig2dest_sph,nmax,J=J_scat)
203 # N.B. these are in the local cartesian coordinates (r̂,θ̂,φ̂)
204 N_dest_0 = np.sum(a_0[:,ň] * N_dest_y, axis=-2)
205 M_dest_0 = np.sum(b_0[:,ň] * M_dest_y, axis=-2)
206 N_dest_plus = np.sum(a_plus[:,ň] * N_dest_y, axis=-2)
207 M_dest_plus = np.sum(b_plus[:,ň] * M_dest_y, axis=-2)
208 N_dest_minus = np.sum(a_minus[:,ň]* N_dest_y, axis=-2)
209 M_dest_minus = np.sum(b_minus[:,ň]* M_dest_y, axis=-2)
210 prefac = math.sqrt(1/(4*3*π))#/ε_0
211 G_sourcez_dest = prefac * (N_dest_0+M_dest_0)
212 G_sourcex_dest = prefac * (N_dest_minus+M_dest_minus-N_dest_plus-M_dest_plus)/math.sqrt(2)
213 G_sourcey_dest = prefac * (N_dest_minus+M_dest_minus+N_dest_plus+M_dest_plus)/(1j*math.sqrt(2))
214 G_source_dest = np.array([G_sourcex_dest, G_sourcey_dest, G_sourcez_dest])
215 # To global cartesian coordinates:
216 G_source_dest = sph_loccart2cart(G_source_dest, sph=orig2dest_sph, axis=-1)
217 return G_source_dest
219 #@ujit
220 def G_Mie_scat_cart(source_cart, dest_cart, a, nmax, k_i, k_e, μ_i=1, μ_e=1, J_ext=1, J_scat=3):
222 TODO
224 RH, RV, TH, TV = mie_coefficients(a=a, nmax=nmax, k_i=k_i, k_e=k_e, μ_i=μ_i, μ_e=μ_e, J_ext=J_ext, J_scat=J_scat)
225 return G_Mie_scat_precalc_cart_new(source_cart, dest_cart, RH, RV, a, nmax, k_i, k_e, μ_i, μ_e, J_ext, J_scat)
228 #TODO
229 def cross_section_Mie_precalc():
230 pass
232 def cross_section_Mie(a, nmax, k_i, k_e, μ_i, μ_e,):
233 pass
238 ####################
239 # Array simulations
240 ####################
244 def scatter_plane_wave(omega, epsilon_b, positions, Tmatrices, k_dirs, E_0s, #saveto = None
247 Solves the plane wave linear scattering problem for a structure of "non-touching" particles
248 for one frequency and arbitrary number K of incoming plane waves.
250 Parameters
251 ----------
252 omega : positive number
253 The frequency of the field.
254 epsilon_b : complex number
255 Permittivity of the background medium (which has to be isotropic).
256 positions : (N,3)-shaped real array
257 Cartesian positions of the particles.
258 TMatrices : (N,2,nelem,2,nelem) or compatible
259 The T-matrices in the "Taylor convention" describing the scattering on a single nanoparticle.
260 If all the particles are identical and equally oriented, only one T-matrix can be given.
261 nelems = (lMax + 2) * lMax, where lMax is the highest multipole order to which the scattering
262 is calculated.
263 k_dirs : (K,3)-shaped real array or compatible
264 The direction of the incident field wave vector, normalized to one.
265 E_0s : (K,3)-shaped complex array or compatible
266 The electric intensity amplitude of the incident field.
268 Returns
269 -------
270 ab : (K, N, 2, nelem)-shaped complex array
271 The a (electric wave), b (magnetic wave) coefficients of the outgoing field for each particle
272 # Fuck this, it will be wiser to make separate function to calculate those from ab:
273 # sigma_xxx : TODO (K, 2, nelem)
274 # TODO partial (TODO which?) cross-section for each type of outgoing waves, summed over all
275 # nanoparticles (total cross section is given by the sum of this.)
277 nelem = TMatrices.shape[-1]
278 if ((nelem != TMatrices.shape[-3]) or (2 != TMatrices.shape[-2]) or (2 != TMatrices.shape[-4])):
279 raise ValueError('The T-matrices must be of shape (N, 2, nelem, 2, nelem) but are of shape %s' % (str(TMatrices.shape),))
280 lMax = nelem2lMax(nelem)
281 if not lMax:
282 raise ValueError('The "nelem" dimension of T-matrix has invalid value (%d).' % nelem)
283 # TODO perhaps more checks.
284 raise Error('Not implemented.')
285 pass
290 import warnings
291 #@ujit
292 def scatter_plane_wave_rectarray(omega, epsilon_b, xN, yN, xd, yd, TMatrices, k_dirs, E_0s,
293 return_pq_0 = False, return_pq= False, return_xy = False, watch_time = False):
295 Solves the plane wave linear scattering problem for a rectangular array of particles
296 for one frequency and arbitrary number K of incoming plane waves.
298 Parameters
299 ----------
300 omega : positive number
301 The frequency of the field.
302 epsilon_b : complex number
303 Permittivity of the background medium (which has to be isotropic).
304 xN, yN : positive integers
305 Particle numbers in the x and y dimensions
306 xd, yd : positive numbers
307 Periodicities in the x and y direction
308 TMatrices : (xN, yN,2,nelem,2,nelem) or compatible or (2,nelem,2,nelem)
309 The T-matrices in the "Taylor convention" describing the scattering on a single nanoparticle.
310 If all the particles are identical and equally oriented, only one T-matrix can be given.
311 nelems = (lMax + 2) * lMax, where lMax is the highest multipole order to which the scattering
312 is calculated.
313 Electric wave index is 0, magnetic wave index is 1.
314 k_dirs : (K,3)-shaped real array or compatible
315 The direction of the incident field wave vector, normalized to one.
316 E_0s : (K,3)-shaped complex array or compatible
317 The electric intensity amplitude of the incident field.
318 return_pq_0 : bool
319 Return also the multipole decomposition coefficients of the incoming plane wave.
320 return_pq : bool NOT IMPLEMENTED
321 Return also the multipole decomposition coefficients of the field incoming to each
322 particle (inc. the field scattered from other particles.
323 return_xy : bool
324 Return also the cartesian x, y positions of the particles.
325 watch_time : bool
326 Inform about the progress on stderr
328 Returns
329 -------
330 ab : (K, xN, yN, 2, nelem)-shaped complex array
331 The a (electric wave), b (magnetic wave) coefficients of the outgoing field for each particle.
332 If none of return_pq or return_xy is set, the array is not enclosed in a tuple.
333 pq_0 : (K, xN, yn, 2, nelem)-shaped complex array
334 The p_0 (electric wave), b_0 (magnetic wave) coefficients of the incoming plane wave for each particle.
335 pq : (K, xN, yN, 2, nelem)-shaped complex array NOT IMPLEMENTED
336 The p (electric wave), q (magnetic wave) coefficients of the total exciting field
337 for each particle (including the field scattered from other particles)
338 x, y : (xN, yN)-shaped real array
339 The x,y positions of the nanoparticles.
341 if (watch_time):
342 timec = time.time()
343 print('%.4f: running scatter_plane_wave_rectarray' % timec, file = sys.stderr)
344 sys.stderr.flush()
345 nelem = TMatrices.shape[-1]
346 if ((nelem != TMatrices.shape[-3]) or (2 != TMatrices.shape[-2]) or (2 != TMatrices.shape[-4])):
347 raise ValueError('The T-matrices must be of shape (N, 2, nelem, 2, nelem) but are of shape %s' % (str(TMatrices.shape),))
348 lMax = nelem2lMax(nelem)
349 if not lMax:
350 raise ValueError('The "nelem" dimension of T-matrix has invalid value (%d).' % nelem)
351 if (watch_time):
352 print('xN = %d, yN = %d, lMax = %d' % (xN, yN, lMax), file = sys.stderr)
353 sys.stderr.flush()
354 # TODO perhaps more checks.
355 k_out = omega * math.sqrt(epsilon_b) / c # wave number
356 my, ny = get_mn_y(lMax)
357 N = yN * xN
359 J_scat=3
360 J_ext=1
362 # Do something with this ugly indexing crap
363 xind, yind = np.meshgrid(np.arange(xN),np.arange(yN), indexing='ij')
364 xind = xind.flatten()
365 yind = yind.flatten()
366 xyind = np.stack((xind, yind, np.zeros((xind.shape),dtype=int)),axis=-1)
367 cart_lattice=xyind * np.array([xd, yd, 0])
368 x=cart_lattice[:,0]
369 y=cart_lattice[:,1]
370 xyind = xyind[:,0:2]
372 # Lattice speedup
373 if (watch_time):
374 timec = time.time()
375 print('%.4f: calculating the %d translation matrix elements' % (timec, 8*nelem*nelem*xN*yN), file = sys.stderr)
376 sys.stderr.flush()
377 Agrid = np.zeros((nelem, 2*xN, 2*yN, nelem),dtype=np.complex_)
378 Bgrid = np.zeros((nelem, 2*xN, 2*yN, nelem),dtype=np.complex_)
379 for yl in range(nelem): # source
380 for xij in range(2*xN):
381 for yij in range(2*yN):
382 for yj in range(nelem): #dest
383 if((yij != yN) or (xij != xN)):
384 d_l2j = cart2sph(np.array([(xij-xN)*xd, (yij-yN)*yd, 0]))
385 Agrid[yj, xij, yij, yl] =(my[yj],ny[yj],my[yl],ny[yl],kdlj=d_l2j[0]*k_out,θlj=d_l2j[1],φlj=d_l2j[2],r_ge_d=False,J=J_scat)
386 Bgrid[yj, xij, yij, yl] =(my[yj],ny[yj],my[yl],ny[yl],kdlj=d_l2j[0]*k_out,θlj=d_l2j[1],φlj=d_l2j[2],r_ge_d=False,J=J_scat)
388 # Translation coefficient matrix T
389 if (watch_time):
390 timecold = timec
391 timec = time.time()
392 print('%4f: translation matrix elements calculated (elapsed %.2f s), filling the matrix'
393 % (timec, timec-timecold), file = sys.stderr)
394 sys.stderr.flush()
395 transmat = np.zeros((xN* yN, 2, nelem, xN* yN, 2, nelem),dtype=np.complex_)
396 for l in range(N):
397 xil, yil = xyind[l]
398 for j in range(N):
399 xij, yij = xyind[j]
400 if (l!=j):
401 transmat[j,0,:,l,0,:] = Agrid[:, xij - xil + xN, yij - yil + yN, :]
402 transmat[j,0,:,l,1,:] = Bgrid[:, xij - xil + xN, yij - yil + yN, :]
403 transmat[j,1,:,l,0,:] = Bgrid[:, xij - xil + xN, yij - yil + yN, :]
404 transmat[j,1,:,l,1,:] = Agrid[:, xij - xil + xN, yij - yil + yN, :]
405 Agrid = None
406 Bgrid = None
407 if (watch_time):
408 timecold = timec
409 timec = time.time()
410 print('%4f: translation matrix filled (elapsed %.2f s), building the interaction matrix'
411 % (timec, timec-timecold), file=sys.stderr)
412 sys.stderr.flush()
414 # Now we solve a linear problem (1 - M T) A = M P_0 where M is the T-matrix :-)
415 MT = np.empty((N,2,nelem,N,2,nelem),dtype=np.complex_)
417 TMatrices = np.broadcast_to(TMatrices, (xN, yN, 2, nelem, 2, nelem))
418 for j in range(N): # I wonder how this can be done without this loop...
419 xij, yij = xyind[j]
420 MT[j] = np.tensordot(TMatrices[xij, yij],transmat[j],axes=([-2,-1],[0,1]))
421 MT.shape = (N*2*nelem, N*2*nelem)
422 leftmatrix = np.identity(N*2*nelem) - MT
423 MT = None
424 if (watch_time):
425 timecold = timec
426 timec = time.time()
427 print('%.4f: interaction matrix complete (elapsed %.2f s)' % (timec, timec-timecold),
428 file=sys.stderr)
429 sys.stderr.flush()
431 if ((1 == k_dirs.ndim) and (1 == E_0s.ndim)):
432 k_cart = k_dirs * k_out # wave vector of the incident plane wave
433 pq_0 = np.zeros((N,2,nelem), dtype=np.complex_)
434 p_y0, q_y0 = plane_pq_y(lMax, k_cart, E_0s)
435 pq_0[:,0,:] = np.exp(1j*np.sum(k_cart[ň,:]*cart_lattice,axis=-1))[:, ň] * p_y0[ň, :]
436 pq_0[:,1,:] = np.exp(1j*np.sum(k_cart[ň,:]*cart_lattice,axis=-1))[:, ň] * q_y0[ň, :]
437 if (return_pq_0):
438 pq_0_arr = pq_0
439 MP_0 = np.empty((N,2,nelem),dtype=np.complex_)
440 #if (watch_time):
441 # print('%4f: building the interaction matrix' % time.time(), file=sys.stderr)
443 for j in range(N): # I wonder how this can be done without this loop...
444 MP_0[j] = np.tensordot(TMatrices[xij, yij],pq_0[j],axes=([-2,-1],[-2,-1]))
445 MP_0.shape = (N*2*nelem,)
447 if (watch_time):
448 timecold = time.time()
449 print('%4f: solving the scattering problem for single incoming wave' % timecold,
450 file = sys.stderr)
451 sys.stderr.flush()
452 ab = np.linalg.solve(leftmatrix, MP_0)
453 if watch_time:
454 timec = time.time()
455 print('%4f: solved (elapsed %.2f s)' % (timec, timec-timecold), file=sys.stderr)
456 sys.stderr.flush()
458 ab.shape = (xN, yN, 2, nelem)
459 else:
460 # handle "broadcasting" for k, E
461 if 1 == k_dirs.ndim:
462 k_dirs = k_dirs[ň,:]
463 if 1 == E_0s.ndim:
464 E_0s = E_0s[ň,:]
465 K = max(E_0s.shape[-2], k_dirs.shape[-2])
466 k_dirs = np.broadcast_to(k_dirs,(K,3))
467 E_0s = np.broadcast_to(E_0s, (K,3))
469 # А ну, чики-брики и в дамки!
470 if watch_time:
471 timecold = time.time()
472 print('%.4f: factorizing the interaction matrix' % timecold, file=sys.stderr)
473 sys.stderr.flush()
474 lupiv = scipy.linalg.lu_factor(leftmatrix, overwrite_a=True)
475 leftmatrix = None
476 if watch_time:
477 timec = time.time()
478 print('%.4f: factorization complete (elapsed %.2f s)' % (timec, timec-timecold),
479 file = sys.stderr)
480 print('%.4f: solving the scattering problem for %d incoming waves' % (timec, K),
481 file=sys.stderr)
482 sys.stderr.flush()
483 timecold = timec
485 if (return_pq_0):
486 pq_0_arr = np.zeros((K,N,2,nelem), dtype=np.complex_)
487 ab = np.empty((K,N*2*nelem), dtype=complex)
488 for ki in range(K):
489 k_cart = k_dirs[ki] * k_out
490 pq_0 = np.zeros((N,2,nelem), dtype=np.complex_)
491 p_y0, q_y0 = plane_pq_y(lMax, k_cart, E_0s[ki])
492 pq_0[:,0,:] = np.exp(1j*np.sum(k_cart[ň,:]*cart_lattice,axis=-1))[:, ň] * p_y0[ň, :]
493 pq_0[:,1,:] = np.exp(1j*np.sum(k_cart[ň,:]*cart_lattice,axis=-1))[:, ň] * q_y0[ň, :]
494 if (return_pq_0):
495 pq_0_arr[ki] = pq_0
496 MP_0 = np.empty((N,2,nelem),dtype=np.complex_)
497 for j in range(N): # I wonder how this can be done without this loop...
498 MP_0[j] = np.tensordot(TMatrices[xij, yij],pq_0[j],axes=([-2,-1],[-2,-1]))
499 MP_0.shape = (N*2*nelem,)
501 ab[ki] = scipy.linalg.lu_solve(lupiv, MP_0)
502 ab.shape = (K, xN, yN, 2, nelem)
503 if watch_time:
504 timec = time.time()
505 print('%.4f: done (elapsed %.2f s)' % (timec, timec-timecold),file = sys.stderr)
506 sys.stderr.flush()
507 if not (return_pq_0 + return_pq + return_xy):
508 return ab
509 returnlist = [ab]
510 if (return_pq_0):
511 pq_0_arr.shape = ab.shape
512 returnlist.append(pq_0_arr)
513 if (return_pq):
514 warnings.warn("return_pq not implemented, ignoring")
515 # returnlist.append(pq_arr)
516 if (return_xy):
517 returnlist.append(x)
518 returnlist.append(y)
519 return tuple(returnlist)
522 import warnings
523 #@ujit
524 def scatter_constmultipole_rectarray(omega, epsilon_b, xN, yN, xd, yd, TMatrices, pq_0_c = 1,
525 return_pq= False, return_xy = False, watch_time = False):
527 Solves the plane wave linear scattering problem for a rectangular array of particles
528 for one frequency and constant exciting spherical waves throughout the array.
530 Parameters
531 ----------
532 omega : positive number
533 The frequency of the field.
534 epsilon_b : complex number
535 Permittivity of the background medium (which has to be isotropic).
536 xN, yN : positive integers
537 Particle numbers in the x and y dimensions
538 xd, yd : positive numbers
539 Periodicities in the x and y direction
540 TMatrices : (xN, yN,2,nelem,2,nelem) or compatible or (2,nelem,2,nelem)
541 The T-matrices in the "Taylor convention" describing the scattering on a single nanoparticle.
542 If all the particles are identical and equally oriented, only one T-matrix can be given.
543 nelems = (lMax + 2) * lMax, where lMax is the highest multipole order to which the scattering
544 is calculated.
545 Electric wave index is 0, magnetic wave index is 1.
546 pq_0_c : (nelem)-shaped complex array or compatible
547 The initial excitation coefficients for the ("complex") multipole waves, in Taylor's convention.
548 return_pq : bool NOT IMPLEMENTED
549 Return also the multipole decomposition coefficients of the field incoming to each
550 particle (inc. the field scattered from other particles.
551 return_xy : bool
552 Return also the cartesian x, y positions of the particles.
553 watch_time : bool
554 Inform about the progress on stderr
556 Returns
557 -------
558 ab : (nelem, xN, yN, 2, nelem)-shaped complex array
559 The a (electric wave), b (magnetic wave) coefficients of the outgoing field for each particle.
560 If none of return_pq or return_xy is set, the array is not enclosed in a tuple.
561 pq : (nelem, xN, yN, 2, nelem)-shaped complex array NOT IMPLEMENTED
562 The p (electric wave), q (magnetic wave) coefficients of the total exciting field
563 for each particle (including the field scattered from other particles)
564 x, y : (xN, yN)-shaped real array
565 The x,y positions of the nanoparticles.
567 if (watch_time):
568 timec = time.time()
569 print('%.4f: running scatter_plane_wave_rectarray' % timec, file = sys.stderr)
570 sys.stderr.flush()
571 nelem = TMatrices.shape[-1]
572 if ((nelem != TMatrices.shape[-3]) or (2 != TMatrices.shape[-2]) or (2 != TMatrices.shape[-4])):
573 raise ValueError('The T-matrices must be of shape (N, 2, nelem, 2, nelem) but are of shape %s' % (str(TMatrices.shape),))
574 lMax = nelem2lMax(nelem)
575 if not lMax:
576 raise ValueError('The "nelem" dimension of T-matrix has invalid value (%d).' % nelem)
577 if (watch_time):
578 print('xN = %d, yN = %d, lMax = %d' % (xN, yN, lMax), file = sys.stderr)
579 sys.stderr.flush()
580 # TODO perhaps more checks.
581 k_out = omega * math.sqrt(epsilon_b) / c # wave number
582 my, ny = get_mn_y(lMax)
583 N = yN * xN
585 J_scat=3
586 J_ext=1
588 # Do something with this ugly indexing crap
589 xind, yind = np.meshgrid(np.arange(xN),np.arange(yN), indexing='ij')
590 xind = xind.flatten()
591 yind = yind.flatten()
592 xyind = np.stack((xind, yind, np.zeros((xind.shape),dtype=int)),axis=-1)
593 cart_lattice=xyind * np.array([xd, yd, 0])
594 x=cart_lattice[:,0]
595 y=cart_lattice[:,1]
596 xyind = xyind[:,0:2]
598 # Lattice speedup
599 if (watch_time):
600 timec = time.time()
601 print('%.4f: calculating the %d translation matrix elements' % (timec, 8*nelem*nelem*xN*yN), file = sys.stderr)
602 sys.stderr.flush()
603 Agrid = np.zeros((nelem, 2*xN, 2*yN, nelem),dtype=np.complex_)
604 Bgrid = np.zeros((nelem, 2*xN, 2*yN, nelem),dtype=np.complex_)
605 for yl in range(nelem): # source
606 for xij in range(2*xN):
607 for yij in range(2*yN):
608 for yj in range(nelem): #dest
609 if((yij != yN) or (xij != xN)):
610 d_l2j = cart2sph(np.array([(xij-xN)*xd, (yij-yN)*yd, 0]))
611 Agrid[yj, xij, yij, yl] =(my[yj],ny[yj],my[yl],ny[yl],kdlj=d_l2j[0]*k_out,θlj=d_l2j[1],φlj=d_l2j[2],r_ge_d=False,J=J_scat)
612 Bgrid[yj, xij, yij, yl] =(my[yj],ny[yj],my[yl],ny[yl],kdlj=d_l2j[0]*k_out,θlj=d_l2j[1],φlj=d_l2j[2],r_ge_d=False,J=J_scat)
614 # Translation coefficient matrix T
615 if (watch_time):
616 timecold = timec
617 timec = time.time()
618 print('%4f: translation matrix elements calculated (elapsed %.2f s), filling the matrix'
619 % (timec, timec-timecold), file = sys.stderr)
620 sys.stderr.flush()
621 transmat = np.zeros((xN* yN, 2, nelem, xN* yN, 2, nelem),dtype=np.complex_)
622 for l in range(N):
623 xil, yil = xyind[l]
624 for j in range(N):
625 xij, yij = xyind[j]
626 if (l!=j):
627 transmat[j,0,:,l,0,:] = Agrid[:, xij - xil + xN, yij - yil + yN, :]
628 transmat[j,0,:,l,1,:] = Bgrid[:, xij - xil + xN, yij - yil + yN, :]
629 transmat[j,1,:,l,0,:] = Bgrid[:, xij - xil + xN, yij - yil + yN, :]
630 transmat[j,1,:,l,1,:] = Agrid[:, xij - xil + xN, yij - yil + yN, :]
631 Agrid = None
632 Bgrid = None
633 if (watch_time):
634 timecold = timec
635 timec = time.time()
636 print('%4f: translation matrix filled (elapsed %.2f s), building the interaction matrix'
637 % (timec, timec-timecold), file=sys.stderr)
638 sys.stderr.flush()
640 # Now we solve a linear problem (1 - M T) A = M P_0 where M is the T-matrix :-)
641 MT = np.empty((N,2,nelem,N,2,nelem),dtype=np.complex_)
643 TMatrices = np.broadcast_to(TMatrices, (xN, yN, 2, nelem, 2, nelem))
644 for j in range(N): # I wonder how this can be done without this loop...
645 xij, yij = xyind[j]
646 MT[j] = np.tensordot(TMatrices[xij, yij],transmat[j],axes=([-2,-1],[0,1]))
647 MT.shape = (N*2*nelem, N*2*nelem)
648 leftmatrix = np.identity(N*2*nelem) - MT
649 MT = None
650 if (watch_time):
651 timecold = timec
652 timec = time.time()
653 print('%.4f: interaction matrix complete (elapsed %.2f s)' % (timec, timec-timecold),
654 file=sys.stderr)
655 sys.stderr.flush()
657 # А ну, чики-брики и в дамки!
658 if watch_time:
659 timecold = time.time()
660 print('%.4f: factorizing the interaction matrix' % timecold, file=sys.stderr)
661 sys.stderr.flush()
662 lupiv = scipy.linalg.lu_factor(leftmatrix, overwrite_a=True)
663 leftmatrix = None
664 if watch_time:
665 timec = time.time()
666 print('%.4f: factorization complete (elapsed %.2f s)' % (timec, timec-timecold),
667 file = sys.stderr)
668 print('%.4f: solving the scattering problem for %d incoming multipoles' % (timec, nelem*2),
669 file=sys.stderr)
670 sys.stderr.flush()
671 timecold = timec
673 if(pq_0_c == 1):
674 pq_0_c = np.full((2,nelem),1)
675 ab = np.empty((2,nelem,N*2*nelem), dtype=complex)
676 for N_or_M in range(2):
677 for yy in range(nelem):
678 pq_0 = np.zeros((2,nelem), dtype=np.complex_)
679 pq_0[N_or_M,yy] = pq_0_c[N_or_M,yy]
680 pq_0 = np.broadcast_to(pq_0, (N, 2, nelem))
681 MP_0 = np.empty((N,2,nelem),dtype=np.complex_)
682 for j in range(N): # I wonder how this can be done without this loop...
683 xij, yij = xyind[j]
684 MP_0[j] = np.tensordot(TMatrices[xij, yij],pq_0[j],axes=([-2,-1],[-2,-1]))
685 MP_0.shape = (N*2*nelem,)
687 ab[N_or_M, yy] = scipy.linalg.lu_solve(lupiv, MP_0)
688 ab.shape = (2,nelem, xN, yN, 2, nelem)
689 if watch_time:
690 timec = time.time()
691 print('%.4f: done (elapsed %.2f s)' % (timec, timec-timecold),file = sys.stderr)
692 sys.stderr.flush()
693 if not (return_pq + return_xy):
694 return ab
695 returnlist = [ab]
696 if (return_pq):
697 warnings.warn("return_pq not implemented, ignoring")
698 # returnlist.append(pq_arr)
699 if (return_xy):
700 returnlist.append(x)
701 returnlist.append(y)
702 return tuple(returnlist)
716 # -------------- hexagonal lattice translation coefficients ---------------------------
717 # Implementation using (C) one-by-one AB-coefficient calculation ufunc
718 def hexlattice_precalc_AB_save2(file, lMax, k_hexside, maxlayer, circular=True, savepointinfo = False, J_scat=3):
719 params = {
720 'lMax' : lMax,
721 'k_hexside' : k_hexside,
722 'maxlayer' : maxlayer,
723 'circular' : circular,
724 'savepointinfo' : savepointinfo,
725 'J_scat' : J_scat
727 tpdict = generate_trianglepoints(maxlayer, v3d=True, circular=circular, sixthindices=True, mirrorindices=True)
728 tphcdict = generate_trianglepoints_hexcomplement(maxlayer, v3d=True, circular=circular, thirdindices=True, mirrorindices=True)
729 my, ny = get_mn_y(lMax)
730 nelem = len(my)
731 a_self_nm = np.empty((tpdict['nmi'].shape[0],nelem,nelem), dtype=complex)
732 b_self_nm = np.empty((tpdict['nmi'].shape[0],nelem,nelem), dtype=complex)
733 a_self_m0 = np.empty((tpdict['mi'].shape[1],nelem,nelem), dtype=complex)
734 b_self_m0 = np.empty((tpdict['mi'].shape[1],nelem,nelem), dtype=complex)
735 a_d2u_nm = np.empty((tphcdict['nmi'].shape[0],nelem,nelem), dtype=complex)
736 b_d2u_nm = np.empty((tphcdict['nmi'].shape[0],nelem,nelem), dtype=complex)
737 a_d2u_m0 = np.empty((tphcdict['mi'].shape[1],nelem,nelem), dtype=complex)
738 b_d2u_m0 = np.empty((tphcdict['mi'].shape[1],nelem,nelem), dtype=complex)
740 k_0 = k_hexside*_s3 # not really a wave vector here because of the normalisation!
741 tc = trans_calculator(lMax)
743 y = np.arange(nelem)
745 points = tpdict['points'][tpdict['nmi']]
746 d_i2j = cart2sph(points)
747 a_self_nm, b_self_nm = tc.get_AB(my[nx,:,nx],ny[nx,:,nx],my[nx,nx,:],ny[nx,nx,:],k_0*d_i2j[:,nx,nx,0],d_i2j[:,nx,nx,1],d_i2j[:,nx,nx,2],False,J_scat)
749 points = tpdict['points'][tpdict['mi'][0]]
750 d_i2j = cart2sph(points)
751 a_self_m0, b_self_m0 = tc.get_AB(my[nx,:,nx],ny[nx,:,nx],my[nx,nx,:],ny[nx,nx,:],k_0*d_i2j[:,nx,nx,0],d_i2j[:,nx,nx,1],d_i2j[:,nx,nx,2],False,J_scat)
753 points = tphcdict['points'][tphcdict['nmi']]
754 d_i2j = cart2sph(points)
755 a_d2u_nm, b_d2u_nm = tc.get_AB(my[nx,:,nx],ny[nx,:,nx],my[nx,nx,:],ny[nx,nx,:],k_0*d_i2j[:,nx,nx,0],d_i2j[:,nx,nx,1],d_i2j[:,nx,nx,2],False,J_scat)
757 points = tphcdict['points'][tphcdict['mi'][0]]
758 d_i2j = cart2sph(points)
759 a_d2u_m0, b_d2u_m0 = tc.get_AB(my[nx,:,nx],ny[nx,:,nx],my[nx,nx,:],ny[nx,nx,:],k_0*d_i2j[:,nx,nx,0],d_i2j[:,nx,nx,1],d_i2j[:,nx,nx,2],False,J_scat)
761 tosave = {
762 'a_self_nm' : a_self_nm,
763 'a_self_m0' : a_self_m0,
764 'b_self_nm' : b_self_nm,
765 'b_self_m0' : b_self_m0,
766 'a_d2u_nm' : a_d2u_nm,
767 'a_d2u_m0' : a_d2u_m0,
768 'b_d2u_nm' : b_d2u_nm,
769 'b_d2u_m0' : b_d2u_m0,
770 'precalc_params' : params
772 if savepointinfo:
773 tosave['tp_points'] = tpdict['points'],
774 tosave['tp_si'] = tpdict['si'],
775 tosave['tp_mi'] = tpdict['mi'],
776 tosave['tp_nmi'] = tpdict['nmi']
777 tosave['tphc_points'] = tphcdict['points'],
778 tosave['tphc_ti'] = tphcdict['ti'],
779 tosave['tphc_mi'] = tphcdict['mi'],
780 tosave['tphc_nmi'] = tphcdict['nmi']
781 np.savez(file, **tosave)
785 # The oldest implementation, using the super-inefficient pure python translation coefficients
786 def hexlattice_precalc_AB_save_purepy(file, lMax, k_hexside, maxlayer, circular=True, savepointinfo = False, J_scat=3):
787 params = {
788 'lMax' : lMax,
789 'k_hexside' : k_hexside,
790 'maxlayer' : maxlayer,
791 'circular' : circular,
792 'savepointinfo' : savepointinfo,
793 'J_scat' : J_scat
795 tpdict = generate_trianglepoints(maxlayer, v3d=True, circular=circular, sixthindices=True, mirrorindices=True)
796 tphcdict = generate_trianglepoints_hexcomplement(maxlayer, v3d=True, circular=circular, thirdindices=True, mirrorindices=True)
797 my, ny = get_mn_y(lMax)
798 nelem = len(my)
799 a_self_nm = np.empty((tpdict['nmi'].shape[0],nelem,nelem), dtype=complex)
800 b_self_nm = np.empty((tpdict['nmi'].shape[0],nelem,nelem), dtype=complex)
801 a_self_m0 = np.empty((tpdict['mi'].shape[1],nelem,nelem), dtype=complex)
802 b_self_m0 = np.empty((tpdict['mi'].shape[1],nelem,nelem), dtype=complex)
803 a_d2u_nm = np.empty((tphcdict['nmi'].shape[0],nelem,nelem), dtype=complex)
804 b_d2u_nm = np.empty((tphcdict['nmi'].shape[0],nelem,nelem), dtype=complex)
805 a_d2u_m0 = np.empty((tphcdict['mi'].shape[1],nelem,nelem), dtype=complex)
806 b_d2u_m0 = np.empty((tphcdict['mi'].shape[1],nelem,nelem), dtype=complex)
808 k_0 = k_hexside*_s3 # not really a wave vector here because of the normalisation!
810 points = tpdict['points'][tpdict['nmi']]
811 for j in range(points.shape[0]):
812 d_i2j = cart2sph(points[j])
813 for yi in range(nelem):
814 for yj in range(nelem):
815 a_self_nm[j, yj, yi] =(my[yj],ny[yj],my[yi],ny[yi],kdlj=d_i2j[0]*k_0,θlj=d_i2j[1],φlj=d_i2j[2],r_ge_d=False,J=J_scat)
816 b_self_nm[j, yj, yi] =(my[yj],ny[yj],my[yi],ny[yi],kdlj=d_i2j[0]*k_0,θlj=d_i2j[1],φlj=d_i2j[2],r_ge_d=False,J=J_scat)
817 points = tpdict['points'][tpdict['mi'][0]]
818 for j in range(points.shape[0]):
819 d_i2j = cart2sph(points[j])
820 for yi in range(nelem):
821 for yj in range(nelem):
822 a_self_m0[j, yj, yi] =(my[yj],ny[yj],my[yi],ny[yi],kdlj=d_i2j[0]*k_0,θlj=d_i2j[1],φlj=d_i2j[2],r_ge_d=False,J=J_scat)
823 b_self_m0[j, yj, yi] =(my[yj],ny[yj],my[yi],ny[yi],kdlj=d_i2j[0]*k_0,θlj=d_i2j[1],φlj=d_i2j[2],r_ge_d=False,J=J_scat)
824 points = tphcdict['points'][tphcdict['nmi']]
825 for j in range(points.shape[0]):
826 d_i2j = cart2sph(points[j])
827 for yi in range(nelem):
828 for yj in range(nelem):
829 a_d2u_nm[j, yj, yi] =(my[yj],ny[yj],my[yi],ny[yi],kdlj=d_i2j[0]*k_0,θlj=d_i2j[1],φlj=d_i2j[2],r_ge_d=False,J=J_scat)
830 b_d2u_nm[j, yj, yi] =(my[yj],ny[yj],my[yi],ny[yi],kdlj=d_i2j[0]*k_0,θlj=d_i2j[1],φlj=d_i2j[2],r_ge_d=False,J=J_scat)
831 points = tphcdict['points'][tphcdict['mi'][0]]
832 for j in range(points.shape[0]):
833 d_i2j = cart2sph(points[j])
834 for yi in range(nelem):
835 for yj in range(nelem):
836 a_d2u_m0[j, yj, yi] =(my[yj],ny[yj],my[yi],ny[yi],kdlj=d_i2j[0]*k_0,θlj=d_i2j[1],φlj=d_i2j[2],r_ge_d=False,J=J_scat)
837 b_d2u_m0[j, yj, yi] =(my[yj],ny[yj],my[yi],ny[yi],kdlj=d_i2j[0]*k_0,θlj=d_i2j[1],φlj=d_i2j[2],r_ge_d=False,J=J_scat)
838 tosave = {
839 'a_self_nm' : a_self_nm,
840 'a_self_m0' : a_self_m0,
841 'b_self_nm' : b_self_nm,
842 'b_self_m0' : b_self_m0,
843 'a_d2u_nm' : a_d2u_nm,
844 'a_d2u_m0' : a_d2u_m0,
845 'b_d2u_nm' : b_d2u_nm,
846 'b_d2u_m0' : b_d2u_m0,
847 'precalc_params' : params
849 if savepointinfo:
850 tosave['tp_points'] = tpdict['points'],
851 tosave['tp_si'] = tpdict['si'],
852 tosave['tp_mi'] = tpdict['mi'],
853 tosave['tp_nmi'] = tpdict['nmi']
854 tosave['tphc_points'] = tphcdict['points'],
855 tosave['tphc_ti'] = tphcdict['ti'],
856 tosave['tphc_mi'] = tphcdict['mi'],
857 tosave['tphc_nmi'] = tphcdict['nmi']
858 np.savez(file, **tosave)
860 def hexlattice_precalc_AB_loadunwrap(file, tpdict = None, tphcdict = None, return_points = False):
861 npz = np.load(file)
862 precalc_params = npz['precalc_params'][()]
863 my, ny = get_mn_y(precalc_params['lMax'])
864 nelem = len(my)
865 # this I should have made more universal...
866 if precalc_params['savepointinfo']:
867 if not tpdict:
868 tpdict = {
869 'points' : npz['tp_points'],
870 'si' : npz['tp_si'],
871 'mi' : npz['tp_mi'],
872 'nmi' : npz['tp_nmi'],
874 if not tphcdict:
875 tphcdict = {
876 'points' : npz['tphc_points'],
877 'ti' : npz['tphc_ti'],
878 'mi' : npz['tphc_mi'],
879 'nmi' : npz['tphc_nmi']
881 else:
882 if not tpdict:
883 tpdict = generate_trianglepoints(maxlayer = precalc_params['maxlayer'], v3d=True,
884 circular=precalc_params['circular'], sixthindices=True, mirrorindices=True)
885 if not tphcdict:
886 tphcdict = generate_trianglepoints_hexcomplement(maxlayer=precalc_params['maxlayer'], v3d=True,
887 circular=precalc_params['circular'], thirdindices=True, mirrorindices=True)
889 # For some obscure reason, I keep getting trailing single-dimension in the beginning for these arrays
890 for a in (tpdict['points'], tphcdict['points'], tpdict['si'], tpdict['mi'],
891 tphcdict['ti'], tphcdict['mi']):
892 if len(a.shape) > 2:
893 a.shape = a.shape[1::]
895 self_tr = tpdict['points']
896 d2u_tr = tphcdict['points']
897 if len(self_tr.shape)>2:
898 self_tr = np.reshape(self_tr, self_tr.shape[1::])
899 if len(d2u_tr.shape)>2:
900 d2u_tr = np.reshape(d2u_tr, d2u_tr.shape[1::])
901 u2d_tr = -d2u_tr
902 a_self = np.empty((self_tr.shape[0],nelem,nelem), dtype=complex)
903 b_self = np.empty((self_tr.shape[0],nelem,nelem), dtype=complex)
904 a_d2u = np.empty(( d2u_tr.shape[0],nelem,nelem), dtype=complex)
905 b_d2u = np.empty(( d2u_tr.shape[0],nelem,nelem), dtype=complex)
906 a_self[tpdict['nmi']]=npz['a_self_nm']
907 a_self[tpdict['mi'][0]]=npz['a_self_m0']
908 b_self[tpdict['nmi']]=npz['b_self_nm']
909 b_self[tpdict['mi'][0]]=npz['b_self_m0']
910 mirrorangles = cart2sph(self_tr[tpdict['mi'][1]])[:,2] - cart2sph(self_tr[tpdict['mi'][0]])[:,2]
911 a_self[tpdict['mi'][1],:,:] = a_self[tpdict['mi'][0],:,:] * np.exp(1j*mirrorangles[:,nx,nx]*(my[nx,nx,:]-my[nx,:,nx]))
912 b_self[tpdict['mi'][1],:,:] = b_self[tpdict['mi'][0],:,:] * np.exp(1j*mirrorangles[:,nx,nx]*(my[nx,nx,:]-my[nx,:,nx]))
913 for i in range(1,6):
914 a_self[tpdict['si'][i],:,:] = a_self[tpdict['si'][0],:,:] * np.exp(1j*math.pi/3*i*(my[nx,:]-my[:,nx]))
915 b_self[tpdict['si'][i],:,:] = b_self[tpdict['si'][0],:,:] * np.exp(1j*math.pi/3*i*(my[nx,:]-my[:,nx]))
916 a_d2u[tphcdict['nmi']]=npz['a_d2u_nm']
917 a_d2u[tphcdict['mi'][0]]=npz['a_d2u_m0']
918 b_d2u[tphcdict['nmi']]=npz['b_d2u_nm']
919 b_d2u[tphcdict['mi'][0]]=npz['b_d2u_m0']
920 mirrorangles = cart2sph(self_tr[tphcdict['mi'][1]])[:,2] - cart2sph(self_tr[tphcdict['mi'][0]])[:,2]
921 a_d2u[tphcdict['mi'][1],:,:] = a_d2u[tphcdict['mi'][0],:,:] * np.exp(1j*mirrorangles[:,nx,nx]*(my[nx,nx,:]-my[nx,:,nx]))
922 b_d2u[tphcdict['mi'][1],:,:] = b_d2u[tphcdict['mi'][0],:,:] * np.exp(1j*mirrorangles[:,nx,nx]*(my[nx,nx,:]-my[nx,:,nx]))
923 for i in (1,-1):
924 a_d2u[tphcdict['ti'][i],:,:] = a_d2u[tphcdict['ti'][0],:,:] * np.exp(i*2j*math.pi/3*(my[nx,:]-my[:,nx]))
925 b_d2u[tphcdict['ti'][i],:,:] = b_d2u[tphcdict['ti'][0],:,:] * np.exp(i*2j*math.pi/3*(my[nx,:]-my[:,nx]))
926 a_u2d = a_d2u * (-1)**(my[nx,:]-my[:,nx])
927 b_u2d = b_d2u * (-1)**(my[nx,:]-my[:,nx])
928 d = {
929 'a_self' : a_self,
930 'b_self' : b_self,
931 'a_d2u' : a_d2u,
932 'b_d2u' : b_d2u,
933 'a_u2d' : a_u2d,
934 'b_u2d' : b_u2d,
936 for k in precalc_params.keys():
937 d[k] = precalc_params[k]
938 if return_points:
939 d['d2u_tr'] = tphcdict['points']
940 d['u2d_tr'] = -tphcdict['points']
941 d['self_tr'] = tpdict['points']
942 return d
944 def hexlattice_get_AB(lMax, k_hexside, maxlayer, circular=True, return_points = True, J_scat=3):
945 params = {
946 'lMax' : lMax,
947 'k_hexside' : k_hexside,
948 'maxlayer' : maxlayer,
949 'circular' : circular,
950 'savepointinfo' : return_points, # should I delete this key?
951 'J_scat' : J_scat
953 tpdict = generate_trianglepoints(maxlayer, v3d=True, circular=circular, sixthindices=True, mirrorindices=True)
954 tphcdict = generate_trianglepoints_hexcomplement(maxlayer, v3d=True, circular=circular, thirdindices=True, mirrorindices=True)
955 my, ny = get_mn_y(lMax)
956 nelem = len(my)
957 a_self_nm = np.empty((tpdict['nmi'].shape[0],nelem,nelem), dtype=complex)
958 b_self_nm = np.empty((tpdict['nmi'].shape[0],nelem,nelem), dtype=complex)
959 a_self_m0 = np.empty((tpdict['mi'].shape[1],nelem,nelem), dtype=complex)
960 b_self_m0 = np.empty((tpdict['mi'].shape[1],nelem,nelem), dtype=complex)
961 a_d2u_nm = np.empty((tphcdict['nmi'].shape[0],nelem,nelem), dtype=complex)
962 b_d2u_nm = np.empty((tphcdict['nmi'].shape[0],nelem,nelem), dtype=complex)
963 a_d2u_m0 = np.empty((tphcdict['mi'].shape[1],nelem,nelem), dtype=complex)
964 b_d2u_m0 = np.empty((tphcdict['mi'].shape[1],nelem,nelem), dtype=complex)
966 k_0 = k_hexside*_s3 # not really a wave vector here because of the normalisation!
967 tc = trans_calculator(lMax)
969 y = np.arange(nelem)
971 points = tpdict['points'][tpdict['nmi']]
972 d_i2j = cart2sph(points)
973 a_self_nm, b_self_nm = tc.get_AB_arrays(k_0*d_i2j[:,0],d_i2j[:,1],d_i2j[:,2],np.array([False]),J_scat)
975 points = tpdict['points'][tpdict['mi'][0]]
976 d_i2j = cart2sph(points)
977 a_self_m0, b_self_m0 = tc.get_AB_arrays(k_0*d_i2j[:,0],d_i2j[:,1],d_i2j[:,2],np.array([False]),J_scat)
979 points = tphcdict['points'][tphcdict['nmi']]
980 d_i2j = cart2sph(points)
981 a_d2u_nm, b_d2u_nm = tc.get_AB_arrays(k_0*d_i2j[:,0],d_i2j[:,1],d_i2j[:,2],np.array([False]),J_scat)
983 points = tphcdict['points'][tphcdict['mi'][0]]
984 d_i2j = cart2sph(points)
985 a_d2u_m0, b_d2u_m0 = tc.get_AB_arrays(k_0*d_i2j[:,0],d_i2j[:,1],d_i2j[:,2],np.array([False]),J_scat)
987 tosave = {
988 'a_self_nm' : a_self_nm,
989 'a_self_m0' : a_self_m0,
990 'b_self_nm' : b_self_nm,
991 'b_self_m0' : b_self_m0,
992 'a_d2u_nm' : a_d2u_nm,
993 'a_d2u_m0' : a_d2u_m0,
994 'b_d2u_nm' : b_d2u_nm,
995 'b_d2u_m0' : b_d2u_m0,
996 'precalc_params' : params
998 if savepointinfo:
999 tosave['tp_points'] = tpdict['points'],
1000 tosave['tp_si'] = tpdict['si'],
1001 tosave['tp_mi'] = tpdict['mi'],
1002 tosave['tp_nmi'] = tpdict['nmi']
1003 tosave['tphc_points'] = tphcdict['points'],
1004 tosave['tphc_ti'] = tphcdict['ti'],
1005 tosave['tphc_mi'] = tphcdict['mi'],
1006 tosave['tphc_nmi'] = tphcdict['nmi']
1007 np.savez(file, **tosave)
1009 self_tr = tpdict['points']
1010 d2u_tr = tphcdict['points']
1011 if len(self_tr.shape)>2:
1012 self_tr = np.reshape(self_tr, self_tr.shape[1::])
1013 if len(d2u_tr.shape)>2:
1014 d2u_tr = np.reshape(d2u_tr, d2u_tr.shape[1::])
1015 u2d_tr = -d2u_tr
1016 a_self = np.empty((self_tr.shape[0],nelem,nelem), dtype=complex)
1017 b_self = np.empty((self_tr.shape[0],nelem,nelem), dtype=complex)
1018 a_d2u = np.empty(( d2u_tr.shape[0],nelem,nelem), dtype=complex)
1019 b_d2u = np.empty(( d2u_tr.shape[0],nelem,nelem), dtype=complex)
1020 a_self[tpdict['nmi']]=a_self_nm
1021 a_self[tpdict['mi'][0]]=a_self_m0
1022 b_self[tpdict['nmi']]=b_self_nm
1023 b_self[tpdict['mi'][0]]=b_self_m0
1024 mirrorangles = cart2sph(self_tr[tpdict['mi'][1]])[:,2] - cart2sph(self_tr[tpdict['mi'][0]])[:,2]
1025 a_self[tpdict['mi'][1],:,:] = a_self[tpdict['mi'][0],:,:] * np.exp(1j*mirrorangles[:,nx,nx]*(my[nx,nx,:]-my[nx,:,nx]))
1026 b_self[tpdict['mi'][1],:,:] = b_self[tpdict['mi'][0],:,:] * np.exp(1j*mirrorangles[:,nx,nx]*(my[nx,nx,:]-my[nx,:,nx]))
1027 for i in range(1,6):
1028 a_self[tpdict['si'][i],:,:] = a_self[tpdict['si'][0],:,:] * np.exp(1j*math.pi/3*i*(my[nx,:]-my[:,nx]))
1029 b_self[tpdict['si'][i],:,:] = b_self[tpdict['si'][0],:,:] * np.exp(1j*math.pi/3*i*(my[nx,:]-my[:,nx]))
1030 a_d2u[tphcdict['nmi']]=a_d2u_nm
1031 a_d2u[tphcdict['mi'][0]]=a_d2u_m0
1032 b_d2u[tphcdict['nmi']]=b_d2u_nm
1033 b_d2u[tphcdict['mi'][0]]=b_d2u_m0
1034 mirrorangles = cart2sph(self_tr[tphcdict['mi'][1]])[:,2] - cart2sph(self_tr[tphcdict['mi'][0]])[:,2]
1035 a_d2u[tphcdict['mi'][1],:,:] = a_d2u[tphcdict['mi'][0],:,:] * np.exp(1j*mirrorangles[:,nx,nx]*(my[nx,nx,:]-my[nx,:,nx]))
1036 b_d2u[tphcdict['mi'][1],:,:] = b_d2u[tphcdict['mi'][0],:,:] * np.exp(1j*mirrorangles[:,nx,nx]*(my[nx,nx,:]-my[nx,:,nx]))
1037 for i in (1,-1):
1038 a_d2u[tphcdict['ti'][i],:,:] = a_d2u[tphcdict['ti'][0],:,:] * np.exp(i*2j*math.pi/3*(my[nx,:]-my[:,nx]))
1039 b_d2u[tphcdict['ti'][i],:,:] = b_d2u[tphcdict['ti'][0],:,:] * np.exp(i*2j*math.pi/3*(my[nx,:]-my[:,nx]))
1040 a_u2d = a_d2u * (-1)**(my[nx,:]-my[:,nx])
1041 b_u2d = b_d2u * (-1)**(my[nx,:]-my[:,nx])
1042 d = {
1043 'a_self' : a_self,
1044 'b_self' : b_self,
1045 'a_d2u' : a_d2u,
1046 'b_d2u' : b_d2u,
1047 'a_u2d' : a_u2d,
1048 'b_u2d' : b_u2d,
1050 for k in params.keys():
1051 d[k] = params[k]
1052 if return_points:
1053 d['d2u_tr'] = tphcdict['points']
1054 d['u2d_tr'] = -tphcdict['points']
1055 d['self_tr'] = tpdict['points']
1056 return d