qpms/qpms_p.py

   1 """
   2 This file contains mostly legacy code, but is still kept for reference. Avoid using this.
   3 """
   4
   5 import numpy as np
   6 from .qpms_c import *
   7 ň = np.newaxis
   8 import scipy
   9 from .constants import ε_0, c, pi, π, e, ℏ, μ_0
  10 eV = e
  11 from scipy.special import lpmn, lpmv, spherical_jn, spherical_yn, poch, gammaln, factorial
  12 import math
  13 import cmath
  14 import os
  15
  16 """
  17 '''
  18 Try to import numba. Its pre-0.28.0 versions can not handle functions
  19 containing utf8 identifiers, so we keep track about that.
  20 '''
  21 try:
  22     import numba
  23     use_jit = True
  24     if numba.__version__ >= '0.28.0':
  25         use_jit_utf8 = True
  26     else:
  27         use_jit_utf8 = False
  28 except ImportError:
  29     use_jit = False
  30     use_jit_utf8 = False
  31 '''
  32 Accordingly, we define our own jit decorator that handles
  33 different versions of numba or does nothing if numba is not
  34 present. Note that functions that include unicode identifiers
  35 must be decorated with #@ujit
  36 '''
  37 #def dummywrap(f):
  38 #    return f
  39 def jit(f):
  40     if use_jit:
  41         return numba.jit(f)
  42     else:
  43         return f
  44 def ujit(f):
  45     if use_jit_utf8:
  46         return numba.jit(f)
  47     else:
  48         return f
  49 """
  50
  51 # Coordinate transforms for arrays of "arbitrary" shape
  52 #@ujit
  53 def cart2sph(cart,axis=-1, allow2d=False):
  54     cart = np.array(cart, copy=False)
  55     if cart.shape[axis] == 3:
  56         [x, y, z] = np.split(cart,3,axis=axis)
  57         r = np.linalg.norm(cart,axis=axis,keepdims=True)
  58         r_zero = np.logical_not(r)
  59         θ = np.arccos(z/(r+r_zero))
  60     elif cart.shape[axis] == 2 and allow2d:
  61         [x, y] = np.split(cart,2,axis=axis)
  62         r = np.linalg.norm(cart,axis=axis,keepdims=True)
  63         r_zero = np.logical_not(r)
  64         θ = np.broadcast_to(np.pi/2, x.shape)
  65     else:
  66         raise ValueError("The converted array has to have dimension 3 "
  67                          "(or 2 if allow2d==True)"
  68                          " along the given axis, not %d" % cart.shape[axis])
  69
  70     φ = np.arctan2(y,x) # arctan2 handles zeroes correctly itself
  71     return np.concatenate((r,θ,φ),axis=axis)
  72
  73 #@ujit
  74 def sph2cart(sph, axis=-1):
  75     sph = np.array(sph, copy=False)
  76     if (sph.shape[axis] != 3):
  77         raise ValueError("The converted array has to have dimension 3"
  78                          " along the given axis")
  79     [r,θ,φ] = np.split(sph,3,axis=axis)
  80     sinθ = np.sin(θ)
  81     x = r * sinθ * np.cos(φ)
  82     y = r * sinθ * np.sin(φ)
  83     z = r * np.cos(θ)
  84     return np.concatenate((x,y,z),axis=axis)
  85
  86 #@ujit
  87 def sph_loccart2cart(loccart, sph, axis=-1):
  88     """
  89     Transformation of vector specified in local orthogonal coordinates
  90     (tangential to spherical coordinates – basis r̂,θ̂,φ̂) to global cartesian
  91     coordinates (basis x̂,ŷ,ẑ)
  92     SLOW FOR SMALL ARRAYS
  93
  94     Parameters
  95     ----------
  96     loccart: ... TODO
  97         the transformed vector in the local orthogonal coordinates
  98
  99     sph: ... TODO
 100         the point (in spherical coordinates) at which the locally
 101         orthogonal basis is evaluated
 102
 103     Returns
 104     -------
 105     output: ... TODO
 106         The coordinates of the vector in global cartesian coordinates
 107     """
 108     loccart = np.array(loccart, copy=False)
 109     sph = np.array(sph, copy=False)
 110     if (loccart.shape[axis] != 3):
 111         raise ValueError("The converted array has to have dimension 3"
 112                          " along the given axis")
 113     [r,θ,φ] = np.split(sph,3,axis=axis)
 114     sinθ = np.sin(θ)
 115     cosθ = np.cos(θ)
 116     sinφ = np.sin(φ)
 117     cosφ = np.cos(φ)
 118
 119     #x = r * sinθ * cosφ
 120     #y = r * sinθ * sinφ
 121     #z = r * cosθ
 122     r̂x = sinθ * cosφ
 123     r̂y = sinθ * sinφ
 124     r̂z = cosθ
 125     θ̂x = cosθ * cosφ
 126     θ̂y = cosθ * sinφ
 127     θ̂z = -sinθ
 128     φ̂x = -sinφ
 129     φ̂y = cosφ
 130     φ̂z = np.zeros(φ̂y.shape)
 131     r̂ = np.concatenate((r̂x,r̂y,r̂z),axis=axis)
 132     θ̂ = np.concatenate((θ̂x,θ̂y,θ̂z),axis=axis)
 133     φ̂ = np.concatenate((φ̂x,φ̂y,φ̂z),axis=axis)
 134     [inr̂,inθ̂,inφ̂] = np.split(loccart,3,axis=axis)
 135     out=inr̂*r̂+inθ̂*θ̂+inφ̂*φ̂
 136     return out
 137
 138 #@ujit
 139 def sph_loccart_basis(sph, sphaxis=-1, cartaxis=None):
 140     """
 141     Returns the local cartesian basis in terms of global cartesian basis.
 142     sphaxis refers to the original dimensions
 143     TODO doc
 144     """
 145     if(cartaxis is None):
 146         cartaxis = sph.ndim # default to last axis
 147     [r,θ,φ] = np.split(sph,3,axis=sphaxis)
 148     sinθ = np.sin(θ)
 149     cosθ = np.cos(θ)
 150     sinφ = np.sin(φ)
 151     cosφ = np.cos(φ)
 152
 153     #x = r * sinθ * cosφ
 154     #y = r * sinθ * sinφ
 155     #z = r * cosθ
 156     r̂x = sinθ * cosφ
 157     r̂y = sinθ * sinφ
 158     r̂z = cosθ
 159     θ̂x = cosθ * cosφ
 160     θ̂y = cosθ * sinφ
 161     θ̂z = -sinθ
 162     φ̂x = -sinφ
 163     φ̂y = cosφ
 164     φ̂z = np.zeros(φ̂y.shape)
 165     #r̂ = np.concatenate((r̂x,r̂y,r̂z),axis=axis)
 166     #θ̂ = np.concatenate((θ̂x,θ̂y,θ̂z),axis=axis)
 167     #φ̂ = np.concatenate((φ̂x,φ̂y,φ̂z),axis=axis)
 168     x = np.expand_dims(np.concatenate((r̂x,θ̂x,φ̂x), axis=sphaxis),axis=cartaxis)
 169     y = np.expand_dims(np.concatenate((r̂y,θ̂y,φ̂y), axis=sphaxis),axis=cartaxis)
 170     z = np.expand_dims(np.concatenate((r̂z,θ̂z,φ̂z), axis=sphaxis),axis=cartaxis)
 171     out = np.concatenate((x,y,z),axis=cartaxis)
 172     return out
 173
 174
 175 #@jit
 176 def nelem2lMax(nelem):
 177     """
 178     Auxiliary inverse function to nelem(lMax) = (lMax + 2) * lMax. Returns 0 if
 179     it nelem does not come from positive integer lMax.
 180     """
 181     lMax = round(math.sqrt(1+nelem) - 1)
 182     if ((lMax < 1) or ((lMax + 2) * lMax != nelem)):
 183         return 0
 184     else:
 185         return lMax
 186
 187 def lMax2nelem(lMax):
 188     return lMax * (lMax+2)
 189
 190 def lpy(nmax, z):
 191     # TODO TRUEVECTORIZE
 192     # TODO DOC
 193     nelem = nmax * (nmax+2)
 194     z = np.array(z, copy=False)
 195     P_y = np.empty(z.shape + (nelem,), dtype=np.float_)
 196     dP_y = np.empty(z.shape + (nelem,), dtype=np.float_)
 197     for i in np.ndindex(z.shape):
 198         P_y[i], dP_y[i] = lpy1(nmax, z[i])
 199     return (P_y, dP_y)
 200
 201 def lpy1(nmax, z):
 202     """
 203     Associated legendre function and its derivatative at z in the 'y-indexing'.
 204     (Without Condon-Shortley phase AFAIK.)
 205     NOT THOROUGHLY TESTED
 206
 207     Parameters
 208     ----------
 209
 210     nmax: int
 211         The maximum order to which the Legendre functions will be evaluated..
 212
 213     z: float
 214         The point at which the Legendre functions are evaluated.
 215
 216     output: (P_y, dP_y) TODO
 217         y-indexed legendre polynomials and their derivatives
 218
 219     """
 220     pmn_plus, dpmn_plus = lpmn(nmax, nmax, z)
 221     pmn_minus, dpmn_minus = lpmn(-nmax, nmax, z)
 222     nelem = nmax * nmax + 2*nmax
 223     P_y = np.empty((nelem,), dtype=np.float_)
 224     dP_y = np.empty((nelem,), dtype=np.float_)
 225     mn_p_y, mn_n_y = get_y_mn_unsigned(nmax)
 226     mn_plus_mask = (mn_p_y >= 0)
 227     mn_minus_mask = (mn_n_y >= 0)
 228     #print( mn_n_y[mn_minus_mask])
 229     P_y[mn_p_y[mn_plus_mask]] = pmn_plus[mn_plus_mask]
 230     P_y[mn_n_y[mn_minus_mask]] = pmn_minus[mn_minus_mask]
 231     dP_y[mn_p_y[mn_plus_mask]] = dpmn_plus[mn_plus_mask]
 232     dP_y[mn_n_y[mn_minus_mask]] = dpmn_minus[mn_minus_mask]
 233     return (P_y, dP_y)
 234
 235 def vswf_yr(pos_sph,lMax,J=1):
 236     """
 237     Normalized vector spherical wavefunctions $\widetilde{M}_{mn}^{j}$,
 238     $\widetilde{N}_{mn}^{j}$ as in [1, (2.40)].
 239
 240     Parameters
 241     ----------
 242
 243     pos_sph : np.array(dtype=float, shape=(someshape,3))
 244         The positions where the spherical vector waves are to be
 245         evaluated. The last axis corresponds to the individual
 246         points (r,θ,φ). The radial coordinate r is dimensionless,
 247         assuming that it has already been multiplied by the
 248         wavenumber.
 249
 250     nmax : int
 251         The maximum order to which the VSWFs are evaluated.
 252
 253     Returns
 254     -------
 255
 256     output : np.array(dtype=complex, shape=(someshape,nmax*nmax + 2*nmax,3))
 257         Spherical vector wave functions evaluated at pos_sph,
 258         in the local basis (r̂,θ̂,φ̂). The last indices correspond
 259         to m, n (in the ordering given by mnindex()), and basis
 260         vector index, respectively.
 261
 262     [1] Jonathan M. Taylor. Optical Binding Phenomena: Observations and
 263     Mechanisms.
 264     """
 265     # TODO TRUEVECTORIZE
 266     pos_sph = np.array(pos_sph, copy = False)
 267     nelem = (lMax + 2) * lMax
 268     M_y = np.empty(pos_sph.shape[:-1] + (nelem,3), dtype = np.complex_)
 269     N_y = np.empty(pos_sph.shape[:-1] + (nelem,3), dtype = np.complex_)
 270     for i in np.ndindex(pos_sph.shape[:-1]):
 271         M_y[i], N_y[i] = vswf_yr1(pos_sph[i], lMax, J) # non-vectorised function
 272     return (M_y, N_y)
 273
 274 #@jit
 275 def _sph_zn_1(n,z,derivative=False):
 276     return spherical_jn(n,z,derivative)
 277 #@jit
 278 def _sph_zn_2(n,z,derivative=False):
 279     return spherical_yn(n,z,derivative)
 280 #@jit
 281 def _sph_zn_3(n,z,derivative=False):
 282     besj=spherical_jn(n,z,derivative)
 283     besy=spherical_yn(n,z,derivative)
 284     return besj + 1j*besy
 285 #@jit
 286 def _sph_zn_4(n,z,derivative=False):
 287     besj=spherical_jn(n,z,derivative)
 288     besy=spherical_yn(n,z,derivative)
 289     return besj + 1j*besy
 290 _sph_zn = [_sph_zn_1,_sph_zn_2,_sph_zn_3,_sph_zn_4]
 291
 292 # computes bessel/hankel functions for orders from 0 up to n;
 293 #@jit
 294 def zJn(n, z, J=1):
 295     return (_sph_zn[J-1](n=np.arange(n+1),z=z,derivative=False), _sph_zn[J-1](n=np.arange(n+1),z=z,derivative=True))
 296
 297
 298
 299 # The following 4 funs have to be refactored, possibly merged
 300
 301 # FIXME: this can be expressed simply as:
 302 # $$ -\frac{1}{2}\sqrt{\frac{2n+1}{4\pi}n\left(n+1\right)}(\delta_{m,1}+\delta_{m,-1}) $$
 303 #@ujit
 304 def π̃_zerolim(nmax): # seems OK
 305     """
 306     lim_{θ→ 0-} π̃(cos θ)
 307     """
 308     my, ny = get_mn_y(nmax)
 309     nelems = len(my)
 310     π̃_y = np.zeros((nelems))
 311     plus1mmask = (my == 1)
 312     minus1mmask = (my == -1)
 313     pluslim = -ny*(1+ny)/2
 314     minuslim = 0.5
 315     π̃_y[plus1mmask] = pluslim[plus1mmask]
 316     π̃_y[minus1mmask] = - minuslim
 317     prenorm =  np.sqrt((2*ny + 1)*factorial(ny-my)/(4*π*factorial(ny+my)))
 318     π̃_y = prenorm *     π̃_y
 319     return π̃_y
 320
 321 #@ujit
 322 def π̃_pilim(nmax): # Taky OK, jen to možná není kompatibilní se vzorečky z mathematiky
 323     """
 324     lim_{θ→ π+} π̃(cos θ)
 325     """
 326     my, ny = get_mn_y(nmax)
 327     nelems = len(my)
 328     π̃_y = np.zeros((nelems))
 329     plus1mmask = (my == 1)
 330     minus1mmask = (my == -1)
 331     pluslim = (-1)**ny*ny*(1+ny)/2
 332     minuslim = 0.5*(-1)**ny
 333     π̃_y[plus1mmask] = pluslim[plus1mmask]
 334     π̃_y[minus1mmask] = minuslim[minus1mmask]
 335     prenorm =  np.sqrt((2*ny + 1)*factorial(ny-my)/(4*π*factorial(ny+my)))
 336     π̃_y = prenorm *     π̃_y
 337     return π̃_y
 338
 339 # FIXME: this can be expressed simply as
 340 # $$ -\frac{1}{2}\sqrt{\frac{2n+1}{4\pi}n\left(n+1\right)}(\delta_{m,1}-\delta_{m,-1}) $$
 341 #@ujit
 342 def τ̃_zerolim(nmax):
 343     """
 344     lim_{θ→ 0-} τ̃(cos θ)
 345     """
 346     p0 = π̃_zerolim(nmax)
 347     my, ny = get_mn_y(nmax)
 348     minus1mmask = (my == -1)
 349     p0[minus1mmask] = -p0[minus1mmask]
 350     return p0
 351
 352 #@ujit
 353 def τ̃_pilim(nmax):
 354     """
 355     lim_{θ→  π+} τ̃(cos θ)
 356     """
 357     t = π̃_pilim(nmax)
 358     my, ny = get_mn_y(nmax)
 359     plus1mmask = (my == 1)
 360     t[plus1mmask] = -t[plus1mmask]
 361     return t
 362
 363 #@ujit
 364 def get_π̃τ̃_y1(θ,nmax):
 365     # TODO replace with the limit functions (below) when θ approaches
 366     # the extreme values at about 1e-6 distance
 367     """
 368     (... TODO)
 369
 370     """
 371     if (abs(θ)<1e-6):
 372         return (π̃_zerolim(nmax),τ̃_zerolim(nmax))
 373     if (abs(θ-π)<1e-6):
 374         return (π̃_pilim(nmax),τ̃_pilim(nmax))
 375     my, ny = get_mn_y(nmax)
 376     nelems = len(my)
 377     Py, dPy = lpy(nmax, math.cos(θ))
 378     prenorm =  np.sqrt((2*ny + 1)*factorial(ny-my)/(4*π*factorial(ny+my)))
 379     π̃_y = prenorm * my * Py / math.sin(θ)  # bacha, možné dělení nulou
 380     τ̃_y = prenorm * dPy * (- math.sin(θ))  # TADY BACHA!!!!!!!!!! * (- math.sin(pos_sph[1])) ???
 381     return (π̃_y,τ̃_y)
 382
 383 #@ujit
 384 def vswf_yr1(pos_sph,nmax,J=1):
 385     """
 386     As vswf_yr, but evaluated only at single position (i.e. pos_sph has
 387     to have shape=(3))
 388     """
 389     if (pos_sph[1].imag or pos_sph[2].imag):
 390         raise ValueError("The angles for the spherical wave functions can not be complex")
 391     kr = pos_sph[0] if pos_sph[0].imag else pos_sph[0].real # To supress the idiotic warning in scipy.special.sph_jn
 392     θ = pos_sph[1].real
 393     φ = pos_sph[2].real
 394     my, ny = get_mn_y(nmax)
 395     Py, dPy = lpy(nmax, math.cos(θ))
 396     nelems = nmax*nmax + 2*nmax
 397     # TODO Remove these two lines in production:
 398     if(len(Py) != nelems or len(my) != nelems):
 399         raise ValueError("This is very wrong.")
 400     prenorm =  np.sqrt((2*ny + 1)*factorial(ny-my)/(4*π*factorial(ny+my)))
 401     if (abs(θ)<1e-6): # Ošetření limitního chování derivací Leg. fcí
 402         π̃_y=π̃_zerolim(nmax)
 403         τ̃_y=τ̃_zerolim(nmax)
 404     elif (abs(θ-π)<1e-6):
 405         π̃_y=π̃_pilim(nmax)
 406         τ̃_y=τ̃_pilim(nmax)
 407     else:
 408         π̃_y = prenorm * my * Py / math.sin(θ)
 409         τ̃_y = prenorm * dPy * (- math.sin(θ))  # TADY BACHA!!!!!!!!!! * (- math.sin(pos_sph[1])) ???
 410     z_n, dz_n = zJn(nmax, kr, J=J)
 411     z_y = z_n[ny]
 412     dz_y = dz_n[ny]
 413     eimf_y = np.exp(1j*my*φ) # zbytečné opakování my, lepší by bylo to spočítat jednou a vyindexovat
 414     M̃_y = np.zeros((nelems,3), dtype=np.complex_)
 415     M̃_y[:,1] = 1j * π̃_y * eimf_y * z_y
 416     M̃_y[:,2] =  -   τ̃_y * eimf_y * z_y
 417     Ñ_y = np.empty((nelems,3), dtype=np.complex_)
 418     Ñ_y[:,0] = (ny*(ny+1)/kr) * prenorm * Py * eimf_y * z_y
 419     Ñradial_fac_y = z_y / kr + dz_y
 420     Ñ_y[:,1] =    τ̃_y * eimf_y * Ñradial_fac_y
 421     Ñ_y[:,2] = 1j*π̃_y * eimf_y * Ñradial_fac_y
 422     return(M̃_y, Ñ_y)
 423
 424 #def plane_E_y(nmax):
 425 #    """
 426 #    The E_mn normalization factor as in [1, (3)] WITHOUT the E_0 factor,
 427 #    y-indexed
 428 #
 429 #    (... TODO)
 430 #
 431 #    References
 432 #    ----------
 433 #    [1] Jonathan M. Taylor. Optical Binding Phenomena: Observations and
 434 #    Mechanisms. FUCK, I MADE A MISTAKE: THIS IS FROM 7U
 435 #    """
 436 #    my, ny = get_mn_y(nmax)
 437 #    return 1j**ny * np.sqrt((2*ny+1)*factorial(ny-my) /
 438 #                            (ny*(ny+1)*factorial(ny+my))
 439 #    )
 440 #@ujit
 441 def zplane_pq_y(nmax, betap = 0):
 442     """
 443     The z-propagating plane wave expansion coefficients as in [1, (1.12)]
 444
 445     Taylor's normalization
 446
 447     (... TODO)
 448     """
 449     my, ny = get_mn_y(nmax)
 450     U_y = 4*π * 1j**ny / (ny * (ny+1))
 451     π̃_y = π̃_zerolim(nmax)
 452     τ̃_y = τ̃_zerolim(nmax)
 453
 454     # fixme co je zač ten e_θ ve vzorečku? (zde neimplementováno)
 455     p_y = U_y*(τ̃_y*math.cos(betap) - 1j*math.sin(betap)*π̃_y)
 456     q_y = U_y*(π̃_y*math.cos(betap) - 1j*math.sin(betap)*τ̃_y)
 457     return (p_y, q_y)
 458
 459
 460 #import warnings
 461 #@ujit
 462 def plane_pq_y(nmax, kdir_cart, E_cart):
 463     """
 464     The plane wave expansion coefficients for any direction kdir_cart
 465     and amplitude vector E_cart (which might be complex, depending on
 466     the phase and polarisation state). If E_cart and kdir_cart are
 467     not orthogonal, the result should correspond to the k-normal part
 468     of E_cart.
 469
 470     The Taylor's convention on the expansion is used; therefore,
 471     the electric field is expressed as
 472     E(r) = E_cart * exp(1j * k ∙ r)
 473          = -1j * ∑_y ( p_y * Ñ_y(|k| * r) + q_y * M̃_y(|k| * r))
 474     (note the -1j factor).
 475
 476     Parameters
 477     ----------
 478     nmax: int
 479         The order of the expansion.
 480     kdir_cart: (...,3)-shaped real array
 481         The wave vector (its magnitude does not play a role).
 482     E_cart: (...,3)-shaped complex array
 483         Amplitude of the plane wave
 484
 485     Returns
 486     -------
 487     p_y, q_y:
 488         The expansion coefficients for the electric (Ñ) and magnetic
 489         (M̃) waves, respectively.
 490     """
 491     # TODO TRUEVECTORIZE
 492     if np.iscomplexobj(kdir_cart):
 493         warnings.warn("The direction vector for the plane wave coefficients should be real. I am discarding the imaginary part now.")
 494         kdir_cart = kdir_cart.real
 495
 496     E_cart = np.array(E_cart, copy=False)
 497     k_sph = cart2sph(kdir_cart)
 498     my, ny = get_mn_y(nmax)
 499     nelem = len(my)
 500     U_y = 4*π * 1j**ny / (ny * (ny+1))
 501     θ̂ = sph_loccart2cart(np.array([0,1,0]), k_sph, axis=-1)
 502     φ̂ = sph_loccart2cart(np.array([0,0,1]), k_sph, axis=-1)
 503
 504     # not properly vectorised part:
 505     π̃_y = np.empty(k_sph.shape[:-1] + (nelem,), dtype=np.complex_)
 506     τ̃_y = np.empty(π̃_y.shape, dtype=np.complex_)
 507     for i in np.ndindex(k_sph.shape[:-1]):
 508         π̃_y[i], τ̃_y[i] = get_π̃τ̃_y1(k_sph[i][1], nmax) # this shit is not vectorised
 509
 510     # last indices of the summands: [...,y,local cartesian component index (of E)]
 511     p_y = np.sum( U_y[...,:,ň]
 512                   * np.conj(np.exp(1j*my[...,:,ň]*k_sph[...,ň,ň,2]) * (
 513                      θ̂[...,ň,:]*τ̃_y[...,:,ň] + 1j*φ̂[...,ň,:]*π̃_y[...,:,ň]))
 514                   * E_cart[...,ň,:],
 515               axis=-1)
 516     q_y = np.sum( U_y[...,:,ň]
 517                   * np.conj(np.exp(1j*my[...,:,ň]*k_sph[...,ň,ň,2]) * (
 518                      θ̂[...,ň,:]*π̃_y[...,:,ň] + 1j*φ̂[...,ň,:]*τ̃_y[...,:,ň]))
 519                   * E_cart[...,ň,:],
 520           axis=-1)
 521     return (p_y, q_y)
 522
 523
 524
 525
 526 # Material parameters
 527 #@ujit
 528 def ε_drude(ε_inf, ω_p, γ_p, ω): # RELATIVE permittivity, of course
 529     return ε_inf - ω_p*ω_p/(ω*(ω+1j*γ_p))
 530
 531
 532 # In[8]:
 533
 534 # Mie scattering
 535 #@ujit
 536 def mie_coefficients(a, nmax,  #ω, ε_i, ε_e=1, J_ext=1, J_scat=3
 537                                k_i, k_e, μ_i=1, μ_e=1, J_ext=1, J_scat=3):
 538     """
 539
 540     FIXME test the magnetic case
 541     TODO description
 542     RH concerns the N ("electric") part, RV the M ("magnetic") part
 543     #
 544
 545     Parameters
 546     ----------
 547     a : float
 548         Diameter of the sphere.
 549
 550     nmax : int
 551         To which order (inc. nmax) to compute the coefficients.
 552
 553     ω : float
 554         Frequency of the radiation
 555
 556     ε_i, ε_e, μ_i, μ_e : complex
 557         Relative permittivities and permeabilities of the sphere (_i)
 558         and the environment (_e)
 559         MAGNETIC (μ_i, μ_e != 1)  CASE UNTESTED AND PROBABLY BUGGY
 560
 561     J_ext, J_scat : 1, 2, 3, or 4 (must be different)
 562         Specifies the species of the Bessel/Hankel functions in which
 563         the external incoming (J_ext) and scattered (J_scat) fields
 564         are represented. 1,2,3,4 correspond to j,y,h(1),h(2), respectively.
 565         The returned coefficients are always with respect to the decomposition
 566         of the "external incoming" wave.
 567
 568     Returns
 569     -------
 570     RV == a/p, RH == b/q, TV = d/p, TH = c/q
 571     TODO
 572     what does it return on index 0???
 573     FIXME permeabilities
 574     """
 575     # permittivities are relative!
 576     # cf. worknotes
 577     #print("a, nmax, ε_m, ε_b, ω",a, nmax, ε_m, ε_b, ω)
 578     #k_i = cmath.sqrt(ε_i*μ_i) * ω / c
 579     x_i = k_i * a
 580     #k_e = cmath.sqrt(ε_e*μ_e) * ω / c
 581     x_e = k_e * a
 582     #print("Mie: phase at radius: x_i",x_i,"x_e",x_e)
 583     m = k_i/k_e#cmath.sqrt(ε_i*μ_i/(ε_e*μ_e))
 584     # We "need" the absolute permeabilities for the final formula
 585     # This is not the absolute wave impedance, because only their ratio
 586     # ηi/ηe is important for getting the Mie coefficients.
 587     η_inv_i = k_i / μ_i
 588     η_inv_e = k_e / μ_e
 589     #print("k_m, x_m,k_b,x_b",k_m, x_m,k_b,x_b)
 590     zi, ži = zJn(nmax, x_i, J=1)
 591     #Pi = (zi * x_i)
 592     #Di = (zi + x_i * ži) / Pi # Vzoreček Taylor (2.9)
 593     #ži = zi + x_i * ži
 594     ze, že = zJn(nmax, x_e, J=J_ext)
 595     #Pe = (ze * x_e)
 596     #De = (ze + x_e * že) / Pe # Vzoreček Taylor (2.9)
 597     #že = ze + x_e * že
 598     zs, žs = zJn(nmax, x_e, J=J_scat)
 599     #Ps = (zs * x_e)
 600     #Ds = (zs + x_e * žs) / Ps # Vzoreček Taylor (2.9)
 601     #žs = zs + x_e * zs
 602     #RH = (μ_i*zi*že - μ_e*ze*ži) / (μ_i*zi*žs - μ_e*zs*ži)
 603     #RV = (μ_e*m*m*zi*že - μ_i*ze*ži) / (μ_e*m*m*zi*žs - μ_i*zs*ži)
 604     #TH = (μ_i*ze*žs - μ_i*zs*že) / (μ_i*zi*žs - μ_e*zs*ži)
 605     #TV = (μ_i*m*ze*žs - μ_i*m*zs*že) / (μ_e*m*m*zi*žs - μ_i*zs*ži)
 606     ži = zi/x_i+ži
 607     žs = zs/x_e+žs
 608     že = ze/x_e+že
 609     RV = -((-η_inv_i * že * zi + η_inv_e * ze * ži)/(-η_inv_e * ži * zs + η_inv_i * zi * žs))
 610     RH = -((-η_inv_e * že * zi + η_inv_i * ze * ži)/(-η_inv_i * ži * zs + η_inv_e * zi * žs))
 611     TV = -((-η_inv_e * že * zs + η_inv_e * ze * žs)/( η_inv_e * ži * zs - η_inv_i * zi * žs))
 612     TH = -(( η_inv_e * že * zs - η_inv_e * ze * žs)/(-η_inv_i * ži * zs + η_inv_e * zi * žs))
 613     return (RH, RV, TH, TV)
 614
 615 #@ujit
 616 def G_Mie_scat_precalc_cart_new(source_cart, dest_cart, RH, RV, a, nmax, k_i, k_e, μ_i=1, μ_e=1, J_ext=1, J_scat=3):
 617     """
 618     Implementation according to Kristensson, page 50
 619     My (Taylor's) basis functions are normalized to n*(n+1), whereas Kristensson's to 1
 620     TODO: check possible -1 factors (cf. Kristensson's dagger notation)
 621     """
 622     my, ny = get_mn_y(nmax)
 623     nelem = len(my)
 624     #source to origin
 625     source_sph = cart2sph(source_cart)
 626     source_sph[0] = k_e * source_sph[0]
 627     dest_sph = cart2sph(dest_cart)
 628     dest_sph[0] = k_e * dest_sph[0]
 629     if(dest_sph[0].real >= source_sph[0].real):
 630         lo_sph = source_sph
 631         hi_sph = dest_sph
 632     else:
 633         lo_sph = dest_sph
 634         hi_sph = source_sph
 635     lo_sph = source_sph
 636     hi_sph = dest_sph
 637
 638     M̃lo_y, Ñlo_y = vswf_yr1(lo_sph,nmax,J=J_scat)
 639     lo_loccart_basis = sph_loccart_basis(lo_sph, sphaxis=-1, cartaxis=None)
 640     M̃lo_cart_y = np.sum(M̃lo_y[:,:,ň]*lo_loccart_basis[ň,:,:],axis=-2)
 641     Ñlo_cart_y = np.sum(Ñlo_y[:,:,ň]*lo_loccart_basis[ň,:,:],axis=-2)
 642
 643     M̃hi_y, Ñhi_y = vswf_yr1(hi_sph,nmax,J=J_scat)#J_scat
 644     hi_loccart_basis = sph_loccart_basis(hi_sph, sphaxis=-1, cartaxis=None)
 645     M̃hi_cart_y = np.sum(M̃hi_y[:,:,ň]*hi_loccart_basis[ň,:,:],axis=-2)
 646     Ñhi_cart_y = np.sum(Ñhi_y[:,:,ň]*hi_loccart_basis[ň,:,:],axis=-2)
 647
 648     G_y = (RH[ny][:,ň,ň] * M̃lo_cart_y[:,:,ň].conj() * M̃hi_cart_y[:,ň,:] +
 649            RV[ny][:,ň,ň] * Ñlo_cart_y[:,:,ň].conj() * Ñhi_cart_y[:,ň,:]) / (ny * (ny+1))[:,ň,ň]
 650     return 1j* k_e*np.sum(G_y,axis=0)
 651
 652
 653 # From PRL 112, 253601 (1)
 654 #@ujit
 655 def Grr_Delga(nmax, a, r, k, ε_m, ε_b):
 656     om = k * c
 657     z = (r-a)/a
 658     g0 = om*cmath.sqrt(ε_b)/(6*c*π)
 659     n = np.arange(1,nmax+1)
 660     s = np.sum( (n+1)**2 * (ε_m-ε_b) / ((1+z)**(2*n+4) * (ε_m + ((n+1)/n)*ε_b)))
 661     return (g0 + s * c**2/(4*π*om**2*ε_b*a**3))
 662
 663
 664
 665 # TODOs
 666 # ====
 667 #
 668 # Rewrite the functions zJn, lpy in (at least simulated) universal manner.
 669 # Then universalise the rest
 670 #
 671 # Implement the actual multiple scattering
 672 #
 673 # Test if the decomposition of plane wave works also for absorbing environment (complex k).
 674
 675 # From PRL 112, 253601 (1)
 676 #@ujit
 677 def Grr_Delga(nmax, a, r, k, ε_m, ε_b):
 678     om = k * c
 679     z = (r-a)/a
 680     g0 = om*cmath.sqrt(ε_b)/(6*c*π)
 681     n = np.arange(1,nmax+1)
 682     s = np.sum( (n+1)**2 * (ε_m-ε_b) / ((1+z)**(2*n+4) * (ε_m + ((n+1)/n)*ε_b)))
 683     return (g0 + s * c**2/(4*π*om**2*ε_b*a**3))
 684
 685 #@ujit
 686 def G0_dip_1(r_cart,k):
 687     """
 688     Free-space dyadic Green's function in terms of the spherical vector waves.
 689     FIXME
 690     """
 691     sph = cart2sph(r_cart*k)
 692     pfz = 0.32573500793527994772 # 1./math.sqrt(3.*π)
 693     pf = 0.23032943298089031951 # 1./math.sqrt(6.*π)
 694     M1_y, N1_y = vswf_yr1(sph,nmax = 1,J=3)
 695     loccart_basis = sph_loccart_basis(sph, sphaxis=-1, cartaxis=None)
 696     N1_cart = np.sum(N1_y[:,:,ň]*loccart_basis[ň,:,:],axis=-2)
 697     coeffs_cart = np.array([[pf,-1j*pf,0.],[0.,0.,pfz],[-pf,-1j*pf,0.]]).conj()
 698     return 1j*k*np.sum(coeffs_cart[:,:,ň]*N1_cart[:,ň,:],axis=0)/2.
 699
 700 # Free-space dyadic Green's functions from RMP 70, 2, 447 =: [1]
 701 # (The numerical value is correct only at the regular part, i.e. r != 0)
 702 #@ujit
 703 def _P(z):
 704     return (1-1/z+1/(z*z))
 705 #@ujit
 706 def _Q(z):
 707     return (-1+3/z-3/(z*z))
 708
 709 # [1, (9)] FIXME The sign here is most likely wrong!!!
 710 #@ujit
 711 def G0_analytical(r #cartesian!
 712                   , k):
 713     I=np.identity(3)
 714     rn = sph_loccart2cart(np.array([1.,0.,0.]), cart2sph(r), axis=-1)
 715     rnxrn = rn[...,:,ň] * rn[...,ň,:]
 716     r = np.linalg.norm(r, axis=-1)
 717     #print(_P(1j*k*r).shape,_Q(1j*k*r).shape, rnxrn.shape, I.shape)
 718     return ((-np.exp(1j*k*r)/(4*π*r))[...,ň,ň] *
 719                    (_P(1j*k*r)[...,ň,ň]*I
 720                     +_Q(1j*k*r)[...,ň,ň]*rnxrn
 721                    ))
 722
 723 # [1, (11)]
 724 #@ujit
 725 def G0L_analytical(r, k):
 726     I=np.identity(3)
 727     rn = sph_loccart2cart(np.array([1.,0.,0.]), cart2sph(r), axis=-1)
 728     rnxrn = rn[...,:,ň] * rn[...,ň,:]
 729     r = np.linalg.norm(r, axis=-1)
 730     return (I-3*rnxrn)/(4*π*k*k*r**3)[...,ň,ň]
 731
 732 # [1,(10)]
 733 #@jit
 734 def G0T_analytical(r, k):
 735     return G0_analytical(r,k) - G0L_analytical(r,k)
 736
 737 #@ujit
 738 def G0_sum_1_slow(source_cart, dest_cart, k, nmax):
 739     my, ny = get_mn_y(nmax)
 740     nelem = len(my)
 741     RH = np.full((nelem),1)
 742     RV = RH
 743     return G_Mie_scat_precalc_cart(source_cart, dest_cart, RH, RV, a=0.001, nmax=nmax, k_i=1, k_e=k, μ_i=1, μ_e=1, J_ext=1, J_scat=3)
 744
 745
 746 def loadWfile(fileName, lMax = None, fatForm = True, nparticles = 2):
 747     '''
 748     Load the translation operator lattice sums generated by hexlattice_ewald.c
 749     fatForm has different indices (so that it can be easily multiplied with T-Matrix
 750     and twice the size.
 751     '''
 752     #nparticles = 2 # TODO generalize
 753     um = 1e-6 # micrometer in SI units
 754     data = np.loadtxt(fileName)
 755     nsamples = data.shape[0]
 756     Wdata = data[:,5:]
 757     if (lMax is None):
 758         nelem = int((Wdata.shape[1]/nparticles**2)**.5/2)
 759         lMax = nelem2lMax(nelem)
 760     else:
 761         nelem = lMax2nelem(lMax)
 762     Ws = Wdata[:,::2] + 1j*Wdata[:,1::2]
 763     # indices: destparticle, srcparticle, wavetype, desty, srcy
 764     Ws.shape = (nsamples, nparticles, nparticles, 2, nelem, nelem)
 765     # maybe TODO copy the following so the original data gets freed to save memory
 766     freqs_weirdunits = np.array(data[:,0], copy=True)
 767     k0s = np.array(data[:,2], copy=True)
 768     ks = np.array(data[:,3:5], copy=True)
 769     freqs = freqs_weirdunits * c / um
 770     if fatForm: #indices: (...,) destparticle, desttype, desty, srcparticle, srctype, srcy
 771         Ws2 = np.moveaxis(Ws, [-5,-4,-3,-2,-1], [-4,-2,-5,-3,-1] )
 772         fatWs = np.empty(Ws2.shape[:-5]+(nparticles, 2,nelem, nparticles, 2, nelem),dtype=complex)
 773         fatWs[...,:,0,:,:,0,:] = Ws2[...,0,:,:,:,:] #A
 774         fatWs[...,:,1,:,:,1,:] = Ws2[...,0,:,:,:,:] #A
 775         fatWs[...,:,1,:,:,0,:] = Ws2[...,1,:,:,:,:] #B
 776         fatWs[...,:,0,:,:,1,:] = Ws2[...,1,:,:,:,:] #B
 777         Ws = fatWs
 778     return {
 779             'lMax' : lMax,
 780             'nelem' : nelem,
 781             'npart' : nparticles,
 782             'freqs' : freqs,
 783             'ks' : ks,
 784             'k0_effs' : k0s,
 785             'Ws' : Ws,
 786             'freqs_weirdunits' : freqs_weirdunits,
 787             }
 788
 789 def processWfiles_sameKs(freqfilenames, destfilename, nparticles = 2, lMax = None, f='npz'):
 790     '''
 791     Processes translation operator data in different files; each file is supposed to contain one frequency.
 792     The Ks in the different files are expected to be exactly the same and in the same order;
 793     the result is sorted by frequencies and saved to npz file;
 794     The representation is always "thin", i.e. the elements which are zero due to z-symmetry are skipped
 795     F is either 'npz' or 'd'
 796     'npz' creates npz archive, 'd' creates a directory with individual npy files, where the W-matrix data
 797     are writen using memmap.
 798     '''
 799
 800     if (nparticles is None):
 801         # TODO here read nparticles from file;
 802         raise # NOT IMPLEMENTED
 803
 804     um = 1e-6 # micrometre in SI units
 805
 806     nfiles = len(freqfilenames)
 807
 808     if (nfiles < 1): return ValueError("At least one filename has to be passed")
 809
 810     #Check the first file to get the "constants" set
 811     filename = freqfilenames[0]
 812     data = np.loadtxt(filename, ndmin = 2)
 813     nk_muster = data.shape[0]
 814     Wdata = data[...,5:]
 815
 816     if (lMax is None):
 817         nelem = int((Wdata.shape[1]/nparticles**2)**.5/2)
 818         lMax = nelem2lMax(nelem)
 819     else:
 820         nelem = lMax2nelem(lMax)
 821     ks_muster = np.array(data[...,3:5], copy=True)
 822
 823     if f == 'npz':
 824         allWs = np.empty((nfiles, nk_muster, nparticles, nparticles, 2, nelem, nelem,), dtype=complex)
 825     elif f == 'd':
 826         j = os.path.join
 827         d = destfilename
 828         try:
 829             os.mkdir(destfilename)
 830         except FileExistsError as exc:
 831             pass
 832         allWs = np.lib.format.open_memmap(filename=j(d,"Wdata.npy"), dtype=complex, shape=(nfiles, nk_muster, nparticles, nparticles,2,nelem,nelem,),mode='w+')
 833     else: raise ValueError("f has to be either 'npz' or 'd'")
 834     k0s = np.empty((nfiles,))
 835     freqs_weirdunits = np.empty((nfiles,))
 836     EeVs = np.empty((nfiles,))
 837     succread = 0
 838
 839     for filename in freqfilenames:
 840         data = np.loadtxt(filename, ndmin=2)
 841         if data.shape[0] != nk_muster:
 842             raise ValueError("%s contains different number of lines than %s"%(filename, freqfilenames[0]))
 843         ks_current = np.array(data[:,3:5])
 844         if not np.all(ks_current == ks_muster):
 845             raise ValueError("%s contains different k-vectors than %s"%(filename, freqfilenames[0]))
 846         curWdata = data[...,5:]
 847         curWs = curWdata[...,::2] + 1j*curWdata[...,1::2]
 848         curWs.shape = (nk_muster, nparticles, nparticles, 2, nelem, nelem)
 849         allWs[succread] = curWs
 850         cur_freqs_weirdunits = np.array(data[:,0])
 851         if not np.all(cur_freqs_weirdunits == cur_freqs_weirdunits[0]):
 852             raise ValueError("%s contains various frequencies; we require only one frequency per file"%(filename))
 853         freqs_weirdunits[succread] = cur_freqs_weirdunits[0]
 854         k0s[succread] = data[0,2] # TODO check this as well?
 855         EeVs[succread] = data[0,1]
 856         succread += 1
 857     del allWs
 858     freqs = freqs_weirdunits * c / um
 859
 860     if f == 'npz':
 861         np.savez(destfilename,
 862              lMax = lMax,
 863              nelem = nelem,
 864              npart = nparticles,
 865              nfreqs = succread,
 866              nk = nk_muster,
 867              freqs = freqs[:succread],
 868              freqs_weirdunits = freqs_weirdunits[:succread],
 869              EeVs_orig = EeVs[:succread],
 870              k0s = k0s[:succread],
 871              ks = ks_muster,
 872              Wdata = allWs[:succread]
 873             )
 874     elif f == 'd':
 875         np.save(j(d,'lMax.npy'), lMax)
 876         np.save(j(d,'nelem.npy'), nelem)
 877         np.save(j(d,'npart.npy'), nparticles)
 878         np.save(j(d,'nfreqs.npy'), succread)
 879         np.save(j(d,'nk.npy'), nk_muster)
 880         np.save(j(d,'freqs.npy'), freqs[:succread])
 881         np.save(j(d,'freqs_weirdunits.npy'), freqs_weirdunits[:succread])
 882         np.save(j(d,'EeVs_orig.npy'), EeVs[:succread])
 883         np.save(j(d,'k0s.npy'), k0s[:succread])
 884         np.save(j(d,'ks.npy'), ks_muster)
 885         # Wdata already saved
 886     return
 887
 888
 889 def loadWfile_info(fileName, lMax = None, midk_halfwidth = None, freqlimits = None):
 890     """
 891     Returns all the relevant data from the W file except for the W values themselves
 892     """
 893     slovnik = loadWfile_processed(fileName, lMax = lMax, fatForm = False, midk_halfwidth=midk_halfwidth, freqlimits=freqlimits)
 894     slovnik['Ws'] = None
 895     return slovnik
 896
 897 def loadWfile_processed(fileName, lMax = None, fatForm = True, midk_halfwidth = None, midk_index = None, freqlimits = None, iteratechunk = None):
 898     """
 899     midk_halfwidth: int
 900         if given, takes only the "middle" k-value by index nk//2 and midk_halfwidth values from both sides
 901     midk_index: int
 902         if given, midk_index is taken as the "middle" value instad of nk//2
 903     freqlimit: pair of doubles;
 904         if given, only the values from the given frequency range (in Hz) are returned to save RAM
 905         N.B.: the frequencies in the processed file are expected to be sorted!
 906     iteratechunk: positive int ; NI
 907         if given, this works as a generator with only iteratechunk frequencies processed and returned at each iteration to save memory
 908     """
 909     if os.path.isdir(fileName): # .npy files in a directory
 910         p = fileName
 911         j = os.path.join
 912         nk = np.load(j(p,'nk.npy'))[()]
 913         nfreqs = np.load(j(p,'nfreqs.npy'))[()]
 914         nparticles = np.load(j(p,'npart.npy'))[()]
 915         Ws = np.load(j(p,'Wdata.npy'), mmap_mode='r')
 916         ks = np.load(j(p,'ks.npy'))
 917         freqs = np.load(j(p,'freqs.npy'))
 918         k0s = np.load(j(p,'k0s.npy'))
 919         EeVs_orig = np.load(j(p,'EeVs_orig.npy'))
 920         freqs_weirdunits = np.load(j(p,'freqs_weirdunits.npy'))
 921     else: # npz file
 922         data = np.load(fileName, mmap_mode='r')
 923         nk = data['nk'][()]
 924         nfreqs = data['nfreqs'][()]
 925         nparticles = data['npart'][()]
 926         Ws = data['Wdata']
 927         ks = data['ks']
 928         freqs = data['freqs']
 929         k0s = data['k0s']
 930         EeVs_orig = data['EeVs_orig']
 931         freqs_weirdunits = data['freqs_weirdunits']
 932     if midk_halfwidth is not None:
 933         k_mid_i = nk // 2 if midk_index is None else midk_index
 934         maxk_i = min(k_mid_i + midk_halfwidth, nk)
 935         mink_i = max(0, k_mid_i - midk_halfwidth)
 936         nk = maxk_i + 1 - mink_i
 937         Ws = Ws[:,mink_i:(maxk_i+1)]
 938         ks = ks[mink_i:(maxk_i+1)]
 939         #k_mid_i_new = k_mid_i - mink_i
 940     if lMax is not None:
 941         nelem = lMax2nelem(lMax)
 942         Ws = Ws[...,:nelem,:nelem]
 943     else:
 944         lMax = data['lMax'][()]
 945         nelem = lMax2nelem(lMax)
 946     if freqlimits is not None:
 947         minind = np.searchsorted(freqs, freqlimits[0], side='left')
 948         maxind = np.searchsorted(freqs, freqlimits[1], side='right')
 949         freqs = freqs[minind:maxind]
 950         Ws = Ws[minind:maxind]
 951         k0s = k0s[minind:maxind]
 952         EeVs_orig = EeVs_orig[minind:maxind]
 953         freqs_weirdunints = freqs_weirdunits[minind:maxind]
 954         nfreqs = maxind-minind
 955     if iteratechunk is None: # everyting at once
 956         if fatForm: #indices: (...,) destparticle, desttype, desty, srcparticle, srctype, srcy
 957             Ws2 = np.moveaxis(Ws, [-5,-4,-3,-2,-1], [-4,-2,-5,-3,-1] )
 958             fatWs = np.empty(Ws2.shape[:-5]+(nparticles, 2,nelem, nparticles, 2, nelem),dtype=complex)
 959             fatWs[...,:,0,:,:,0,:] = Ws2[...,0,:,:,:,:] #A
 960             fatWs[...,:,1,:,:,1,:] = Ws2[...,0,:,:,:,:] #A
 961             fatWs[...,:,1,:,:,0,:] = Ws2[...,1,:,:,:,:] #B
 962             fatWs[...,:,0,:,:,1,:] = Ws2[...,1,:,:,:,:] #B
 963             Ws = fatWs
 964         return{
 965                 'lMax': lMax,
 966                 'nelem': nelem,
 967                 'npart': nparticles,
 968                 'nfreqs': nfreqs,
 969                 'nk' : nk,
 970                 'freqs': freqs,
 971                 'freqs_weirdunits': freqs_weirdunits,
 972                 'EeVs_orig': EeVs_orig,
 973                 'k0s': k0s,
 974                 'ks': ks,
 975                 'Ws': Ws,
 976                 }
 977     else:
 978         def gen(lMax, nelem, nparticles, nfreqs, nk, freqs, freqs_weirdunits, EeVs_orig, k0s, ks, Ws, iteratechunk):
 979             starti = 0
 980             while(starti < nfreqs):
 981                 stopi = min(starti+iteratechunk, nfreqs)
 982                 chunkWs = Ws[starti:stopi]
 983                 if fatForm: #indices: (...,) destparticle, desttype, desty, srcparticle, srctype, srcy
 984                     Ws2 = np.moveaxis(chunkWs, [-5,-4,-3,-2,-1], [-4,-2,-5,-3,-1] )
 985                     fatWs = np.empty(Ws2.shape[:-5]+(nparticles, 2,nelem, nparticles, 2, nelem),dtype=complex)
 986                     fatWs[...,:,0,:,:,0,:] = Ws2[...,0,:,:,:,:] #A
 987                     fatWs[...,:,1,:,:,1,:] = Ws2[...,0,:,:,:,:] #A
 988                     fatWs[...,:,1,:,:,0,:] = Ws2[...,1,:,:,:,:] #B
 989                     fatWs[...,:,0,:,:,1,:] = Ws2[...,1,:,:,:,:] #B
 990                     chunkWs = fatWs
 991                 yield {
 992                         'lMax': lMax,
 993                         'nelem': nelem,
 994                         'npart': nparticles,
 995                         'nfreqs': stopi-starti,
 996                         'nk' : nk,
 997                         'freqs': freqs[starti:stopi],
 998                         'freqs_weirdunits': freqs_weirdunits[starti:stopi],
 999                         'EeVs_orig': EeVs_orig[starti:stopi],
1000                         'k0s': k0s[starti:stopi],
1001                         'ks': ks,
1002                         'Ws': chunkWs,
1003                         'chunk_range': (starti, stopi),
1004                         'nfreqs_total': nfreqs
1005                         }
1006                 starti += iteratechunk
1007         return gen(lMax, nelem, nparticles, nfreqs, nk, freqs, freqs_weirdunits, EeVs_orig, k0s, ks, Ws, iteratechunk)
1008
1009
1010