misc/finitesqlatzsym-scatter.py

   1 #!/usr/bin/env python3
   2
   3 import argparse, re, random, string, sys
   4 import subprocess
   5 import warnings
   6 from scipy.constants import hbar, e as eV, pi, c
   7
   8 unitcell_size = 1 # rectangular lattice
   9 unitcell_indices = tuple(range(unitcell_size))
  10
  11 def make_action_sharedlist(opname, listname):
  12     class opAction(argparse.Action):
  13         def __call__(self, parser, args, values, option_string=None):
  14             if (not hasattr(args, listname)) or getattr(args, listname) is None:
  15                 setattr(args, listname, list())
  16             getattr(args,listname).append((opname, values))
  17     return opAction
  18
  19 parser = argparse.ArgumentParser()
  20 #TODO? použít type=argparse.FileType('r') ?
  21 parser.add_argument('--TMatrix', action='store', required=True, help='Path to TMatrix file')
  22 #parser.add_argument('--griddir', action='store', required=True, help='Path to the directory with precalculated translation operators')
  23 parser.add_argument('--output_prefix', '-p', '-o', action='store', required=True, help='Prefix to the npz output (will be appended frequency, hexside and chunkno)')
  24 parser.add_argument('--nosuffix', action='store_true', help='Do not add dimension metadata to the output filenames')
  25 #sizepar = parser.add_mutually_exclusive_group(required=True)
  26 #parser.add_argument('--hexside', action='store', type=float, required=True, help='Lattice hexagon size length')
  27 parser.add_argument('--dx', action='store', type=float, required=True, help='x-direction lattice constant')
  28 parser.add_argument('--dy', action='store', type=float, required=True, help='y-direction lattice constant')
  29
  30 parser.add_argument('--Nx', '--nx', action='store', type=int, required=True, help='Lattice points in the x-direction')
  31 parser.add_argument('--Ny', '--ny', action='store', type=int, required=True, help='Lattice points in the y-direction')
  32
  33 # In these default settings, the area is 2x2 times larger than first BZ
  34 parser.add_argument('--kxmin', action='store', type=float, default=-1., help='TODO')
  35 parser.add_argument('--kxmax', action='store', type=float, default=1., help='TODO')
  36 parser.add_argument('--kymin', action='store', type=float, default=-1., help='TODO')
  37 parser.add_argument('--kymax', action='store', type=float, default=1., help='TODO')
  38
  39 #parser.add_argument('--kdensity', action='store', type=int, default=33, help='Number of k-points per x-axis segment')
  40 parser.add_argument('--kxdensity', action='store', type=int, default=51, help='k-space resolution in the x-direction')
  41 parser.add_argument('--kydensity', action='store', type=int, default=51, help='k-space resolution in the y-direction')
  42
  43 partgrp = parser.add_mutually_exclusive_group()
  44 partgrp.add_argument('--only_TE', action='store_true', help='Calculate only the projection on the E⟂z modes')
  45 partgrp.add_argument('--only_TM', action='store_true', help='Calculate only the projection on the E∥z modes')
  46 partgrp.add_argument('--serial', action='store_true', help='Calculate the TE and TM parts separately to save memory')
  47
  48 parser.add_argument('--nocentre', action='store_true', help='Place the coordinate origin to the left bottom corner rather that to the centre of the array')
  49
  50 parser.add_argument('--plot_TMatrix', action='store_true', help='Visualise TMatrix on the first page of the output')
  51 #parser.add_argument('--SVD_output', action='store', help='Path to output singular value decomposition result')
  52 parser.add_argument('--maxlayer', action='store', type=int, default=100, help='How far to sum the lattice points to obtain the dispersion')
  53 parser.add_argument('--scp_to', action='store', metavar='N', type=str, help='SCP the output files to a given destination')
  54 parser.add_argument('--background_permittivity', action='store', type=float, default=1., help='Background medium relative permittivity (default 1)')
  55 parser.add_argument('--eVfreq', action='store', required=True, type=float, help='Frequency in eV')
  56
  57 parser.add_argument('--chunklen', action='store', type=int, default=3000, help='Number of k-points per output file (default 3000)')
  58 parser.add_argument('--lMax', action='store', type=int, help='Override lMax from the TMatrix file')
  59 #TODO some more sophisticated x axis definitions
  60 #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).')
  61 parser.add_argument('--verbose', '-v', action='count', help='Be verbose (about computation times, mostly)')
  62 popgrp=parser.add_argument_group(title='Operations')
  63 popgrp.add_argument('--tr', dest='ops', action=make_action_sharedlist('tr', 'ops'), default=list()) # the default value for dest can be set once
  64 for i in unitcell_indices:
  65     popgrp.add_argument('--tr%d'%i, dest='ops', action=make_action_sharedlist('tr%d'%i, 'ops'))
  66 popgrp.add_argument('--sym', dest='ops', action=make_action_sharedlist('sym', 'ops'))
  67 for i in unitcell_indices:
  68     popgrp.add_argument('--sym%d'%i, dest='ops', action=make_action_sharedlist('sym%d'%i, 'ops'))
  69 #popgrp.add_argument('--mult', dest='ops', nargs=3, metavar=('INCSPEC', 'SCATSPEC', 'MULTIPLIER'), action=make_action_sharedlist('mult', 'ops'))
  70 #popgrp.add_argument('--mult0', dest='ops', nargs=3, metavar=('INCSPEC', 'SCATSPEC', 'MULTIPLIER'), action=make_action_sharedlist('mult0', 'ops'))
  71 #popgrp.add_argument('--mult1', dest='ops', nargs=3, metavar=('INCSPEC', 'SCATSPEC', 'MULTIPLIER'), action=make_action_sharedlist('mult1', 'ops'))
  72 popgrp.add_argument('--multl', dest='ops', nargs=3, metavar=('INCL[,INCL,...]', 'SCATL[,SCATL,...]', 'MULTIPLIER'), action=make_action_sharedlist('multl', 'ops'))
  73 for i in unitcell_indices:
  74     popgrp.add_argument('--multl%d'%i, dest='ops', nargs=3, metavar=('INCL[,INCL,...]', 'SCATL[,SCATL,...]', 'MULTIPLIER'), action=make_action_sharedlist('multl%d'%i, 'ops'))
  75 #popgrp.add_argument('--multl1', dest='ops', nargs=3, metavar=('INCL[,INCL,...]', 'SCATL[,SCATL,...]', 'MULTIPLIER'), action=make_action_sharedlist('multl1', 'ops'))
  76 parser.add_argument('--frequency_multiplier', action='store', type=float, default=1., help='Multiplies the frequencies in the TMatrix file by a given factor.')
  77 # TODO enable more flexible per-sublattice specification
  78 pargs=parser.parse_args()
  79 if pargs.verbose:
  80     print(pargs, file = sys.stderr)
  81
  82 maxlayer=pargs.maxlayer
  83 eVfreq = pargs.eVfreq
  84 freq = eVfreq*eV/hbar
  85 verbose=pargs.verbose
  86 dy = pargs.dy
  87 dx = pargs.dx
  88 Ny = pargs.Ny
  89 Nx = pargs.Nx
  90
  91 TMatrix_file = pargs.TMatrix
  92
  93 epsilon_b = pargs.background_permittivity #2.3104
  94 #gaussianSigma = pargs.gaussian if pargs.gaussian else None # hexside * 222 / 7
  95 interpfreqfactor = pargs.frequency_multiplier
  96 scp_dest = pargs.scp_to if pargs.scp_to else None
  97 kxdensity = pargs.kxdensity
  98 kydensity = pargs.kydensity
  99 chunklen = pargs.chunklen
 100
 101 ops = list()
 102 opre = re.compile('(tr|sym|copy|multl|mult)(\d*)')
 103 for oparg in pargs.ops:
 104     opm = opre.match(oparg[0])
 105     if opm:
 106         ops.append(((opm.group(2),) if opm.group(2) else unitcell_indices, opm.group(1), oparg[1]))
 107     else:
 108         raise # should not happen
 109 if(verbose):
 110     print(ops, file = sys.stderr)
 111
 112
 113 # -----------------finished basic CLI parsing (except for op arguments) ------------------
 114 from qpms.timetrack import _time_b, _time_e
 115 btime=_time_b(verbose)
 116
 117 import qpms
 118 import numpy as np
 119 import os, warnings, math
 120 from scipy import interpolate
 121 nx = None
 122 s3 = math.sqrt(3)
 123
 124
 125 # specifikace T-matice zde
 126 refind = math.sqrt(epsilon_b)
 127 cdn = c / refind
 128 k_0 = freq * refind / c # = freq / cdn
 129 TMatrices_orig, freqs_orig, freqs_weirdunits_orig, lMaxTM = qpms.loadScuffTMatrices(TMatrix_file)
 130 lMax = lMaxTM
 131 if pargs.lMax:
 132     lMax = pargs.lMax if pargs.lMax else lMaxTM
 133 my, ny = qpms.get_mn_y(lMax)
 134 nelem = len(my)
 135 if pargs.lMax: #force commandline specified lMax
 136     TMatrices_orig = TMatrices_orig[...,0:nelem,:,0:nelem]
 137 TMatrices = np.array(np.broadcast_to(TMatrices_orig[:,nx,:,:,:,:],(len(freqs_orig),unitcell_size,2,nelem,2,nelem)) )
 138 xfl = qpms.xflip_tyty(lMax)
 139 yfl = qpms.yflip_tyty(lMax)
 140 zfl = qpms.zflip_tyty(lMax)
 141 c2rot = qpms.apply_matrix_left(qpms.yflip_yy(3),qpms.xflip_yy(3),-1)
 142
 143 reCN = re.compile('(\d*)C(\d+)')
 144 #TODO C nekonečno
 145
 146 for op in ops:
 147     if op[0] == 'all':
 148         #targets = (0,1)
 149         targets = unitcell_indices
 150     elif isinstance(op[0],int):
 151         targets = (op[0],)
 152     else:
 153         targets = op[0]
 154
 155     if op[1] == 'sym':
 156         mCN = reCN.match(op[2]) # Fuck van Rossum for not having assignments inside expressions
 157         if op[2] == 'σ_z':
 158             for t in targets:
 159                 TMatrices[:,t] = (TMatrices[:,t] + qpms.apply_ndmatrix_left(zfl,qpms.apply_ndmatrix_left(zfl, TMatrices[:,t], (-4,-3)),(-2,-1)))/2
 160         elif op[2] == 'σ_y':
 161             for t in targets:
 162                 TMatrices[:,t] = (TMatrices[:,t] + qpms.apply_ndmatrix_left(yfl,qpms.apply_ndmatrix_left(yfl, TMatrices[:,t], (-4,-3)),(-2,-1)))/2
 163         elif op[2] == 'σ_x':
 164             for t in targets:
 165                 TMatrices[:,t] = (TMatrices[:,t] + qpms.apply_ndmatrix_left(xfl,qpms.apply_ndmatrix_left(xfl, TMatrices[:,t], (-4,-3)),(-2,-1)))/2
 166         elif op[2] == 'C2': # special case of the latter
 167             for t in targets:
 168                 TMatrices[:,t] = (TMatrices[:,t] + qpms.apply_matrix_left(c2rot,qpms.apply_matrix_left(c2rot, TMatrices[:,t], -3),-1))/2
 169         elif mCN:
 170             rotN = int(mCN.group(2))
 171             TMatrix_contribs = np.empty((rotN,TMatrices.shape[0],2,nelem,2,nelem), dtype=np.complex_)
 172             for t in targets:
 173                 for i in range(rotN):
 174                     rotangle = 2*np.pi*i / rotN
 175                     rot =  qpms.WignerD_yy_fromvector(lMax,np.array([0,0,rotangle]))
 176                     rotinv = qpms.WignerD_yy_fromvector(lMax,np.array([0,0,-rotangle]))
 177                     TMatrix_contribs[i] = qpms.apply_matrix_left(rot,qpms.apply_matrix_left(rotinv, TMatrices[:,t], -3),-1)
 178                 TMatrices[:,t] = np.sum(TMatrix_contribs, axis=0) / rotN
 179         else:
 180             raise
 181     elif op[1] == 'tr':
 182         mCN = reCN.match(op[2]) # Fuck van Rossum for not having assignments inside expressions
 183         if op[2] == 'σ_z':
 184             for t in targets:
 185                 TMatrices[:,t] = qpms.apply_ndmatrix_left(zfl,qpms.apply_ndmatrix_left(zfl, TMatrices[:,t], (-4,-3)),(-2,-1))
 186         elif op[2] == 'σ_y':
 187             for t in targets:
 188                 TMatrices[:,t] = qpms.apply_ndmatrix_left(yfl,qpms.apply_ndmatrix_left(yfl, TMatrices[:,t], (-4,-3)),(-2,-1))
 189         elif op[2] == 'σ_x':
 190             for t in targets:
 191                 TMatrices[:,t] = qpms.apply_ndmatrix_left(xfl,qpms.apply_ndmatrix_left(xfl, TMatrices[:,t], (-4,-3)),(-2,-1))
 192         elif op[2] == 'C2':
 193             for t in targets:
 194                 TMatrices[:,t] = qpms.apply_matrix_left(c2rot,qpms.apply_matrix_left(c2rot, TMatrices[:,t], -3),-1)
 195         elif mCN:
 196             rotN = int(mCN.group(2))
 197             power = int(mCN.group(1)) if mCN.group(1) else 1
 198             TMatrix_contribs = np.empty((rotN,TMatrices.shape[0],2,nelem,2,nelem), dtype=np.complex_)
 199             for t in targets:
 200                 rotangle = 2*np.pi*power/rotN
 201                 rot = qpms.WignerD_yy_fromvector(lMax, np.array([0,0,rotangle]))
 202                 rotinv = qpms.WignerD_yy_fromvector(lMax, np.array([0,0,-rotangle]))
 203                 TMatrices[:,t] = qpms.apply_matrix_left(rot, qpms.apply_matrix_left(rotinv, TMatrices[:,t], -3),-1)
 204         else:
 205             raise
 206     elif op[1] == 'copy':
 207         raise # not implemented
 208     elif op[1] == 'mult':
 209         raise # not implemented
 210     elif op[1] == 'multl':
 211         incy = np.full((nelem,), False, dtype=bool)
 212         for incl in op[2][0].split(','):
 213             l = int(incl)
 214             incy += (l == ny)
 215         scaty = np.full((nelem,), False, dtype=bool)
 216         for scatl in op[2][1].split(','):
 217             l = int(scatl)
 218             scaty += (l == ny)
 219         for t in targets:
 220             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])
 221     else:
 222         raise #unknown operation; should not happen
 223
 224 TMatrices_interp = interpolate.interp1d(freqs_orig*interpfreqfactor, TMatrices, axis=0, kind='linear',fill_value="extrapolate")
 225
 226
 227 xpositions = np.arange(Nx) * dx
 228 ypositions = np.arange(Ny) * dy
 229 if not pargs.nocentre:
 230     xpositions -= Nx * dx / 2
 231     ypositions -= Ny * dy / 2
 232 xpositions, ypositions = np.meshgrid(xpositions, ypositions, indexing='ij', copy=False)
 233 positions=np.stack((xpositions.ravel(),ypositions.ravel()), axis=-1)
 234 positions=positions[np.random.permutation(len(positions))]
 235 N = positions.shape[0]
 236
 237 kx = np.linspace(pargs.kxmin, pargs.kxmax, num=pargs.kxdensity, endpoint=True) * 2*np.pi / dx
 238 ky = np.linspace(pargs.kymin, pargs.kymax, num=pargs.kydensity, endpoint=True) * 2*np.pi / dy
 239 kx, ky = np.meshgrid(kx, ky, indexing='ij', copy=False)
 240 kz = np.sqrt(k_0**2 - (kx ** 2 + ky ** 2))
 241
 242 klist_full = np.stack((kx,ky,kz), axis=-1).reshape((-1,3))
 243 TMatrices_om = TMatrices_interp(freq)
 244
 245 chunkn = math.ceil(klist_full.size / 3 / chunklen)
 246
 247 if verbose:
 248     print('Evaluating %d k-points' % klist_full.size + ('in %d chunks'%chunkn) if chunkn>1 else '' , file = sys.stderr)
 249     sys.stderr.flush()
 250
 251 try:
 252     version = qpms.__version__
 253 except NameError:
 254     version = None
 255
 256 metadata = np.array({
 257                 'script': os.path.basename(__file__),
 258                 'version': version,
 259                 'type' : 'Plane wave scattering on a finite rectangular lattice',
 260                 'lMax' : lMax,
 261                 'dx' : dx,
 262                 'dy' : dy,
 263                 'Nx' : Nx,
 264                 'Ny' : Ny,
 265                 #'maxlayer' : maxlayer,
 266                 #'gaussianSigma' : gaussianSigma,
 267                 'epsilon_b' : epsilon_b,
 268                 #'hexside' : hexside,
 269                 'chunkn' : chunkn,
 270                 'chunki' : 0,
 271                 'TMatrix_file' : TMatrix_file,
 272                 'ops' : ops,
 273                 'centred' : not pargs.nocentre
 274                 })
 275
 276
 277 scat = qpms.Scattering_2D_zsym(positions, TMatrices_om, k_0, verbose=verbose)
 278
 279 if pargs.only_TE:
 280     actions = (0,)
 281 elif pargs.only_TM:
 282     actions = (1,)
 283 elif pargs.serial:
 284     actions = (0,1)
 285 else:
 286     actions = (None,)
 287
 288 xu = np.array((1,0,0))
 289 yu = np.array((0,1,0))
 290 zu = np.array((0,0,1))
 291 TEč, TMč = qpms.symz_indexarrays(lMax)
 292
 293 klist_full_2D = klist_full[...,:2]
 294 klist_full_dir = klist_full/np.linalg.norm(klist_full, axis=-1, keepdims=True)
 295 for action in actions:
 296     if action is None:
 297         scat.prepare(verbose=verbose)
 298         actionstring = ''
 299     else:
 300         scat.prepare_partial(action, verbose=verbose)
 301         actionstring = '.TM' if action else '.TE'
 302     for chunki in range(chunkn):
 303         sbtime = _time_b(verbose, step='Solving the scattering problem, chunk %d'%chunki+actionstring)
 304         if pargs.nosuffix:
 305             outfile = pargs.output_prefix + actionstring + (
 306                     ('.%03d' % chunki) if chunkn > 1 else '')
 307         else:
 308             outfile = '%s_%dx%d_%.0fnmx%.0fnm_%.4f%s%s.npz' % (
 309                     pargs.output_prefix, Nx, Ny, dx/1e-9, dy/1e-9,
 310                     eVfreq, actionstring,
 311                     (".%03d" % chunki) if chunkn > 1 else '')
 312
 313         klist = klist_full[chunki * chunklen : (chunki + 1) * chunklen]
 314         klist2d = klist_full_2D[chunki * chunklen : (chunki + 1) * chunklen]
 315         klistdir = klist_full_dir[chunki * chunklen : (chunki + 1) * chunklen]
 316
 317         '''
 318         The following loop is a fuckup that has its roots in the fact that
 319         the function qpms.get_π̃τ̃_y1 in qpms_p.py is not vectorized
 320         (and consequently, neither is plane_pq_y.)
 321
 322         And Scattering_2D_zsym.scatter_partial is not vectorized, either.
 323         '''
 324         if action == 0 or action is None:
 325             xresult = np.full((klist.shape[0], N, nelem), np.nan, dtype=complex)
 326             yresult = np.full((klist.shape[0], N, nelem), np.nan, dtype=complex)
 327         if action == 1 or action is None:
 328             zresult = np.full((klist.shape[0], N, nelem), np.nan, dtype=complex)
 329         for i in range(klist.shape[0]):
 330             if math.isnan(klist[i,2]):
 331                 if(verbose):
 332                     print("%d. momentum %s invalid (k_0=%f), skipping" % (i, str(klist[i]),k_0))
 333                 continue
 334             kdir = klistdir[i]
 335             phases = np.exp(-1j*np.sum(klist2d[i] * positions, axis=-1))
 336             if action == 0 or action is None:
 337                 pq = np.array(qpms.plane_pq_y(lMax, kdir, xu)).ravel()[TEč] * phases[:, nx]
 338                 xresult[i] = scat.scatter_partial(0, pq)
 339                 pq = np.array(qpms.plane_pq_y(lMax, kdir, yu)).ravel()[TEč] * phases[:, nx]
 340                 yresult[i] = scat.scatter_partial(0, pq)
 341             if action == 1 or action is None:
 342                 pq = np.array(qpms.plane_pq_y(lMax, kdir, zu)).ravel()[TMč] * phases[:, nx]
 343                 zresult[i] = scat.scatter_partial(1, pq)
 344         _time_e(sbtime, verbose, step='Solving the scattering problem, chunk %d'%chunki+actionstring)
 345
 346         metadata[()]['chunki'] = chunki
 347         if action is None:
 348             np.savez(outfile, omega = freq, klist = klist,
 349                 metadata=metadata,
 350                 positions=positions,
 351                 ab_x=xresult,
 352                 ab_y=yresult,
 353                 ab_z=zresult
 354                 )
 355         elif action == 0:
 356             np.savez(outfile, omega = freq, klist = klist,
 357                 metadata=metadata,
 358                 positions=positions,
 359                 ab_x=xresult,
 360                 ab_y=yresult,
 361                 )
 362         elif action == 1:
 363             np.savez(outfile, omega = freq, klist = klist,
 364                 metadata=metadata,
 365                 positions=positions,
 366                 ab_z=zresult
 367                 )
 368         else:
 369             raise
 370
 371         if scp_dest:
 372             if outfile:
 373                 subprocess.run(['scp', outfile, scp_dest])
 374     scat.forget_matrices() # free memory in case --serial was used
 375
 376 _time_e(btime, verbose)