qpms/hexpoints_c.pyx

   1 import math
   2 import numpy as np
   3 cimport numpy as np
   4 nx = None
   5
   6 cdef double _s3 = math.sqrt(3)
   7
   8 from scipy.constants import c
   9 from .timetrack import _time_b, _time_e
  10 from .tmatrices import symz_indexarrays
  11 from .hexpoints import hexlattice_get_AB
  12
  13 cpdef hexlattice_zsym_getSVD(int lMax, TMatrices_om, double epsilon_b, double hexside, size_t maxlayer, double omega, klist, gaussianSigma=False, int onlyNmin=0, verbose=False):
  14     cdef np.ndarray[np.npy_double, ndim = 2] klist_c = klist
  15     btime = _time_b(verbose)
  16     cdef size_t nelem = lMax * (lMax + 2)
  17     _n2id = np.identity(2*nelem)
  18     _n2id.shape = (2,nelem,2,nelem)
  19     cdef np.ndarray[np.npy_double, ndim = 4] n2id = _n2id
  20     cdef double nan = float('nan')
  21     k_0 = omega * math.sqrt(epsilon_b) / c
  22     tdic = hexlattice_get_AB(lMax,k_0*hexside,maxlayer)
  23     cdef np.ndarray[np.npy_cdouble, ndim = 3] a_self = tdic['a_self'][:,:nelem,:nelem]
  24     cdef np.ndarray[np.npy_cdouble, ndim = 3] b_self = tdic['b_self'][:,:nelem,:nelem]
  25     cdef np.ndarray[np.npy_cdouble, ndim = 3] a_u2d = tdic['a_u2d'][:,:nelem,:nelem]
  26     cdef np.ndarray[np.npy_cdouble, ndim = 3] b_u2d = tdic['b_u2d'][:,:nelem,:nelem]
  27     cdef np.ndarray[np.npy_cdouble, ndim = 3] a_d2u = tdic['a_d2u'][:,:nelem,:nelem]
  28     cdef np.ndarray[np.npy_cdouble, ndim = 3] b_d2u = tdic['b_d2u'][:,:nelem,:nelem]
  29     cdef np.ndarray[np.npy_double, ndim = 2] unitcell_translations = tdic['self_tr']*hexside*_s3
  30     cdef np.ndarray[np.npy_double, ndim = 2] u2d_translations = tdic['u2d_tr']*hexside*_s3
  31     cdef np.ndarray[np.npy_double, ndim = 2] d2u_translations = tdic['d2u_tr']*hexside*_s3
  32
  33     cdef np.ndarray[np.npy_double, ndim = 1] unitcell_envelope, u2d_envelope, d2u_envelope
  34     if gaussianSigma:
  35         sbtime = _time_b(verbose, step='Calculating gaussian envelope')
  36         unitcell_envelope = np.exp(-np.sum(tdic['self_tr']**2,axis=-1)/(2*gaussianSigma**2))
  37         u2d_envelope = np.exp(-np.sum(tdic['u2d_tr']**2,axis=-1)/(2*gaussianSigma**2))
  38         d2u_envelope = np.exp(-np.sum(tdic['d2u_tr']**2,axis=-1)/(2*gaussianSigma**2))
  39         _time_e(sbtime, verbose, step='Calculating gaussian envelope')
  40
  41     cdef np.ndarray[np.npy_cdouble, ndim = 3] svUfullTElist, svVfullTElist, svUfullTMlist, svVfullTMlist
  42     cdef np.ndarray[np.npy_cdouble, ndim = 2] svSfullTElist, svSfullTMlist
  43     cdef np.ndarray[np.npy_double, ndim = 2] minsvTElist, minsvTMlist
  44
  45     #TMatrices_om = TMatrices_interp(omega)
  46     if(not onlyNmin):
  47         svUfullTElist = np.full((klist_c.shape[0], 2*nelem, 2*nelem), np.nan, dtype=complex)
  48         svVfullTElist = np.full((klist_c.shape[0], 2*nelem, 2*nelem), np.nan, dtype=complex)
  49         svSfullTElist = np.full((klist_c.shape[0], 2*nelem), np.nan, dtype=complex)
  50         svUfullTMlist = np.full((klist_c.shape[0], 2*nelem, 2*nelem), np.nan, dtype=complex)
  51         svVfullTMlist = np.full((klist_c.shape[0], 2*nelem, 2*nelem), np.nan, dtype=complex)
  52         svSfullTMlist = np.full((klist_c.shape[0], 2*nelem), np.nan, dtype=complex)
  53     else:
  54         minsvTElist = np.full((klist_c.shape[0], onlyNmin),np.nan)
  55         minsvTMlist = np.full((klist_c.shape[0], onlyNmin),np.nan)
  56
  57     cdef np.ndarray[np.npy_cdouble] leftmatrixlist = np.full((klist_c.shape[0],2,2,nelem,2,2,nelem),np.nan,dtype=complex)
  58     cdef np.ndarray[np.npy_bool, ndim=1] isNaNlist = np.zeros((klist_c.shape[0]), dtype=bool)
  59
  60     sbtime = _time_b(verbose, step='Initialization of matrices for SVD for a given list of k\'s')
  61     # sem nějaká rozumná smyčka
  62
  63     # declarations for the ki loop:
  64     cdef size_t ki
  65     cdef np.ndarray[np.npy_cdouble, ndim = 1] phases_self
  66     cdef np.ndarray[np.npy_cdouble, ndim = 1] phases_u2d
  67     cdef np.ndarray[np.npy_cdouble, ndim = 1] phases_d2u
  68     cdef np.ndarray[np.npy_cdouble, ndim=6] leftmatrix
  69     cdef np.ndarray[np.npy_double, ndim=1] k
  70     cdef int j
  71
  72     for ki in range(klist_c.shape[0]):
  73         k = klist_c[ki]
  74         if (k_0*k_0 - k[0]*k[0] - k[1]*k[1] < 0):
  75             isNaNlist[ki] = True
  76             continue
  77
  78         phases_self = np.exp(1j*np.tensordot(k,unitcell_translations,axes=(0,-1)))
  79         phases_u2d = np.exp(1j*np.tensordot(k,u2d_translations,axes=(0,-1)))
  80         phases_d2u = np.exp(1j*np.tensordot(k,d2u_translations,axes=(0,-1)))
  81         if gaussianSigma:
  82             phases_self *= unitcell_envelope
  83             phases_u2d *= u2d_envelope
  84             phases_d2u *= d2u_envelope
  85         leftmatrix = np.zeros((2,2,nelem, 2,2,nelem), dtype=complex)
  86         #       0:[u,E<--u,E  ]
  87         #       1:[d,M<--d,M  ]
  88         leftmatrix[0,0,:,0,0,:] = np.tensordot(a_self,phases_self, axes=(0,-1)) # u2u, E2E
  89         leftmatrix[1,0,:,1,0,:] = leftmatrix[0,0,:,0,0,:] # d2d, E2E
  90         leftmatrix[0,1,:,0,1,:] = leftmatrix[0,0,:,0,0,:] # u2u, M2M
  91         leftmatrix[1,1,:,1,1,:] = leftmatrix[0,0,:,0,0,:] # d2d, M2M
  92         leftmatrix[0,0,:,0,1,:] = np.tensordot(b_self,phases_self, axes=(0,-1)) # u2u, M2E
  93         leftmatrix[0,1,:,0,0,:] = leftmatrix[0,0,:,0,1,:] # u2u, E2M
  94         leftmatrix[1,1,:,1,0,:] = leftmatrix[0,0,:,0,1,:] # d2d, E2M
  95         leftmatrix[1,0,:,1,1,:] = leftmatrix[0,0,:,0,1,:] # d2d, M2E
  96         leftmatrix[0,0,:,1,0,:] = np.tensordot(a_d2u, phases_d2u,axes=(0,-1)) #d2u,E2E
  97         leftmatrix[0,1,:,1,1,:] = leftmatrix[0,0,:,1,0,:] #d2u, M2M
  98         leftmatrix[1,0,:,0,0,:] = np.tensordot(a_u2d, phases_u2d,axes=(0,-1)) #u2d,E2E
  99         leftmatrix[1,1,:,0,1,:] = leftmatrix[1,0,:,0,0,:] #u2d, M2M
 100         leftmatrix[0,0,:,1,1,:] = np.tensordot(b_d2u, phases_d2u,axes=(0,-1)) #d2u,M2E
 101         leftmatrix[0,1,:,1,0,:] = leftmatrix[0,0,:,1,1,:] #d2u, E2M
 102         leftmatrix[1,0,:,0,1,:] = np.tensordot(b_u2d, phases_u2d,axes=(0,-1)) #u2d,M2E
 103         leftmatrix[1,1,:,0,0,:] = leftmatrix[1,0,:,0,1,:] #u2d, E2M
 104         #leftmatrix is now the translation matrix T
 105         for j in range(2):
 106             leftmatrix[j] = -np.tensordot(TMatrices_om[j], leftmatrix[j], axes=([-2,-1],[0,1]))
 107             # at this point, jth row of leftmatrix is that of -MT
 108             leftmatrix[j,:,:,j,:,:] += n2id
 109         #now we are done, 1-MT
 110
 111         leftmatrixlist[ki] = leftmatrix
 112
 113     nnlist = np.logical_not(isNaNlist)
 114     leftmatrixlist_s = np.reshape(leftmatrixlist,(klist_c.shape[0], 2*2*nelem,2*2*nelem))[nnlist]
 115     TEc, TMc = symz_indexarrays(lMax, 2)
 116     leftmatrixlist_TE = leftmatrixlist_s[np.ix_(np.arange(leftmatrixlist_s.shape[0]),TEc,TEc)]
 117     leftmatrixlist_TM = leftmatrixlist_s[np.ix_(np.arange(leftmatrixlist_s.shape[0]),TMc,TMc)]
 118     _time_e(sbtime, verbose, step='Initializing matrices for SVD for a given list of k\'s')
 119
 120     sbtime = _time_b(verbose, step='Calculating SVDs for a given list of k\'s.')
 121     if(not onlyNmin):
 122         svUfullTElist[nnlist], svSfullTElist[nnlist], svVfullTElist[nnlist] = np.linalg.svd(leftmatrixlist_TE, compute_uv=True)
 123         svUfullTMlist[nnlist], svSfullTMlist[nnlist], svVfullTMlist[nnlist] = np.linalg.svd(leftmatrixlist_TM, compute_uv=True)
 124         _time_e(sbtime, verbose, step='Calculating SVDs for a given list of k\'s.')
 125         return ((svUfullTElist, svSfullTElist, svVfullTElist), (svUfullTMlist, svSfullTMlist, svVfullTMlist))
 126     else:
 127         minsvTElist[nnlist] = np.linalg.svd(leftmatrixlist_TE, compute_uv=False)[...,-onlyNmin:]
 128         minsvTMlist[nnlist] = np.linalg.svd(leftmatrixlist_TM, compute_uv=False)[...,-onlyNmin:]
 129         _time_e(sbtime, verbose, step='Calculating SVDs for a given list of k\'s.')
 130         return (minsvTElist, minsvTMlist)