misc/dispersion-SVD.py

   1 #!/usr/bin/env python3
   2
   3
   4 import argparse, re, random, string
   5 import subprocess
   6 from scipy.constants import hbar, e as eV, pi, c
   7
   8 def make_action_sharedlist(opname, listname):
   9     class opAction(argparse.Action):
  10         def __call__(self, parser, args, values, option_string=None):
  11             if (not hasattr(args, listname)) or getattr(args, listname) is None:
  12                 setattr(args, listname, list())
  13             getattr(args,listname).append((opname, values))
  14     return opAction
  15
  16 parser = argparse.ArgumentParser()
  17 #TODO? použít type=argparse.FileType('r') ?
  18 parser.add_argument('--TMatrix', action='store', required=True, help='Path to TMatrix file')
  19 parser.add_argument('--griddir', action='store', required=True, help='Path to the directory with precalculated translation operators')
  20 #sizepar = parser.add_mutually_exclusive_group(required=True)
  21 parser.add_argument('--hexside', action='store', type=float, required=True, help='Lattice hexagon size length')
  22 parser.add_argument('--output', action='store', help='Path to output PDF')
  23 parser.add_argument('--store_SVD', action='store_true', help='If specified without --SVD_output, it will save the data in a file named as the PDF output, but with .npz extension instead')
  24 #parser.add_argument('--SVD_output', action='store', help='Path to output singular value decomposition result')
  25 parser.add_argument('--nSV', action='store', metavar='N', type=int, default=1, help='Store and draw N minimun singular values')
  26 parser.add_argument('--scp_to', action='store', metavar='N', type=str, help='SCP the output files to a given destination')
  27 parser.add_argument('--background_permittivity', action='store', type=float, default=1., help='Background medium relative permittivity (default 1)')
  28 parser.add_argument('--sparse', action='store', type=int, help='Skip frequencies for preview')
  29 parser.add_argument('--eVmax', action='store', type=float, help='Skip frequencies above this value')
  30 parser.add_argument('--eVmin', action='store', type=float, help='Skip frequencies below this value')
  31 parser.add_argument('--kdensity', action='store', type=int, default=66, help='Number of k-points per x-axis segment')
  32 parser.add_argument('--lMax', action='store', type=int, help='Override lMax from the TMatrix file')
  33 #TODO some more sophisticated x axis definitions
  34 parser.add_argument('--gaussian', action='store', type=float, metavar='σ', help='Use a gaussian envelope for weighting the interaction matrix contributions (depending on the distance), measured in unit cell lengths (?) FIxME).')
  35 popgrp=parser.add_argument_group(title='Operations')
  36 popgrp.add_argument('--tr', dest='ops', action=make_action_sharedlist('tr', 'ops'), default=list()) # the default value for dest can be set once
  37 popgrp.add_argument('--tr0', dest='ops', action=make_action_sharedlist('tr0', 'ops'))
  38 popgrp.add_argument('--tr1', dest='ops', action=make_action_sharedlist('tr1', 'ops'))
  39 popgrp.add_argument('--sym', dest='ops', action=make_action_sharedlist('sym', 'ops'))
  40 popgrp.add_argument('--sym0', dest='ops', action=make_action_sharedlist('sym0', 'ops'))
  41 popgrp.add_argument('--sym1', dest='ops', action=make_action_sharedlist('sym1', 'ops'))
  42 #popgrp.add_argument('--mult', dest='ops', nargs=3, metavar=('INCSPEC', 'SCATSPEC', 'MULTIPLIER'), action=make_action_sharedlist('mult', 'ops'))
  43 #popgrp.add_argument('--mult0', dest='ops', nargs=3, metavar=('INCSPEC', 'SCATSPEC', 'MULTIPLIER'), action=make_action_sharedlist('mult0', 'ops'))
  44 #popgrp.add_argument('--mult1', dest='ops', nargs=3, metavar=('INCSPEC', 'SCATSPEC', 'MULTIPLIER'), action=make_action_sharedlist('mult1', 'ops'))
  45 popgrp.add_argument('--multl', dest='ops', nargs=3, metavar=('INCL[,INCL,...]', 'SCATL[,SCATL,...]', 'MULTIPLIER'), action=make_action_sharedlist('multl', 'ops'))
  46 popgrp.add_argument('--multl0', dest='ops', nargs=3, metavar=('INCL[,INCL,...]', 'SCATL[,SCATL,...]', 'MULTIPLIER'), action=make_action_sharedlist('multl0', 'ops'))
  47 popgrp.add_argument('--multl1', dest='ops', nargs=3, metavar=('INCL[,INCL,...]', 'SCATL[,SCATL,...]', 'MULTIPLIER'), action=make_action_sharedlist('multl1', 'ops'))
  48 parser.add_argument('--frequency_multiplier', action='store', type=float, default=1., help='Multiplies the frequencies in the TMatrix file by a given factor.')
  49 # TODO enable more flexible per-sublattice specification
  50 pargs=parser.parse_args()
  51 print(pargs)
  52
  53
  54 translations_dir = pargs.griddir
  55 TMatrix_file = pargs.TMatrix
  56 pdfout = pargs.output if pargs.output else (''.join(random.choice(string.ascii_uppercase + string.digits) for _ in range(10)) + '.pdf')
  57 print(pdfout)
  58 if(pargs.store_SVD):
  59     if re.search('.pdf$', pdfout):
  60         svdout = re.sub('.pdf$', r'.npz', pdfout)
  61     else:
  62         svdout = pdfout + '.npz'
  63 else:
  64      svdout = None
  65
  66
  67 hexside = pargs.hexside #375e-9
  68 epsilon_b = pargs.background_permittivity #2.3104
  69 gaussianSigma = pargs.gaussian if pargs.gaussian else None # hexside * 222 / 7
  70 interpfreqfactor = pargs.frequency_multiplier
  71 scp_dest = pargs.scp_to if pargs.scp_to else None
  72 kdensity = pargs.kdensity
  73 minfreq = pargs.eVmin*eV/hbar if pargs.eVmin else None
  74 maxfreq = pargs.eVmax*eV/hbar if pargs.eVmax else None
  75 skipfreq = pargs.sparse if pargs.sparse else None
  76 svn = pargs.nSV
  77
  78 # TODO multiplier operation definitions and parsing
  79 #factor13inc = 10
  80 #factor13scat=10
  81
  82 ops = list()
  83 opre = re.compile('(tr|sym|copy|multl|mult)(\d*)')
  84 for oparg in pargs.ops:
  85     opm = opre.match(oparg[0])
  86     if opm:
  87         ops.append(((opm.group(2),) if opm.group(2) else (0,1), opm.group(1), oparg[1]))
  88     else:
  89         raise # should not happen
  90 print(ops)
  91
  92 #ops = (
  93 #    # co, typ operace (symetrizace / transformace / kopie), specifikace (operace nebo zdroj),
  94 #    # co: 0, 1, (0,1), (0,), (1,), #NI: 'all'
  95 #    # typ operace: sym, tr, copy
  96 #    # specifikace:
  97 #    # sym, tr: 'σ_z', 'σ_y', 'C2'; sym: 'C3',
  98 #    # copy: 0, 1 (zdroj)
  99 #    ((0,1), 'sym', 'σ_z'),
 100 #    #((0,1), 'sym', 'σ_x'),
 101 #    #((0,1), 'sym', 'σ_y'),
 102 #    ((0,1), 'sym', 'C3'),
 103 #    ((1), 'tr', 'C2'),
 104 #
 105 #)
 106
 107 # -----------------finished basic CLI parsing (except for op arguments) ------------------
 108 import time
 109 begtime=time.time()
 110
 111 from matplotlib.path import Path
 112 import matplotlib.patches as patches
 113 import matplotlib.pyplot as plt
 114 import qpms
 115 import numpy as np
 116 import os, sys, warnings, math
 117 from matplotlib import pyplot as plt
 118 from matplotlib.backends.backend_pdf import PdfPages
 119 from scipy import interpolate
 120 nx = None
 121 s3 = math.sqrt(3)
 122
 123
 124
 125 pdf = PdfPages(pdfout)
 126
 127 # In[3]:
 128
 129 # specifikace T-matice zde
 130 cdn = c/ math.sqrt(epsilon_b)
 131 TMatrices_orig, freqs_orig, freqs_weirdunits_orig, lMaxTM = qpms.loadScuffTMatrices(TMatrix_file)
 132 if pargs.lMax:
 133     lMax = pargs.lMax if pargs.lMax else lMaxTM
 134 my, ny = qpms.get_mn_y(lMax)
 135 nelem = len(my)
 136 if pargs.lMax: #force commandline specified lMax
 137     TMatrices_orig = TMatrices_orig[...,0:nelem,:,0:nelem]
 138
 139 ž = np.arange(2*nelem)
 140 tž = ž // nelem
 141 mž = my[ž%nelem]
 142 nž = ny[ž%nelem]
 143 TEž = ž[(mž+nž+tž) % 2 == 0]
 144 TMž = ž[(mž+nž+tž) % 2 == 1]
 145
 146 č = np.arange(2*2*nelem)
 147 žč = č % (2* nelem)
 148 tč = tž[žč]
 149 mč = mž[žč]
 150 nč = nž[žč]
 151 TEč = č[(mč+nč+tč) % 2 == 0]
 152 TMč = č[(mč+nč+tč) % 2 == 1]
 153
 154 TMatrices = np.array(np.broadcast_to(TMatrices_orig[:,nx,:,:,:,:],(len(freqs_orig),2,2,nelem,2,nelem)) )
 155
 156 #TMatrices[:,:,:,:,:,ny==3] *= factor13inc
 157 #TMatrices[:,:,:,ny==3,:,:] *= factor13scat
 158 xfl = qpms.xflip_tyty(lMax)
 159 yfl = qpms.yflip_tyty(lMax)
 160 zfl = qpms.zflip_tyty(lMax)
 161 c2rot = qpms.apply_matrix_left(qpms.yflip_yy(3),qpms.xflip_yy(3),-1)
 162
 163 reCN = re.compile('(\d*)C(\d+)')
 164 #TODO C nekonečno
 165
 166 for op in ops:
 167     if op[0] == 'all':
 168         targets = (0,1)
 169     elif isinstance(op[0],int):
 170         targets = (op[0],)
 171     else:
 172         targets = op[0]
 173
 174     if op[1] == 'sym':
 175         mCN = reCN.match(op[2]) # Fuck van Rossum for not having assignments inside expressions
 176         if op[2] == 'σ_z':
 177             for t in targets:
 178                 TMatrices[:,t] = (TMatrices[:,t] + qpms.apply_ndmatrix_left(zfl,qpms.apply_ndmatrix_left(zfl, TMatrices[:,t], (-4,-3)),(-2,-1)))/2
 179         elif op[2] == 'σ_y':
 180             for t in targets:
 181                 TMatrices[:,t] = (TMatrices[:,t] + qpms.apply_ndmatrix_left(yfl,qpms.apply_ndmatrix_left(yfl, TMatrices[:,t], (-4,-3)),(-2,-1)))/2
 182         elif op[2] == 'σ_x':
 183             for t in targets:
 184                 TMatrices[:,t] = (TMatrices[:,t] + qpms.apply_ndmatrix_left(xfl,qpms.apply_ndmatrix_left(xfl, TMatrices[:,t], (-4,-3)),(-2,-1)))/2
 185         elif op[2] == 'C2': # special case of the latter
 186             for t in targets:
 187                 TMatrices[:,t] = (TMatrices[:,t] + qpms.apply_matrix_left(c2rot,qpms.apply_matrix_left(c2rot, TMatrices[:,t], -3),-1))/2
 188         elif mCN:
 189             rotN = int(mCN.group(2))
 190             TMatrix_contribs = np.empty((rotN,TMatrices.shape[0],2,nelem,2,nelem), dtype=np.complex_)
 191             for t in targets:
 192                 for i in range(rotN):
 193                     rotangle = 2*np.pi*i / rotN
 194                     rot =  qpms.WignerD_yy_fromvector(lMax,np.array([0,0,rotangle]))
 195                     rotinv = qpms.WignerD_yy_fromvector(lMax,np.array([0,0,-rotangle]))
 196                     TMatrix_contribs[i] = qpms.apply_matrix_left(rot,qpms.apply_matrix_left(rotinv, TMatrices[:,t], -3),-1)
 197                 TMatrices[:,t] = np.sum(TMatrix_contribs, axis=0) / rotN
 198         else:
 199             raise
 200     elif op[1] == 'tr':
 201         mCN = reCN.match(op[2]) # Fuck van Rossum for not having assignments inside expressions
 202         if op[2] == 'σ_z':
 203             for t in targets:
 204                 TMatrices[:,t] = qpms.apply_ndmatrix_left(zfl,qpms.apply_ndmatrix_left(zfl, TMatrices[:,t], (-4,-3)),(-2,-1))
 205         elif op[2] == 'σ_y':
 206             for t in targets:
 207                 TMatrices[:,t] = qpms.apply_ndmatrix_left(yfl,qpms.apply_ndmatrix_left(yfl, TMatrices[:,t], (-4,-3)),(-2,-1))
 208         elif op[2] == 'σ_x':
 209             for t in targets:
 210                 TMatrices[:,t] = qpms.apply_ndmatrix_left(xfl,qpms.apply_ndmatrix_left(xfl, TMatrices[:,t], (-4,-3)),(-2,-1))
 211         elif op[2] == 'C2':
 212             for t in targets:
 213                 TMatrices[:,t] = qpms.apply_matrix_left(c2rot,qpms.apply_matrix_left(c2rot, TMatrices[:,t], -3),-1)
 214         elif mCN:
 215             rotN = int(mCN.group(2))
 216             power = int(mCN.group(1)) if mCN.group(1) else 1
 217             TMatrix_contribs = np.empty((rotN,TMatrices.shape[0],2,nelem,2,nelem), dtype=np.complex_)
 218             for t in targets:
 219                 rotangle = 2*np.pi*power/rotN
 220                 rot = qpms.WignerD_yy_fromvector(lMax, np.array([0,0,rotangle]))
 221                 rotinv = qpms.WignerD_yy_fromvector(lMax, np.array([0,0,-rotangle]))
 222                 TMatrices[:,t] = qpms.apply_matrix_left(rot, qpms.apply_matrix_left(rotinv, TMatrices[:,t], -3),-1)
 223         else:
 224             raise
 225     elif op[1] == 'copy':
 226         raise # not implemented
 227     elif op[1] == 'mult':
 228         raise # not implemented
 229     elif op[1] == 'multl':
 230         incy = np.full((nelem,), False, dtype=bool)
 231         for incl in op[2][0].split(','):
 232             l = int(incl)
 233             incy += (l == ny)
 234         scaty = np.full((nelem,), False, dtype=bool)
 235         for scatl in op[2][1].split(','):
 236             l = int(scatl)
 237             scaty += (l == ny)
 238         for t in targets:
 239             TMatrices[np.ix_(np.arange(TMatrices.shape[0]),np.array([t]),np.array([0,1]),scaty,np.array([0,1]),incy)] *= float(op[2][2])
 240     else:
 241         raise #unknown operation; should not happen
 242
 243 TMatrices_interp = interpolate.interp1d(freqs_orig*interpfreqfactor, TMatrices, axis=0, kind='linear',fill_value="extrapolate")
 244
 245
 246
 247 # In[4]:
 248
 249 om = np.linspace(np.min(freqs_orig), np.max(freqs_orig),100)
 250 TMatrix0ip = np.reshape(TMatrices_interp(om)[:,0], (len(om), 2*nelem*2*nelem))
 251 f, axa = plt.subplots(2, 2, figsize=(15,15))
 252
 253 #print(TMatrices.shape)
 254 #plt.plot(om, TMatrices[:,0,0,0,0].imag,'r',om, TMatrices[:,0,0,0,0].real,'r--',om, TMatrices[:,0,2,0,2].imag,'b',om, TMatrices[:,0,2,0,2].real,'b--'))
 255
 256 ax = axa[0,0]
 257 ax2 = ax.twiny()
 258 ax2.set_xlim([ax.get_xlim()[0]/eV*hbar,ax.get_xlim()[1]/eV*hbar])
 259 ax.plot(
 260          om, TMatrix0ip[:,:].imag,'-',om, TMatrix0ip[:,:].real,'--',
 261 )
 262 ax = axa[0,1]
 263 ax2 = ax.twiny()
 264 ax2.set_xlim([ax.get_xlim()[0]/eV*hbar,ax.get_xlim()[1]/eV*hbar])
 265 ax.plot(
 266          om, abs(TMatrix0ip[:,:]),'-'
 267 )
 268 ax.set_yscale('log')
 269
 270 ax = axa[1,1]
 271 ax2 = ax.twiny()
 272 ax2.set_xlim([ax.get_xlim()[0]/eV*hbar,ax.get_xlim()[1]/eV*hbar])
 273 ax.plot(
 274          om, np.unwrap(np.angle(TMatrix0ip[:,:]),axis=0),'-'
 275 )
 276
 277 ax = axa[1,0]
 278 ax.text(0.5,0.5,str(pargs).replace(',',',\n'),horizontalalignment='center',verticalalignment='center',transform=ax.transAxes)
 279 pdf.savefig(f)
 280
 281
 282 # In[ ]:
 283
 284 #kdensity = 66 #defined from cl arguments
 285 bz_0 = np.array((0,0,0.,))
 286 bz_K1 = np.array((1.,0,0))*4*np.pi/3/hexside/s3
 287 bz_K2 = np.array((1./2.,s3/2,0))*4*np.pi/3/hexside/s3
 288 bz_M = np.array((3./4, s3/4,0))*4*np.pi/3/hexside/s3
 289 k0Mlist = bz_0 + (bz_M-bz_0) * np.linspace(0,1,kdensity)[:,nx]
 290 kMK1list = bz_M + (bz_K1-bz_M) * np.linspace(0,1,kdensity)[:,nx]
 291 kK10list = bz_K1 + (bz_0-bz_K1) * np.linspace(0,1,kdensity)[:,nx]
 292 k0K2list = bz_0 + (bz_K2-bz_0) * np.linspace(0,1,kdensity)[:,nx]
 293 kK2Mlist = bz_K2 + (bz_M-bz_K2) * np.linspace(0,1,kdensity)[:,nx]
 294 B1 = 2* bz_K1 - bz_K2
 295 B2 = 2* bz_K2 - bz_K1
 296 klist = np.concatenate((k0Mlist,kMK1list,kK10list,k0K2list,kK2Mlist), axis=0)
 297 kxmaplist = np.concatenate((np.array([0]),np.cumsum(np.linalg.norm(np.diff(klist, axis=0), axis=-1))))
 298
 299
 300 # In[ ]:
 301
 302 n2id = np.identity(2*nelem)
 303 n2id.shape = (2,nelem,2,nelem)
 304 extlistlist = list()
 305 leftmatrixlistlist = list()
 306 minsvTElistlist=list()
 307 minsvTMlistlist=list()
 308 if svdout:
 309     svUfullTElistlist = list()
 310     svVfullTElistlist = list()
 311     svSfullTElistlist = list()
 312     svUfullTMlistlist = list()
 313     svVfullTMlistlist = list()
 314     svSfullTMlistlist = list()
 315
 316 nan = float('nan')
 317 omegalist = list()
 318 filecount = 0
 319 for trfile in os.scandir(translations_dir):
 320     filecount += 1
 321     if (skipfreq and filecount % skipfreq):
 322         continue
 323     try:
 324         npz = np.load(trfile.path, mmap_mode='r')
 325         k_0 = npz['precalc_params'][()]['k_hexside'] / hexside
 326         omega = k_0 * c / math.sqrt(epsilon_b)
 327         if((minfreq and omega < minfreq) or (maxfreq and omega > maxfreq)):
 328             continue
 329     except:
 330         print ("Unexpected error, trying to continue with another file:", sys.exc_info()[0])
 331         continue
 332     try:
 333         tdic = qpms.hexlattice_precalc_AB_loadunwrap(trfile.path, return_points=True)
 334     except:
 335         print ("Unexpected error, trying to continue with another file:", sys.exc_info()[0])
 336         continue
 337     k_0 = tdic['k_hexside'] / hexside
 338     omega = k_0 * c / math.sqrt(epsilon_b)
 339     omegalist.append(omega)
 340     print(filecount, omega/eV*hbar)
 341     sys.stdout.flush()
 342     a_self = tdic['a_self'][:,:nelem,:nelem]
 343     b_self = tdic['b_self'][:,:nelem,:nelem]
 344     a_u2d = tdic['a_u2d'][:,:nelem,:nelem]
 345     b_u2d = tdic['b_u2d'][:,:nelem,:nelem]
 346     a_d2u = tdic['a_d2u'][:,:nelem,:nelem]
 347     b_d2u = tdic['b_d2u'][:,:nelem,:nelem]
 348     unitcell_translations = tdic['self_tr']*hexside*s3
 349     u2d_translations = tdic['u2d_tr']*hexside*s3
 350     d2u_translations = tdic['d2u_tr']*hexside*s3
 351     if gaussianSigma:
 352         unitcell_envelope = np.exp(-np.sum(tdic['self_tr']**2,axis=-1)/(2*gaussianSigma**2))
 353         u2d_envelope = np.exp(-np.sum(tdic['u2d_tr']**2,axis=-1)/(2*gaussianSigma**2))
 354         d2u_envelope = np.exp(-np.sum(tdic['d2u_tr']**2,axis=-1)/(2*gaussianSigma**2))
 355
 356
 357     TMatrices_om = TMatrices_interp(omega)
 358     if svdout:
 359         svUfullTElist = np.full((klist.shape[0], 2*nelem, 2*nelem), np.nan, dtype=complex)
 360         svVfullTElist = np.full((klist.shape[0], 2*nelem, 2*nelem), np.nan, dtype=complex)
 361         svSfullTElist = np.full((klist.shape[0], 2*nelem), np.nan, dtype=complex)
 362         svUfullTMlist = np.full((klist.shape[0], 2*nelem, 2*nelem), np.nan, dtype=complex)
 363         svVfullTMlist = np.full((klist.shape[0], 2*nelem, 2*nelem), np.nan, dtype=complex)
 364         svSfullTMlist = np.full((klist.shape[0], 2*nelem), np.nan, dtype=complex)
 365
 366
 367     minsvTElist = np.full((klist.shape[0], svn),np.nan)
 368     minsvTMlist = np.full((klist.shape[0], svn),np.nan)
 369     leftmatrixlist = np.full((klist.shape[0],2,2,nelem,2,2,nelem),np.nan,dtype=complex)
 370     isNaNlist = np.zeros((klist.shape[0]), dtype=bool)
 371
 372     # sem nějaká rozumná smyčka
 373     for ki in range(klist.shape[0]):
 374         k = klist[ki]
 375         if (k_0*k_0 - k[0]*k[0] - k[1]*k[1] < 0):
 376             isNaNlist[ki] = True
 377             continue
 378
 379         phases_self = np.exp(1j*np.tensordot(k,unitcell_translations,axes=(0,-1)))
 380         phases_u2d = np.exp(1j*np.tensordot(k,u2d_translations,axes=(0,-1)))
 381         phases_d2u = np.exp(1j*np.tensordot(k,d2u_translations,axes=(0,-1)))
 382         if gaussianSigma:
 383             phases_self *= unitcell_envelope
 384             phases_u2d *= u2d_envelope
 385             phases_d2u *= d2u_envelope
 386         leftmatrix = np.zeros((2,2,nelem, 2,2,nelem), dtype=complex)
 387
 388         leftmatrix[0,0,:,0,0,:] = np.tensordot(a_self,phases_self, axes=(0,-1)) # u2u, E2E
 389         leftmatrix[1,0,:,1,0,:] = leftmatrix[0,0,:,0,0,:] # d2d, E2E
 390         leftmatrix[0,1,:,0,1,:] = leftmatrix[0,0,:,0,0,:] # u2u, M2M
 391         leftmatrix[1,1,:,1,1,:] = leftmatrix[0,0,:,0,0,:] # d2d, M2M
 392         leftmatrix[0,0,:,0,1,:] = np.tensordot(b_self,phases_self, axes=(0,-1)) # u2u, M2E
 393         leftmatrix[0,1,:,0,0,:] = leftmatrix[0,0,:,0,1,:] # u2u, E2M
 394         leftmatrix[1,1,:,1,0,:] = leftmatrix[0,0,:,0,1,:] # d2d, E2M
 395         leftmatrix[1,0,:,1,1,:] = leftmatrix[0,0,:,0,1,:] # d2d, M2E
 396         leftmatrix[0,0,:,1,0,:] = np.tensordot(a_d2u, phases_d2u,axes=(0,-1)) #d2u,E2E
 397         leftmatrix[0,1,:,1,1,:] = leftmatrix[0,0,:,1,0,:] #d2u, M2M
 398         leftmatrix[1,0,:,0,0,:] = np.tensordot(a_u2d, phases_u2d,axes=(0,-1)) #u2d,E2E
 399         leftmatrix[1,1,:,0,1,:] = leftmatrix[1,0,:,0,0,:] #u2d, M2M
 400         leftmatrix[0,0,:,1,1,:] = np.tensordot(b_d2u, phases_d2u,axes=(0,-1)) #d2u,M2E
 401         leftmatrix[0,1,:,1,0,:] = leftmatrix[0,0,:,1,1,:] #d2u, E2M
 402         leftmatrix[1,0,:,0,1,:] = np.tensordot(b_u2d, phases_u2d,axes=(0,-1)) #u2d,M2E
 403         leftmatrix[1,1,:,0,0,:] = leftmatrix[1,0,:,0,1,:] #u2d, E2M
 404         #leftmatrix is now the translation matrix T
 405         for j in range(2):
 406             leftmatrix[j] = -np.tensordot(TMatrices_om[j], leftmatrix[j], axes=([-2,-1],[0,1]))
 407             # at this point, jth row of leftmatrix is that of -MT
 408             leftmatrix[j,:,:,j,:,:] += n2id
 409         #now we are done, 1-MT
 410
 411         leftmatrixlist[ki] = leftmatrix
 412
 413
 414     nnlist = np.logical_not(isNaNlist)
 415     leftmatrixlist_s = np.reshape(leftmatrixlist,(klist.shape[0], 2*2*nelem,2*2*nelem))[nnlist]
 416     leftmatrixlist_TE = leftmatrixlist_s[np.ix_(np.arange(leftmatrixlist_s.shape[0]),TEč,TEč)]
 417     leftmatrixlist_TM = leftmatrixlist_s[np.ix_(np.arange(leftmatrixlist_s.shape[0]),TMč,TMč)]
 418     #svarr = np.linalg.svd(leftmatrixlist_TE, compute_uv=False)
 419     #argsortlist = np.argsort(svarr, axis=-1)[...,:svn]
 420     #minsvTElist[nnlist] = svarr[...,argsortlist]
 421     #minsvTElist[nnlist] = np.amin(np.linalg.svd(leftmatrixlist_TE, compute_uv=False), axis=-1)
 422     if svdout:
 423         svUfullTElist[nnlist], svSfullTElist[nnlist], svVfullTElist[nnlist] = np.linalg.svd(leftmatrixlist_TE, compute_uv=True)
 424         svUfullTMlist[nnlist], svSfullTMlist[nnlist], svVfullTMlist[nnlist] = np.linalg.svd(leftmatrixlist_TM, compute_uv=True)
 425         svUfullTElistlist.append(svUfullTElist)
 426         svVfullTElistlist.append(svVfullTElist)
 427         svSfullTElistlist.append(svSfullTElist)
 428         svUfullTMlistlist.append(svUfullTMlist)
 429         svVfullTMlistlist.append(svVfullTMlist)
 430         svSfullTMlistlist.append(svSfullTMlist)
 431     minsvTElist[nnlist] = np.linalg.svd(leftmatrixlist_TE, compute_uv=False)[...,-svn:]
 432     #svarr = np.linalg.svd(leftmatrixlist_TM, compute_uv=False)
 433     #argsortlist = np.argsort(svarr, axis=-1)[...,:svn]
 434     #minsvTMlist[nnlist] = svarr[...,argsortlist]
 435     #minsvTMlist[nnlist] = np.amin(np.linalg.svd(leftmatrixlist_TM, compute_uv=False), axis=-1)
 436     minsvTMlist[nnlist] = np.linalg.svd(leftmatrixlist_TM, compute_uv=False)[...,-svn:]
 437     minsvTMlistlist.append(minsvTMlist)
 438     minsvTElistlist.append(minsvTElist)
 439
 440 minsvTElistarr = np.array(minsvTElistlist)
 441 minsvTMlistarr = np.array(minsvTMlistlist)
 442 del minsvTElistlist, minsvTMlistlist
 443 if svdout:
 444     svUfullTElistarr = np.array(svUfullTElistlist)
 445     svVfullTElistarr = np.array(svVfullTElistlist)
 446     svSfullTElistarr = np.array(svSfullTElistlist)
 447     del svUfullTElistlist, svVfullTElistlist, svSfullTElistlist
 448     svUfullTMlistarr = np.array(svUfullTMlistlist)
 449     svVfullTMlistarr = np.array(svVfullTMlistlist)
 450     svSfullTMlistarr = np.array(svSfullTMlistlist)
 451     del svUfullTMlistlist, svVfullTMlistlist, svSfullTMlistlist
 452 omegalist = np.array(omegalist)
 453 # order to make the scatter plots "nice"
 454 omegaorder = np.argsort(omegalist)
 455 omegalist = omegalist[omegaorder]
 456 minsvTElistarr = minsvTElistarr[omegaorder]
 457 minsvTMlistarr = minsvTMlistarr[omegaorder]
 458 if svdout:
 459     svUfullTElistarr = svUfullTElistarr[omegaorder]
 460     svVfullTElistarr = svVfullTElistarr[omegaorder]
 461     svSfullTElistarr = svSfullTElistarr[omegaorder]
 462     svUfullTMlistarr = svUfullTMlistarr[omegaorder]
 463     svVfullTMlistarr = svVfullTMlistarr[omegaorder]
 464     svSfullTMlistarr = svSfullTMlistarr[omegaorder]
 465     np.savez(svdout, omega = omegalist, klist = klist, bzpoints = np.array([bz_0, bz_K1, bz_K2, bz_M, B1, B2]),
 466                      uTE = svUfullTElistarr,
 467                      vTE = svVfullTElistarr,
 468                      sTE = svSfullTElistarr,
 469                      uTM = svUfullTMlistarr,
 470                      vTM = svVfullTMlistarr,
 471                      sTM = svSfullTMlistarr,
 472     )
 473
 474
 475 omlist = np.broadcast_to(omegalist[:,nx], minsvTElistarr[...,0].shape)
 476 kxmlarr = np.broadcast_to(kxmaplist[nx,:], minsvTElistarr[...,0].shape)
 477 klist = np.concatenate((k0Mlist,kMK1list,kK10list,k0K2list,kK2Mlist), axis=0)
 478
 479
 480 # In[ ]:
 481 for minN in reversed(range(svn)):
 482     f, ax = plt.subplots(1, figsize=(20,15))
 483     sc = ax.scatter(kxmlarr, omlist/eV*hbar, c = np.sqrt(minsvTMlistarr[...,minN]), s =40, lw=0)
 484     ax.plot(kxmaplist, np.linalg.norm(klist,axis=-1)*cdn/eV*hbar, '-',
 485            kxmaplist, np.linalg.norm(klist+B1, axis=-1)*cdn/eV*hbar, '-',
 486            kxmaplist, np.linalg.norm(klist+B2, axis=-1)*cdn/eV*hbar, '-',
 487            kxmaplist, np.linalg.norm(klist-B2, axis=-1)*cdn/eV*hbar, '-',
 488            kxmaplist, np.linalg.norm(klist-B1, axis=-1)*cdn/eV*hbar, '-',
 489            kxmaplist, np.linalg.norm(klist+B2-B1, axis=-1)*cdn/eV*hbar, '-',
 490             kxmaplist, np.linalg.norm(klist-B2+B1, axis=-1)*cdn/eV*hbar, '-',
 491             kxmaplist, np.linalg.norm(klist-B2-B1, axis=-1)*cdn/eV*hbar, '-',
 492             kxmaplist, np.linalg.norm(klist+B2+B1, axis=-1)*cdn/eV*hbar, '-',
 493             kxmaplist, np.linalg.norm(klist-2*B1, axis=-1)*cdn/eV*hbar, '-',
 494             kxmaplist, np.linalg.norm(klist-2*B2, axis=-1)*cdn/eV*hbar, '-',
 495             kxmaplist, np.linalg.norm(klist-2*B2-B1, axis=-1)*cdn/eV*hbar, '-',
 496             kxmaplist, np.linalg.norm(klist-2*B1-B2, axis=-1)*cdn/eV*hbar, '-',
 497             kxmaplist, np.linalg.norm(klist-2*B1-2*B2, axis=-1)*cdn/eV*hbar, '-',
 498     #        kxmaplist, np.linalg.norm(klist+2*B2-B1, axis=-1)*cdn, '-',
 499     #        kxmaplist, np.linalg.norm(klist+2*B1-B2, axis=-1)*cdn, '-',
 500            )
 501     ax.set_xlim([np.min(kxmlarr),np.max(kxmlarr)])
 502     #ax.set_ylim([2.15,2.30])
 503     ax.set_ylim([np.min(omlist/eV*hbar),np.max(omlist/eV*hbar)])
 504     ax.set_xticks([0, kxmaplist[len(k0Mlist)-1], kxmaplist[len(k0Mlist)+len(kMK1list)-1], kxmaplist[len(k0Mlist)+len(kMK1list)+len(kK10list)-1], kxmaplist[len(k0Mlist)+len(kMK1list)+len(kK10list)+len(k0K2list)-1], kxmaplist[len(k0Mlist)+len(kMK1list)+len(kK10list)+len(k0K2list)+len(kK2Mlist)-1]])
 505     ax.set_xticklabels(['Γ', 'M', 'K', 'Γ', 'K\'','M'])
 506     f.colorbar(sc)
 507
 508     pdf.savefig(f)
 509
 510
 511     # In[ ]:
 512
 513     f, ax = plt.subplots(1, figsize=(20,15))
 514     sc = ax.scatter(kxmlarr, omlist/eV*hbar, c = np.sqrt(minsvTElistarr[...,minN]), s =40, lw=0)
 515     ax.plot(kxmaplist, np.linalg.norm(klist,axis=-1)*cdn/eV*hbar, '-',
 516            kxmaplist, np.linalg.norm(klist+B1, axis=-1)*cdn/eV*hbar, '-',
 517            kxmaplist, np.linalg.norm(klist+B2, axis=-1)*cdn/eV*hbar, '-',
 518            kxmaplist, np.linalg.norm(klist-B2, axis=-1)*cdn/eV*hbar, '-',
 519            kxmaplist, np.linalg.norm(klist-B1, axis=-1)*cdn/eV*hbar, '-',
 520            kxmaplist, np.linalg.norm(klist+B2-B1, axis=-1)*cdn/eV*hbar, '-',
 521             kxmaplist, np.linalg.norm(klist-B2+B1, axis=-1)*cdn/eV*hbar, '-',
 522             kxmaplist, np.linalg.norm(klist-B2-B1, axis=-1)*cdn/eV*hbar, '-',
 523             kxmaplist, np.linalg.norm(klist+B2+B1, axis=-1)*cdn/eV*hbar, '-',
 524             kxmaplist, np.linalg.norm(klist-2*B1, axis=-1)*cdn/eV*hbar, '-',
 525             kxmaplist, np.linalg.norm(klist-2*B2, axis=-1)*cdn/eV*hbar, '-',
 526             kxmaplist, np.linalg.norm(klist-2*B2-B1, axis=-1)*cdn/eV*hbar, '-',
 527             kxmaplist, np.linalg.norm(klist-2*B1-B2, axis=-1)*cdn/eV*hbar, '-',
 528             kxmaplist, np.linalg.norm(klist-2*B1-2*B2, axis=-1)*cdn/eV*hbar, '-',
 529     #        kxmaplist, np.linalg.norm(klist+2*B2-B1, axis=-1)*cdn, '-',
 530     #        kxmaplist, np.linalg.norm(klist+2*B1-B2, axis=-1)*cdn, '-',
 531            )
 532     ax.set_xlim([np.min(kxmlarr),np.max(kxmlarr)])
 533     #ax.set_ylim([2.15,2.30])
 534     ax.set_ylim([np.min(omlist/eV*hbar),np.max(omlist/eV*hbar)])
 535     ax.set_xticks([0, kxmaplist[len(k0Mlist)-1], kxmaplist[len(k0Mlist)+len(kMK1list)-1], kxmaplist[len(k0Mlist)+len(kMK1list)+len(kK10list)-1], kxmaplist[len(k0Mlist)+len(kMK1list)+len(kK10list)+len(k0K2list)-1], kxmaplist[len(k0Mlist)+len(kMK1list)+len(kK10list)+len(k0K2list)+len(kK2Mlist)-1]])
 536     ax.set_xticklabels(['Γ', 'M', 'K', 'Γ', 'K\'','M'])
 537     f.colorbar(sc)
 538
 539     pdf.savefig(f)
 540 pdf.close()
 541
 542 if scp_dest:
 543     subprocess.run(['scp', pdfout, scp_dest])
 544     if svdout:
 545         subprocess.run(['scp', svdout, scp_dest])
 546
 547 print(time.strftime("%H.%M:%S",time.gmtime(time.time()-begtime)))