misc/dispersion_hex_chunks.py

   1 #!/usr/bin/env python3
   2
   3 import argparse, re, random, string
   4 import subprocess
   5 from scipy.constants import hbar, e as eV, pi, c
   6
   7 def make_action_sharedlist(opname, listname):
   8     class opAction(argparse.Action):
   9         def __call__(self, parser, args, values, option_string=None):
  10             if (not hasattr(args, listname)) or getattr(args, listname) is None:
  11                 setattr(args, listname, list())
  12             getattr(args,listname).append((opname, values))
  13     return opAction
  14
  15 parser = argparse.ArgumentParser()
  16 #TODO? použít type=argparse.FileType('r') ?
  17 parser.add_argument('--TMatrix', action='store', required=True, help='Path to TMatrix file')
  18 #parser.add_argument('--griddir', action='store', required=True, help='Path to the directory with precalculated translation operators')
  19 parser.add_argument('--output_prefix', action='store', required=True, help='Prefix to the npz output (will be appended frequency, hexside and chunkno)')
  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('--plot_TMatrix', action='store_true', help='Visualise TMatrix on the first page of the output')
  23 #parser.add_argument('--SVD_output', action='store', help='Path to output singular value decomposition result')
  24 parser.add_argument('--maxlayer', action='store', type=int, default=100, help='How far to sum the lattice points to obtain the dispersion')
  25 parser.add_argument('--scp_to', action='store', metavar='N', type=str, help='SCP the output files to a given destination')
  26 parser.add_argument('--background_permittivity', action='store', type=float, default=1., help='Background medium relative permittivity (default 1)')
  27 parser.add_argument('--eVfreq', action='store', required=True, type=float, help='Frequency in eV')
  28 parser.add_argument('--kdensity', action='store', type=int, default=33, help='Number of k-points per x-axis segment')
  29 parser.add_argument('--chunklen', action='store', type=int, default=1000, help='Number of k-points per output file (default 1000)')
  30 parser.add_argument('--lMax', action='store', type=int, help='Override lMax from the TMatrix file')
  31 #TODO some more sophisticated x axis definitions
  32 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).')
  33 parser.add_argument('--verbose', '-v', action='count', help='Be verbose (about computation times, mostly)')
  34 popgrp=parser.add_argument_group(title='Operations')
  35 popgrp.add_argument('--tr', dest='ops', action=make_action_sharedlist('tr', 'ops'), default=list()) # the default value for dest can be set once
  36 popgrp.add_argument('--tr0', dest='ops', action=make_action_sharedlist('tr0', 'ops'))
  37 popgrp.add_argument('--tr1', dest='ops', action=make_action_sharedlist('tr1', 'ops'))
  38 popgrp.add_argument('--sym', dest='ops', action=make_action_sharedlist('sym', 'ops'))
  39 popgrp.add_argument('--sym0', dest='ops', action=make_action_sharedlist('sym0', 'ops'))
  40 popgrp.add_argument('--sym1', dest='ops', action=make_action_sharedlist('sym1', 'ops'))
  41 #popgrp.add_argument('--mult', dest='ops', nargs=3, metavar=('INCSPEC', 'SCATSPEC', 'MULTIPLIER'), action=make_action_sharedlist('mult', 'ops'))
  42 #popgrp.add_argument('--mult0', dest='ops', nargs=3, metavar=('INCSPEC', 'SCATSPEC', 'MULTIPLIER'), action=make_action_sharedlist('mult0', 'ops'))
  43 #popgrp.add_argument('--mult1', dest='ops', nargs=3, metavar=('INCSPEC', 'SCATSPEC', 'MULTIPLIER'), action=make_action_sharedlist('mult1', 'ops'))
  44 popgrp.add_argument('--multl', dest='ops', nargs=3, metavar=('INCL[,INCL,...]', 'SCATL[,SCATL,...]', 'MULTIPLIER'), action=make_action_sharedlist('multl', 'ops'))
  45 popgrp.add_argument('--multl0', dest='ops', nargs=3, metavar=('INCL[,INCL,...]', 'SCATL[,SCATL,...]', 'MULTIPLIER'), action=make_action_sharedlist('multl0', 'ops'))
  46 popgrp.add_argument('--multl1', dest='ops', nargs=3, metavar=('INCL[,INCL,...]', 'SCATL[,SCATL,...]', 'MULTIPLIER'), action=make_action_sharedlist('multl1', 'ops'))
  47 parser.add_argument('--frequency_multiplier', action='store', type=float, default=1., help='Multiplies the frequencies in the TMatrix file by a given factor.')
  48 # TODO enable more flexible per-sublattice specification
  49 pargs=parser.parse_args()
  50 print(pargs)
  51
  52 maxlayer=pargs.maxlayer
  53 hexside=pargs.hexside
  54 eVfreq = pargs.eVfreq
  55 freq = eVfreq*eV/hbar
  56 verbose=pargs.verbose
  57
  58 TMatrix_file = pargs.TMatrix
  59
  60 epsilon_b = pargs.background_permittivity #2.3104
  61 gaussianSigma = pargs.gaussian if pargs.gaussian else None # hexside * 222 / 7
  62 interpfreqfactor = pargs.frequency_multiplier
  63 scp_dest = pargs.scp_to if pargs.scp_to else None
  64 kdensity = pargs.kdensity
  65 chunklen = pargs.chunklen
  66
  67 ops = list()
  68 opre = re.compile('(tr|sym|copy|multl|mult)(\d*)')
  69 for oparg in pargs.ops:
  70     opm = opre.match(oparg[0])
  71     if opm:
  72         ops.append(((opm.group(2),) if opm.group(2) else (0,1), opm.group(1), oparg[1]))
  73     else:
  74         raise # should not happen
  75 print(ops)
  76
  77
  78 # -----------------finished basic CLI parsing (except for op arguments) ------------------
  79 from qpms.timetrack import _time_b, _time_e
  80 btime=_time_b(verbose)
  81
  82 import qpms
  83 import numpy as np
  84 import os, sys, warnings, math
  85 from scipy import interpolate
  86 nx = None
  87 s3 = math.sqrt(3)
  88
  89
  90 # specifikace T-matice zde
  91 cdn = c/ math.sqrt(epsilon_b)
  92 TMatrices_orig, freqs_orig, freqs_weirdunits_orig, lMaxTM = qpms.loadScuffTMatrices(TMatrix_file)
  93 lMax = lMaxTM
  94 if pargs.lMax:
  95     lMax = pargs.lMax if pargs.lMax else lMaxTM
  96 my, ny = qpms.get_mn_y(lMax)
  97 nelem = len(my)
  98 if pargs.lMax: #force commandline specified lMax
  99     TMatrices_orig = TMatrices_orig[...,0:nelem,:,0:nelem]
 100
 101 TMatrices = np.array(np.broadcast_to(TMatrices_orig[:,nx,:,:,:,:],(len(freqs_orig),2,2,nelem,2,nelem)) )
 102
 103 #TMatrices[:,:,:,:,:,ny==3] *= factor13inc
 104 #TMatrices[:,:,:,ny==3,:,:] *= factor13scat
 105 xfl = qpms.xflip_tyty(lMax)
 106 yfl = qpms.yflip_tyty(lMax)
 107 zfl = qpms.zflip_tyty(lMax)
 108 c2rot = qpms.apply_matrix_left(qpms.yflip_yy(3),qpms.xflip_yy(3),-1)
 109
 110 reCN = re.compile('(\d*)C(\d+)')
 111 #TODO C nekonečno
 112
 113 for op in ops:
 114     if op[0] == 'all':
 115         targets = (0,1)
 116     elif isinstance(op[0],int):
 117         targets = (op[0],)
 118     else:
 119         targets = op[0]
 120
 121     if op[1] == 'sym':
 122         mCN = reCN.match(op[2]) # Fuck van Rossum for not having assignments inside expressions
 123         if op[2] == 'σ_z':
 124             for t in targets:
 125                 TMatrices[:,t] = (TMatrices[:,t] + qpms.apply_ndmatrix_left(zfl,qpms.apply_ndmatrix_left(zfl, TMatrices[:,t], (-4,-3)),(-2,-1)))/2
 126         elif op[2] == 'σ_y':
 127             for t in targets:
 128                 TMatrices[:,t] = (TMatrices[:,t] + qpms.apply_ndmatrix_left(yfl,qpms.apply_ndmatrix_left(yfl, TMatrices[:,t], (-4,-3)),(-2,-1)))/2
 129         elif op[2] == 'σ_x':
 130             for t in targets:
 131                 TMatrices[:,t] = (TMatrices[:,t] + qpms.apply_ndmatrix_left(xfl,qpms.apply_ndmatrix_left(xfl, TMatrices[:,t], (-4,-3)),(-2,-1)))/2
 132         elif op[2] == 'C2': # special case of the latter
 133             for t in targets:
 134                 TMatrices[:,t] = (TMatrices[:,t] + qpms.apply_matrix_left(c2rot,qpms.apply_matrix_left(c2rot, TMatrices[:,t], -3),-1))/2
 135         elif mCN:
 136             rotN = int(mCN.group(2))
 137             TMatrix_contribs = np.empty((rotN,TMatrices.shape[0],2,nelem,2,nelem), dtype=np.complex_)
 138             for t in targets:
 139                 for i in range(rotN):
 140                     rotangle = 2*np.pi*i / rotN
 141                     rot =  qpms.WignerD_yy_fromvector(lMax,np.array([0,0,rotangle]))
 142                     rotinv = qpms.WignerD_yy_fromvector(lMax,np.array([0,0,-rotangle]))
 143                     TMatrix_contribs[i] = qpms.apply_matrix_left(rot,qpms.apply_matrix_left(rotinv, TMatrices[:,t], -3),-1)
 144                 TMatrices[:,t] = np.sum(TMatrix_contribs, axis=0) / rotN
 145         else:
 146             raise
 147     elif op[1] == 'tr':
 148         mCN = reCN.match(op[2]) # Fuck van Rossum for not having assignments inside expressions
 149         if op[2] == 'σ_z':
 150             for t in targets:
 151                 TMatrices[:,t] = qpms.apply_ndmatrix_left(zfl,qpms.apply_ndmatrix_left(zfl, TMatrices[:,t], (-4,-3)),(-2,-1))
 152         elif op[2] == 'σ_y':
 153             for t in targets:
 154                 TMatrices[:,t] = qpms.apply_ndmatrix_left(yfl,qpms.apply_ndmatrix_left(yfl, TMatrices[:,t], (-4,-3)),(-2,-1))
 155         elif op[2] == 'σ_x':
 156             for t in targets:
 157                 TMatrices[:,t] = qpms.apply_ndmatrix_left(xfl,qpms.apply_ndmatrix_left(xfl, TMatrices[:,t], (-4,-3)),(-2,-1))
 158         elif op[2] == 'C2':
 159             for t in targets:
 160                 TMatrices[:,t] = qpms.apply_matrix_left(c2rot,qpms.apply_matrix_left(c2rot, TMatrices[:,t], -3),-1)
 161         elif mCN:
 162             rotN = int(mCN.group(2))
 163             power = int(mCN.group(1)) if mCN.group(1) else 1
 164             TMatrix_contribs = np.empty((rotN,TMatrices.shape[0],2,nelem,2,nelem), dtype=np.complex_)
 165             for t in targets:
 166                 rotangle = 2*np.pi*power/rotN
 167                 rot = qpms.WignerD_yy_fromvector(lMax, np.array([0,0,rotangle]))
 168                 rotinv = qpms.WignerD_yy_fromvector(lMax, np.array([0,0,-rotangle]))
 169                 TMatrices[:,t] = qpms.apply_matrix_left(rot, qpms.apply_matrix_left(rotinv, TMatrices[:,t], -3),-1)
 170         else:
 171             raise
 172     elif op[1] == 'copy':
 173         raise # not implemented
 174     elif op[1] == 'mult':
 175         raise # not implemented
 176     elif op[1] == 'multl':
 177         incy = np.full((nelem,), False, dtype=bool)
 178         for incl in op[2][0].split(','):
 179             l = int(incl)
 180             incy += (l == ny)
 181         scaty = np.full((nelem,), False, dtype=bool)
 182         for scatl in op[2][1].split(','):
 183             l = int(scatl)
 184             scaty += (l == ny)
 185         for t in targets:
 186             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])
 187     else:
 188         raise #unknown operation; should not happen
 189
 190 TMatrices_interp = interpolate.interp1d(freqs_orig*interpfreqfactor, TMatrices, axis=0, kind='linear',fill_value="extrapolate")
 191
 192 klist_full = qpms.generate_trianglepoints(kdensity, v3d=True, include_origin=True)*3*math.pi/(3*kdensity*hexside)
 193 TMatrices_om = TMatrices_interp(freq)
 194
 195 chunkn = math.ceil(klist_full.shape[0] / chunklen)
 196
 197 if verbose:
 198     print('Evaluating %d k-points in %d chunks' % (klist_full.shape[0], chunkn), file = sys.stderr)
 199     sys.stderr.flush()
 200
 201 metadata = np.array({
 202                 'lMax' : lMax,
 203                 'maxlayer' : maxlayer,
 204                 'gaussianSigma' : gaussianSigma,
 205                 'epsilon_b' : epsilon_b,
 206                 'hexside' : hexside,
 207                 'chunkn' : chunkn,
 208                 'TMatrix_file' : TMatrix_file,
 209                 'ops' : ops,
 210                 })
 211
 212 for chunki in range(chunkn):
 213     svdout = '%s_%dnm_%.4f_c%03d.npz' % (pargs.output_prefix, hexside/1e-9, eVfreq, chunki)
 214
 215     klist = klist_full[chunki * chunklen : (chunki + 1) * chunklen]
 216
 217     svdres = qpms.hexlattice_zsym_getSVD(lMax=lMax, TMatrices_om=TMatrices_om, epsilon_b=epsilon_b, hexside=hexside, maxlayer=maxlayer,
 218             omega=freq, klist=klist, gaussianSigma=gaussianSigma, onlyNmin=False, verbose=verbose)
 219
 220     #((svUfullTElist, svSfullTElist, svVfullTElist), (svUfullTMlist, svSfullTMlist, svVfullTMlist)) = svdres
 221
 222     np.savez(svdout, omega = freq, klist = klist,
 223             metadata=metadata,
 224                      uTE = svdres[0][0],
 225                      vTE = svdres[0][2],
 226                      sTE = svdres[0][1],
 227                      uTM = svdres[1][0],
 228                      vTM = svdres[1][2],
 229                      sTM = svdres[1][1],
 230
 231     )
 232     svdres=None
 233
 234     if scp_dest:
 235         if svdout:
 236             subprocess.run(['scp', svdout, scp_dest])
 237
 238 _time_e(btime, verbose)
 239 #print(time.strftime("%H.%M:%S",time.gmtime(time.time()-begtime)))