qpms/tmatrices.py

   1 import numpy as np
   2 use_moble_quaternion = False
   3 try:
   4     import quaternion, spherical_functions as sf # because of the Wigner matrices. These imports are SLOW.
   5     use_moble_quaternion = True
   6 except ImportError:
   7     use_moble_quaternion = False
   8
   9 import re
  10 from scipy import interpolate
  11 from scipy.constants import hbar, e as eV, pi, c
  12 from .cycommon import get_mn_y, get_nelem
  13 from .cyquaternions import CQuat
  14 ň = np.newaxis
  15 from .types import NormalizationT, TMatrixSpec
  16
  17 # Transformations of spherical bases
  18 def WignerD_mm(l, quat):
  19     """
  20     Calculates Wigner D matrix (as an numpy (2*l+1,2*l+1)-shaped array)
  21     for order l, and a rotation given by quaternion quat.
  22
  23     This represents the rotation of spherical vector basis
  24     TODO doc
  25     """
  26
  27     if use_moble_quaternion:
  28         indices = np.array([ [l,i,j] for i in range(-l,l+1) for j in range(-l,l+1)])
  29         Delems = sf.Wigner_D_element(quat, indices).reshape(2*l+1,2*l+1)
  30         return Delems
  31     else:
  32         Delems = np.zeros((2*l+1, 2*l+1), dtype=complex)
  33         for i in range(-l,l+1):
  34             for j in range(-l,l+1):
  35                 Delems[i,j] = quat.wignerDelem(l, i, j)
  36         return Delems
  37
  38
  39 def WignerD_mm_fromvector(l, vect):
  40     """
  41     TODO doc
  42     """
  43     if use_moble_quaternion:
  44         return WignerD_mm(l, quaternion.from_rotation_vector(vect))
  45     else:
  46         return WignerD_mm(l, CQuat.from_rotvector(vect))
  47
  48
  49 def WignerD_yy(lmax, quat):
  50     """
  51     TODO doc
  52     """
  53     my, ny = get_mn_y(lmax)
  54     Delems = np.zeros((len(my),len(my)),dtype=complex)
  55     b_in = 0
  56     e_in =  None
  57     for l in range(1,lmax+1):
  58         e_in = b_in + 2*l+1
  59         Delems[b_in:e_in,b_in:e_in] = WignerD_mm(l, quat)
  60         b_in = e_in
  61     return Delems
  62
  63 def WignerD_yy_fromvector(lmax, vect):
  64     """
  65     TODO doc
  66     """
  67     if use_moble_quaternion:
  68         return WignerD_yy(lmax, quaternion.from_rotation_vector(vect))
  69     else:
  70         return WignerD_yy(lMax, CQuat.from_rotvector(vect))
  71
  72 def identity_yy(lmax):
  73     """
  74     TODO doc
  75     """
  76     return np.eye(lMax2nelem(lMax))
  77
  78 def identity_tyty(lmax):
  79     """
  80     TODO doc
  81     """
  82     nelem = lMax2nelem(lmax)
  83     return np.eye(2*nelem).reshape((2,nelem,2,nelem))
  84
  85 def xflip_yy(lmax):
  86     """
  87     TODO doc
  88     xflip = δ(m + m') δ(l - l')
  89     (i.e. ones on the (m' m) antidiagonal
  90     """
  91     my, ny = get_mn_y(lmax)
  92     elems = np.zeros((len(my),len(my)),dtype=int)
  93     b_in = 0
  94     e_in =  None
  95     for l in range(1,lmax+1):
  96         e_in = b_in + 2*l+1
  97         elems[b_in:e_in,b_in:e_in] = np.eye(2*l+1)[::-1,:]
  98         b_in = e_in
  99     return elems
 100
 101 def xflip_tyy(lmax):
 102     fl_yy = xflip_yy(lmax)
 103     return np.array([fl_yy,-fl_yy])
 104
 105 def xflip_tyty(lmax):
 106     fl_yy = xflip_yy(lmax)
 107     nelem = fl_yy.shape[0]
 108     fl_tyty = np.zeros((2,nelem,2,nelem),dtype=int)
 109     fl_tyty[0,:,0,:] = fl_yy
 110     fl_tyty[1,:,1,:] = -fl_yy
 111     return fl_tyty
 112
 113 def yflip_yy(lmax):
 114     """
 115     TODO doc
 116     yflip = rot(z,pi/2) * xflip * rot(z,-pi/2)
 117           = δ(m + m') δ(l - l') * (-1)**m
 118     """
 119     my, ny = get_mn_y(lmax)
 120     elems = xflip_yy(lmax)
 121     elems[(my % 2)==1] = elems[(my % 2)==1] * -1 # Obvious sign of tiredness (this is correct but ugly; FIXME)
 122     return elems
 123
 124 def yflip_tyy(lmax):
 125     fl_yy = yflip_yy(lmax)
 126     return np.array([fl_yy,-fl_yy])
 127
 128 def yflip_tyty(lmax):
 129     fl_yy = yflip_yy(lmax)
 130     nelem = fl_yy.shape[0]
 131     fl_tyty = np.zeros((2,nelem,2,nelem),dtype=int)
 132     fl_tyty[0,:,0,:] = fl_yy
 133     fl_tyty[1,:,1,:] = -fl_yy
 134     return fl_tyty
 135
 136 def zflip_yy(lmax):
 137     """
 138     TODO doc
 139     zflip = (-1)^(l+m)
 140     """
 141     my, ny = get_mn_y(lmax)
 142     elems = np.zeros((len(my), len(my)), dtype=int)
 143     b_in = 0
 144     e_in = None
 145     for l in range(1,lmax+1):
 146         e_in = b_in + 2*l+1
 147         elems[b_in:e_in,b_in:e_in] = np.diag([(-1)**i for i in range(e_in-b_in)])
 148         b_in = e_in
 149     return elems
 150
 151 def zflip_tyy(lmax):
 152     fl_yy = zflip_yy(lmax)
 153     return np.array([fl_yy,-fl_yy])
 154
 155 def zflip_tyty(lmax):
 156     fl_yy = zflip_yy(lmax)
 157     nelem = fl_yy.shape[0]
 158     fl_tyty = np.zeros((2,nelem,2,nelem),dtype=int)
 159     fl_tyty[0,:,0,:] = fl_yy
 160     fl_tyty[1,:,1,:] = -fl_yy
 161     return fl_tyty
 162
 163 def zrotN_yy(N, lMax):
 164     return WignerD_yy_fromvector(lMax, np.array([0,0,pi * (2/N)]))
 165
 166 def op_yy2tyty(yyop):
 167     '''
 168     Broadcasts an yy operator to tyty operator without considering mirroring effects.
 169     Good (maybe only) for rotations.
 170     '''
 171     return np.moveaxis(np.eye(2)[:,:,ň,ň] * yyop, 2,1)
 172
 173 def zrotN_tyty(N, lMax):
 174     return op_yy2tyty(zrotN_yy(N, lMax))
 175
 176 def parity_yy(lmax):
 177     """
 178     Parity operator (flip in x,y,z)
 179     parity = (-1)**l
 180     """
 181     my, ny = get_mn_y(lmax)
 182     return np.diag((-1)**ny)
 183
 184 # BTW parity (xyz-flip) is simply (-1)**ny
 185
 186
 187
 188 #----------------------------------------------------#
 189 # Loading T-matrices from scuff-tmatrix output files #
 190 #----------------------------------------------------#
 191
 192 # We don't really need this particular function anymore, but...
 193 def _scuffTMatrixConvert_EM_01(EM):
 194     #print(EM)
 195     if (EM == b'E'):
 196         return 1
 197     elif (EM == b'M'):
 198         return 0
 199     else:
 200         return None
 201
 202 def loadScuffTMatrices(fileName, normalisation = 1, version = 'old', freqscale = None, order = None):
 203     """
 204     TODO doc
 205
 206     version describes version of scuff-em. It is either 'old' or 'new'.
 207     default order is ('N','M') with 'old' version, ('M','N') with 'new'
 208     """
 209     oldversion = (version == 'old')
 210     μm = 1e-6
 211     table = np.genfromtxt(fileName,
 212                   converters={1: _scuffTMatrixConvert_EM_01, 4: _scuffTMatrixConvert_EM_01} if oldversion else None,
 213                   skip_header = 0 if oldversion else 5,
 214                   usecols = None if oldversion else (0, 2, 3, 4, 6, 7, 8, 9, 10),
 215                   dtype=[('freq', '<f8'),
 216                          ('outc_type', '<i8'), ('outc_l', '<i8'), ('outc_m', '<i8'),
 217                          ('inc_type', '<i8'), ('inc_l', '<i8'), ('inc_m', '<i8'),
 218                          ('Treal', '<f8'), ('Timag', '<f8')]
 219                    if oldversion else
 220                         [('freq', '<f8'),
 221                          ('outc_l', '<i8'), ('outc_m', '<i8'),  ('outc_type', '<i8'),
 222                          ('inc_l', '<i8'), ('inc_m', '<i8'), ('inc_type', '<i8'),
 223                          ('Treal', '<f8'), ('Timag', '<f8')]
 224                     )
 225     lMax=np.max(table['outc_l'])
 226     my,ny = get_mn_y(lMax)
 227     nelem = len(ny)
 228     TMatrix_sz = nelem**2 * 4 # number of rows for each frequency: nelem * nelem spherical incides, 2 * 2 E/M types
 229     freqs_weirdunits = table['freq'][::TMatrix_sz].copy()
 230     freqs = freqs_weirdunits * c / μm
 231     # The iteration in the TMatrix file goes in this order (the last one iterates fastest, i.e. in the innermost loop):
 232     # freq outc_l outc_m outc_type inc_l inc_m inc_type
 233     # The l,m mapping is the same as is given by my get_mn_y function, so no need to touch that
 234     TMatrices_tmp_real = table['Treal'].reshape(len(freqs), nelem, 2, nelem, 2)
 235     TMatrices_tmp_imag = table['Timag'].reshape(len(freqs), nelem, 2, nelem, 2)
 236     # There are two přoblems with the previous matrices. First, we want to have the
 237     # type indices first, so we want a shape (len(freqs), 2, nelem, 2, nelem) as in the older code.
 238     # Second, M-waves come first, so they have now 0-valued index, and E-waves have 1-valued index,
 239     # which we want to be inverted.
 240     TMatrices = np.zeros((len(freqs),2,nelem,2,nelem),dtype=complex)
 241     reorder = (0,1)
 242     if ((order == ('N', 'M')) and not oldversion): # reverse order for the new version
 243         reorder = (1,0)
 244     # TODO reverse order for the old version...
 245     for inc_type in (0,1):
 246         for outc_type in (0,1):
 247             TMatrices[:,reorder[outc_type],:,reorder[inc_type],:] = TMatrices_tmp_real[:,:,outc_type,:,inc_type]+1j*TMatrices_tmp_imag[:,:,outc_type,:,inc_type]
 248     # IMPORTANT: now we are going from Reid's/Kristensson's/Jackson's/whoseever convention to Taylor's convention
 249     # TODO make these consistent with what is defined in qpms_types.h (implement all possibilities)
 250     if normalisation == 1:
 251         TMatrices[:,:,:,:,:] = TMatrices[:,:,:,:,:] * np.sqrt(ny*(ny+1))[ň,ň,ň,ň,:] / np.sqrt(ny*(ny+1))[ň,ň,:,ň,ň]
 252     elif normalisation == 2: # Kristensson?
 253         pass
 254     if freqscale is not None:
 255         freqs *= freqscale
 256         freqs_weirdunits *= freqscale
 257     return (TMatrices, freqs, freqs_weirdunits, lMax)
 258
 259
 260 # misc tensor maniputalion
 261 def apply_matrix_left(matrix, tensor, axis):
 262     """
 263     TODO doc
 264     Apply square matrix to a given axis of a tensor, so that the result retains the shape
 265     of the original tensor. The summation goes over the second index of the matrix and the
 266     given tensor axis.
 267     """
 268     tmp = np.tensordot(matrix, tensor, axes=(-1,axis))
 269     return np.moveaxis(tmp, 0, axis)
 270
 271 def apply_ndmatrix_left(matrix,tensor,axes):
 272     """
 273     Generalized apply_matrix_left, the matrix can have more (2N) abstract dimensions,
 274     like M[i,j,k,...z,i,j,k,...,z]. N axes have to be specified in a tuple, corresponding
 275     to the axes 0,1,...N-1 of the matrix
 276     """
 277     N = len(axes)
 278     matrix = np.tensordot(matrix, tensor, axes=([-N+axn for axn in range(N)],axes))
 279     matrix = np.moveaxis(matrix, range(N), axes)
 280     return matrix
 281
 282 def apply_ndmatrix_right(tensor, matrix, axes):
 283     """
 284     Right-side analogue of apply_ndmatrix_lift.
 285     Multiplies a tensor with a 2N-dimensional matrix, conserving the axis order.
 286     """
 287     N = len(axes)
 288     matrix = np.tensordot(tensor, matrix, axes = (axes, range(N)))
 289     matrix = np.moveaxis(matrix, [-N+axn for axn in range(N)], axes)
 290     return matrix
 291
 292 def ndmatrix_Hconj(matrix):
 293     """
 294     For 2N-dimensional matrix, swap the first N and last N matrix, and complex conjugate.
 295     """
 296     twoN = len(matrix.shape)
 297     if not twoN % 2 == 0:
 298         raise ValueError("The matrix has to have even number of axes.")
 299     N = twoN//2
 300     matrix = np.moveaxis(matrix, range(N), range(N, 2*N))
 301     return matrix.conj()
 302
 303 def symz_indexarrays(lMax, npart = 1):
 304     """
 305     Returns indices that are used for separating the in-plane E ('TE' in the photonic crystal
 306     jargon) and perpendicular E ('TM' in the photonic crystal jargon) modes
 307     in the z-mirror symmetric systems.
 308
 309     Parameters
 310     ----------
 311     lMax : int
 312         The maximum degree cutoff for the T-matrix to which these indices will be applied.
 313
 314     npart : int
 315         Number of particles (TODO better description)
 316
 317     Returns
 318     -------
 319     TEč, TMč : (npart * 2 * nelem)-shaped bool ndarray
 320         Mask arrays corresponding to the 'TE' and 'TM' modes, respectively.
 321     """
 322     my, ny = get_mn_y(lMax)
 323     nelem = len(my)
 324     ž = np.arange(2*nelem) # single particle spherical wave indices
 325     tž = ž // nelem # tž == 0: electric waves, tž == 1: magnetic waves
 326     mž = my[ž%nelem]
 327     nž = ny[ž%nelem]
 328     TEž = ž[(mž+nž+tž) % 2 == 0]
 329     TMž = ž[(mž+nž+tž) % 2 == 1]
 330
 331     č = np.arange(npart*2*nelem) # spherical wave indices for multiple particles (e.g. in a unit cell)
 332     žč = č % (2* nelem)
 333     tč = tž[žč]
 334     mč = mž[žč]
 335     nč = nž[žč]
 336     TEč = č[(mč+nč+tč) % 2 == 0]
 337     TMč = č[(mč+nč+tč) % 2 == 1]
 338     return (TEč, TMč)
 339
 340
 341 """
 342 Processing T-matrix related operations from scripts
 343 ===================================================
 344
 345 see also scripts_common.py
 346
 347 The unit cell is defined by a dict particle_specs and a list TMatrix_specs.
 348 This particular module has to provide the T-matrices according to what is defined
 349 in TMatrix_specs
 350
 351 TMatrix_specs is a list of tuples (lMax_override, TMatrix_path, ops)
 352 where
 353   - TMatrix_path is path to the file generated by scuff-tmatrix
 354   - lMax_override is int or None; if it is int and less than the lMax found from the T-matrix file,
 355     lMax_override is used as the order cutoff for the output T-matrix.
 356   - ops is an iterable of tuples (optype, opargs) where currently optype can be 'sym' or 'tr'
 357     for symmetrization operation or some other transformation.
 358 """
 359
 360 #TODO FEATURE more basic group symmetry operations, cf. http://symmetry.otterbein.edu/tutorial/index.html
 361
 362 # This is for finite „fractional“ rotations along the z-axis (mCN means rotation of 2π*(m/N))
 363 reCN = re.compile('(\d*)C(\d+)') # TODO STYLE make this regexp also accept the 3*C_5-type input, eqiv. to 3C5
 364
 365 def get_TMatrix_fromspec(tmatrix_spec):
 366     ''' TODO doc
 367     returns (TMatrices, freqs, lMax)
 368     '''
 369     lMax_override, tmpath, ops = tmatrix_spec
 370     TMatrices, freqs, freqs_weirdunits, lMax = loadScuffTMatrices(tmpath)
 371     if lMax_override is not None and (lMax_override < lMax_orig):
 372         nelem = get_nelem(lMax_override)
 373         TMatrices = TMatrices[...,0:nelem,:,0:nelem]
 374         lMax = lMax_override
 375
 376     for (optype, opargs) in ops:
 377         if optype == 'sym':
 378             mCN = reCN.match(opargs)
 379             if opargs == 'C2' or opargs == 'C_2':
 380                 opmat = apply_matrix_left(yflip_yy(lMax), xflip_yy(lMax), -1)
 381                 TMatrices = (TMatrices + apply_matrix_left(opmat, apply_matrix_left(opmat, TMatrices, -3), -1))/2
 382             elif opargs == 'σ_x' or opargs == 'σx':
 383                 opmat = xflip_tyty(lMax)
 384                 TMatrices = (TMatrices + apply_ndmatrix_left(opmat, apply_ndmatrix_left(opmat, TMatrices, (-4,-3)),(-2,-1)))/2
 385             elif opargs == 'σ_y' or opargs == 'σy':
 386                 opmat = yflip_tyty(lMax)
 387                 TMatrices = (TMatrices + apply_ndmatrix_left(opmat, apply_ndmatrix_left(opmat, TMatrices, (-4,-3)),(-2,-1)))/2
 388             elif opargs == 'σ_z' or opargs == 'σz':
 389                 opmat = zflip_tyty(lMax)
 390                 TMatrices = (TMatrices + apply_ndmatrix_left(opmat, apply_ndmatrix_left(opmat, TMatrices, (-4,-3)),(-2,-1)))/2
 391             elif mCN:
 392                 rotN = int(mCN.group(2)) # the possible m is ignored
 393                 TMatrix_contribs = np.empty((rotN,)+TMatrices.shape, dtype=np.complex_)
 394                 for i in range(rotN):
 395                     rotangle = 2*np.pi*i / rotN
 396                     rot = WignerD_yy_fromvector(lMax, np.array([0,0,rotangle]))
 397                     rotinv = WignerD_yy_fromvector(lMax, np.array([0,0,-rotangle]))
 398                     TMatrix_contribs[i] = apply_matrix_left(rot, apply_matrix_left(rotinv, TMatrices, -3), -1)
 399                 TMatrices = np.sum(TMatrix_contribs, axis=0) / rotN
 400             elif opargs == 'C_inf' or opargs == 'Cinf' or opargs == 'C_∞' or opargs == 'C∞':
 401                 raise ValueError('not implemented: ', opargs)  # TODO FEATURE
 402             else:
 403                 raise ValueError('not implemented: ', opargs)
 404         elif optype == 'tr':
 405             mCN = reCN.match(opargs)
 406             if opargs == 'C2' or opargs == 'C_2':
 407                 opmat = apply_matrix_left(yflip_yy(lMax), xflip_yy(lMax), -1)
 408                 TMatrices = apply_matrix_left(opmat, apply_matrix_left(opmat, TMatrices, -3), -1)/2
 409             elif opargs == 'σ_x' or opargs == 'σx':
 410                 opmat = xflip_tyty(lMax)
 411                 TMatrices = apply_ndmatrix_left(opmat, apply_ndmatrix_left(opmat, TMatrices, (-4,-3)),(-2,-1))/2
 412             elif opargs == 'σ_y' or opargs == 'σy':
 413                 opmat = yflip_tyty(lMax)
 414                 TMatrices = apply_ndmatrix_left(opmat, apply_ndmatrix_left(opmat, TMatrices, (-4,-3)),(-2,-1))/2
 415             elif opargs == 'σ_z' or opargs == 'σz':
 416                 opmat = zflip_tyty(lMax)
 417                 TMatrices = apply_ndmatrix_left(opmat, apply_ndmatrix_left(opmat, TMatrices, (-4,-3)),(-2,-1))/2
 418             elif mCN:
 419                 rotN = int(mCN.group(2))
 420                 power = int(mCN.group(1)) if mCN.group(1) else 1
 421                 rotangle = 2*np.pi*power / rotN
 422                 rot = WignerD_yy_fromvector(lMax, np.array([0,0,rotangle]))
 423                 rotinv = WignerD_yy_fromvector(lMax, np.array([0,0,-rotangle]))
 424                 TMatrices = apply_matrix_left(rot, apply_matrix_left(rotinv, TMatrices, -3), -1)
 425             else:
 426                 raise ValueError('not implemented: ', opargs)
 427         else:
 428             raise ValueError('not implemented: ', optype)
 429     return (TMatrices, freqs, lMax)
 430
 431 class TMatrix(TMatrixSpec):
 432     '''
 433     TODO doc
 434
 435     TODO support for different/multiple interpolators
 436     '''
 437     def __init__(self, tmatrix_spec):
 438         #self.specification = tmatrix_spec
 439         self.lMax_override = tmatrix_spec.lMax_override
 440         self.tmatrix_path = tmatrix_spec.tmatrix_path
 441         self.ops = tmatrix_spec.ops
 442         self.tmdata, self.freqs, self.lMax = get_TMatrix_fromspec(tmatrix_spec)
 443         self.nelem = get_nelem(self.lMax)
 444         #self._interpolators = dict()
 445         self.default_interpolator = interpolate.interp1d(self.freqs,
 446             self.tmdata, axis=0, kind='linear', fill_value='extrapolate')
 447         self.normalization = NormalizationT.TAYLOR # TODO others are not supported by the loading functions
 448
 449     def atfreq(self, freq):
 450         freqarray = np.array(freq, copy=False)
 451         if freqarray.shape: # not just a scalar
 452             tm_interp = np.empty(freqarray.shape + (2, self.nelem, 2, self.nelem), dtype=np.complex_)
 453             for i in np.ndindex(freqarray.shape):
 454                 tm_interp[i] = self.default_interpolator(freqarray[i])
 455             return tm_interp
 456         else: # scalar
 457             return self.default_interpolator(freq)
 458
 459     __getitem__ = atfreq # might be changed later, use atfreq to be sure
 460
 461 def perform_tmspecs(tmspecs):
 462     """Takes a sequence of TMatrixSpec or TMatrix instances and returns
 463     a list of corresponding TMatrix instances"""
 464     return [(tmspec if hasattr(tmspec, "tmdata") else TMatrix(tmspec))
 465                 for tmspec in tmspecs]
 466