qpms/scripts_common.py

   1 '''
   2 Mostly legacy code; new scripts use mostly argproc.py
   3 '''
   4 import warnings
   5 import argparse
   6 #import sys # for debugging purpose, TODO remove in production
   7 import os # because of path
   8 from .types import TMatrixOp, TMatrixSpec, ParticleSpec, LatticeSpec
   9 import collections
  10
  11 ''' # REMOVE IN PRODUCTION
  12 ParticleSpec = collections.namedtuple('ParticleSpec', ['label', 'position', 'tmatrix_spec'])
  13 TMatrixOp = collections.namedtuple('TMatrixOp',
  14                 ['optype', 'content'])
  15 TMatrixSpec = collections.namedtuple('TMatrixSpec',
  16         ['lMax_override', 'tmatrix_path', 'ops'])
  17 '''
  18
  19 __TODOs__ = '''
  20     - Checking validity of T-matrix ops (the arguments of --tr, --sym or similar) according to what is implemented
  21       in tmatrices.py.
  22     - Implement a more user-friendly way to define the lattice base vectors and positions of the particles.
  23       cf. https://stackoverflow.com/questions/2371436/evaluating-a-mathematical-expression-in-a-string/2371789
  24     - low priority: allow to perform some more custom operations on T-Matrix, using some kind of parsing from the previous point
  25     - Autodetect symmetries
  26
  27 '''
  28
  29 def make_action_sharedlist(opname, listname):
  30     class opAction(argparse.Action):
  31         def __call__(self, parser, args, values, option_string=None):
  32             if (not hasattr(args, listname)) or getattr(args, listname) is None:
  33                 setattr(args, listname, list())
  34             getattr(args,listname).append((opname, values))
  35     return opAction
  36
  37
  38 def add_argparse_k_output_options(parser):
  39     parser.add_argument('--kdensity', '--k_density', action='store', type=int, nargs='+', default=33, help='Number of k-points per x-axis segment FIXME DESCRIPTION')
  40     parser.add_argument('--bz_coverage', action='store', type=float, default=1., help='Brillouin zone coverage in relative length (default 1 for whole 1. BZ)')
  41     parser.add_argument('--bz_edge_width', action='store', type=float, default=0., help='Width of the more densely covered belt along the 1. BZ edge in relative lengths')
  42     parser.add_argument('--bz_edge_factor', action='store', type=float, default=8., help='Relative density of the belt along the 1. BZ edge w.r.t. k_density (default==8)')
  43     parser.add_argument('--bz_edge_twoside', action='store_true', help='Compute also the parts of the densely covered edge belt outside the 1. BZ')
  44     parser.add_argument('--bz_corner_width', action='store', type=float, default=0., help='Size of the more densely covered subcell along the 1. BZ corners in relative lengths')
  45     parser.add_argument('--bz_corner_factor', action='store', type=float, default=16., help='Relative density of the subcell along the 1. BZ corner w.r.t. k_density (default==16)')
  46     parser.add_argument('--bz_corner_twoside', action='store_true', help='Compute also the parts of the densely covered subcell outside the 1. BZ')
  47     return
  48
  49 def add_argparse_unitcell_definitions(parser):
  50     #TODO create some user-friendlier way to define lattice vectors, cf. https://stackoverflow.com/questions/2371436/evaluating-a-mathematical-expression-in-a-string/2371789
  51     parser.add_argument('--lattice_base', nargs=4, action='store', type=float, required=True, help='Lattice basis vectors x1, y1, x2, y2')
  52     parser.add_argument('--particle', '-p', nargs='+', action=make_action_sharedlist('particle', 'particlespec'), help='Particle label, coordinates x,y, and (optionally) path to the T-Matrix.')
  53     parser.add_argument('--TMatrix', '-t', nargs='+', action=make_action_sharedlist('TMatrix_path', 'particlespec'), help='Path to TMatrix file')
  54     parser.add_argument('--background_permittivity', action='store', type=float, default=1., help='Background medium relative permittivity (default 1)')
  55     parser.add_argument('--lMax', action=make_action_sharedlist('lMax', 'particlespec'), nargs='+', help='Override lMax from the TMatrix file')
  56     popgrp=parser.add_argument_group(title='Operations')
  57     popgrp.add_argument('--tr', dest='ops', nargs='+', action=make_action_sharedlist('tr', 'ops'), default=list()) # the default value for dest can be set once
  58     popgrp.add_argument('--sym', dest='ops', nargs='+', action=make_action_sharedlist('sym', 'ops'))
  59     #popgrp.add_argument('--mult', dest='ops', nargs=3, metavar=('INCSPEC', 'SCATSPEC', 'MULTIPLIER'), action=make_action_sharedlist('mult', 'ops'))
  60     #popgrp.add_argument('--multl', dest='ops', nargs=3, metavar=('INCL[,INCL,...]', 'SCATL[,SCATL,...]', 'MULTIPLIER'), action=make_action_sharedlist('multl', 'ops'))
  61     return
  62
  63 def add_argparse_infinite_lattice_options(parser):
  64     parser.add_argument('--maxlayer', action='store', type=int, default=100, help='How far to sum the lattice points to obtain the dispersion')
  65     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).')
  66     return
  67
  68 def add_argparse_output_options(parser):
  69     parser.add_argument('--output_prefix', action='store', required=True, help='Prefix to the npz output (will be appended frequency, hexside and chunkno)')
  70     parser.add_argument('--plot_TMatrix', action='store_true', help='Visualise TMatrix on the first page of the output')
  71     parser.add_argument('--scp_to', action='store', metavar='N', type=str, help='SCP the output files to a given destination')
  72     parser.add_argument('--chunklen', action='store', type=int, default=1000, help='Number of k-points per output file (default 1000)')
  73     return
  74
  75 def add_argparse_common_options(parser):
  76     parser.add_argument('--eVfreq', action='store', required=True, type=float, help='Frequency in eV')
  77     parser.add_argument('--verbose', '-v', action='count', help='Be verbose (about computation times, mostly)')
  78     parser.add_argument('--frequency_multiplier', action='store', type=float, default=1., help='Multiplies the frequencies in the TMatrix file by a given factor.')
  79
  80 def arg_preprocess_particles(pargs, d=None, return_tuple=False):
  81     '''
  82     Nanoparticle position and T-matrix path parsing
  83
  84     parser: ArgumentParser on which add_argparse_unitcell_definitions() and whose
  85     parse_args() has been called.
  86
  87     returns a list of ParticleSpec objects
  88     '''
  89     TMatrix_paths = dict()
  90     lMax_overrides = dict()
  91     default_TMatrix_path = None
  92     default_lMax_override = None
  93     if not any((arg_type == 'particle') for (arg_type, arg_content) in pargs.particlespec):
  94         # no particles positions given: suppose only one per unit cell, in the cell origin
  95         positions = {None: (0.0)}
  96     else:
  97         positions = dict()
  98     for arg_type, arg_content in pargs.particlespec:
  99         if arg_type == 'particle': # --particle option
 100             if  3 <= len(arg_content) <= 4:
 101                 try:
 102                     positions[arg_content[0]] = (float(arg_content[1]), float(arg_content[2]))
 103                 except ValueError as e:
 104                     e.args += ("second and third argument of --particle must be valid floats, given: ", arg_content)
 105                     raise
 106                 if len(arg_content) == 4:
 107                     if arg_content[0] in TMatrix_paths:
 108                         warnings.warn('T-matrix path for particle \'%s\' already specified.'
 109                         'Overriding with the last value.' % arg_content[0], SyntaxWarning)
 110                     TMatrix_paths[arg_content[0]] = arg_content[3]
 111
 112             else:
 113                 raise ValueError("--particle expects 3 or 4 arguments, %d given: " % len(arg_content), arg_content)
 114         elif arg_type == 'TMatrix_path': # --TMatrix option
 115             if len(arg_content) == 1: # --TMatrix default_path
 116                 if default_TMatrix_path is not None:
 117                     warnings.warn('Default T-matrix path already specified. Overriding with the last value.', SyntaxWarning)
 118                 default_TMatrix_path = arg_content[0]
 119             elif len(arg_content) > 1: # --TMatrix label [label2 [...]] path
 120                 for label in arg_content[:-1]:
 121                     if label in TMatrix_paths.keys():
 122                         warnings.warn('T-matrix path for particle \'%s\' already specified.'
 123                         'Overriding with the last value.' % label, SyntaxWarning)
 124                     TMatrix_paths[label] = arg_content[-1]
 125         elif arg_type == 'lMax': # --lMax option
 126             if len(arg_content) == 1: # --lMax default_lmax_override
 127                 if default_lMax_override is not None:
 128                     warnings.warn('Default lMax override value already specified. Overriding the last value.', SyntaxWarning)
 129                 default_lMax_override = int(arg_content[-1])
 130             else:
 131                 for label in arg_content[:-1]:
 132                     if label in lMax_overrides.keys:
 133                         warnings.warn('lMax override for particle \'%s\' already specified.'
 134                             'overriding with the last value.' % label, SyntaxWarning)
 135                     lMax_overrides[label] = int(arg_content[-1])
 136         else: assert False, 'unknown option type'
 137     # Check the info from positions and TMatrix_paths and lMax_overrides
 138     if not set(TMatrix_paths.keys()) <= set(positions.keys()):
 139         raise ValueError("T-Matrix path(s) for particle(s) labeled %s was given, but not their positions"
 140                 % str(set(TMatrix_paths.keys()) - set(positions.keys())))
 141     if not set(lMax_overrides.keys()) <= set(positions.keys()):
 142         raise ValueError("lMax override(s) for particle(s) labeled %s was given, but not their positions"
 143                 %str(set(lMax_overrides.keys()) - set(positions.keys())))
 144     if (set(TMatrix_paths.keys()) != set(positions.keys())) and default_TMatrix_path is None:
 145         raise ValueError("Position(s) of particles(s) labeled %s was given without their T-matrix"
 146             " and no default T-matrix was specified"
 147             % str(set(positions.keys()) - set(TMatrix_paths.keys())))
 148     # Fill default_TMatrix_path to those that don't have its own
 149     for label in (set(positions.keys()) - set(TMatrix_paths.keys())):
 150         TMatrix_paths[label] = default_TMatrix_path
 151     for path in TMatrix_paths.values():
 152         if not os.path.exists(path):
 153             raise ValueError("Cannot access T-matrix file %s. Does it exist?" % path)
 154
 155     # Assign (pre-parse) the T-matrix operations to individual particles
 156     ops = dict()
 157     for label in positions.keys(): ops[label] = list()
 158     for optype, arg_content in pargs.ops:
 159         # if, no label given, apply to all, otherwise on the specifield particles
 160         for label in (positions.keys() if len(arg_content) == 1 else arg_content[:-1]):
 161             try:
 162                 ops[label].append(TMatrixOp(optype, arg_content[-1]))
 163             except KeyError as e:
 164                 e.args += 'Specified operation on undefined particle labeled \'%s\'' % label
 165                 raise
 166
 167     #### Collect all the info about the particles / their T-matrices into one list ####
 168     # get rid of the non-unique T-matrix specs (so there is only one instance living
 169     # of each different TMatrixSpec, possibly with multiple references to it
 170     TMatrix_specs = dict((spec, spec)
 171             for spec in (TMatrixSpec(
 172                     lMax_overrides[label] if label in lMax_overrides.keys() else None,
 173                      TMatrix_paths[label],
 174                      tuple(ops[label]))
 175                 for label in positions.keys())
 176             )
 177     # particles_specs contains (label, (xpos, ypos), tmspec per element)
 178     particles_specs = [ParticleSpec(label, positions[label],
 179         TMatrix_specs[(lMax_overrides[label] if label in lMax_overrides.keys() else None,
 180                        TMatrix_paths[label],
 181                        tuple(ops[label]))]
 182         ) for label in positions.keys()]
 183
 184     return particles_specs
 185
 186 '''
 187 import argparse, re, random, string
 188 import subprocess
 189 from scipy.constants import hbar, e as eV, pi, c
 190 import warnings
 191
 192 parser = argparse.ArgumentParser()
 193 pargs=parser.parse_args()
 194 print(pargs)
 195
 196 exit(0) ###
 197
 198 maxlayer=pargs.maxlayer
 199 #DEL hexside=pargs.hexside
 200 eVfreq = pargs.eVfreq
 201 freq = eVfreq*eV/hbar
 202 verbose=pargs.verbose
 203
 204 #DEL TMatrix_file = pargs.TMatrix
 205
 206 epsilon_b = pargs.background_permittivity #2.3104
 207 gaussianSigma = pargs.gaussian if pargs.gaussian else None # hexside * 222 / 7
 208 interpfreqfactor = pargs.frequency_multiplier
 209 scp_dest = pargs.scp_to if pargs.scp_to else None
 210 kdensity = pargs.kdensity
 211 chunklen = pargs.chunklen
 212
 213 # -----------------finished basic CLI parsing (except for op arguments) ------------------
 214 from qpms.timetrack import _time_b, _time_e
 215 btime=_time_b(verbose)
 216
 217 import qpms
 218 import numpy as np
 219 import os, sys, warnings, math
 220 from scipy import interpolate
 221 nx = None
 222 s3 = math.sqrt(3)
 223
 224
 225 # specifikace T-matice zde
 226 cdn = c/ math.sqrt(epsilon_b)
 227 TMatrices_orig, freqs_orig, freqs_weirdunits_orig, lMaxTM = qpms.loadScuffTMatrices(TMatrix_file)
 228 lMax = lMaxTM
 229 if pargs.lMax:
 230     lMax = pargs.lMax if pargs.lMax else lMaxTM
 231 my, ny = qpms.get_mn_y(lMax)
 232 nelem = len(my)
 233 if pargs.lMax: #force commandline specified lMax
 234     TMatrices_orig = TMatrices_orig[...,0:nelem,:,0:nelem]
 235
 236 TMatrices = np.array(np.broadcast_to(TMatrices_orig[:,nx,:,:,:,:],(len(freqs_orig),2,2,nelem,2,nelem)) )
 237
 238 #TMatrices[:,:,:,:,:,ny==3] *= factor13inc
 239 #TMatrices[:,:,:,ny==3,:,:] *= factor13scat
 240 xfl = qpms.xflip_tyty(lMax)
 241 yfl = qpms.yflip_tyty(lMax)
 242 zfl = qpms.zflip_tyty(lMax)
 243 c2rot = qpms.apply_matrix_left(qpms.yflip_yy(3),qpms.xflip_yy(3),-1)
 244
 245 reCN = re.compile('(\d*)C(\d+)')
 246 #TODO C nekonečno
 247
 248 for op in ops:
 249     if op[0] == 'all':
 250         targets = (0,1)
 251     elif isinstance(op[0],int):
 252         targets = (op[0],)
 253     else:
 254         targets = op[0]
 255
 256     if op[1] == 'sym':
 257         mCN = reCN.match(op[2]) # Fuck van Rossum for not having assignments inside expressions
 258         if op[2] == 'σ_z':
 259             for t in targets:
 260                 TMatrices[:,t] = (TMatrices[:,t] + qpms.apply_ndmatrix_left(zfl,qpms.apply_ndmatrix_left(zfl, TMatrices[:,t], (-4,-3)),(-2,-1)))/2
 261         elif op[2] == 'σ_y':
 262             for t in targets:
 263                 TMatrices[:,t] = (TMatrices[:,t] + qpms.apply_ndmatrix_left(yfl,qpms.apply_ndmatrix_left(yfl, TMatrices[:,t], (-4,-3)),(-2,-1)))/2
 264         elif op[2] == 'σ_x':
 265             for t in targets:
 266                 TMatrices[:,t] = (TMatrices[:,t] + qpms.apply_ndmatrix_left(xfl,qpms.apply_ndmatrix_left(xfl, TMatrices[:,t], (-4,-3)),(-2,-1)))/2
 267         elif op[2] == 'C2': # special case of the latter
 268             for t in targets:
 269                 TMatrices[:,t] = (TMatrices[:,t] + qpms.apply_matrix_left(c2rot,qpms.apply_matrix_left(c2rot, TMatrices[:,t], -3),-1))/2
 270         elif mCN:
 271             rotN = int(mCN.group(2))
 272             TMatrix_contribs = np.empty((rotN,TMatrices.shape[0],2,nelem,2,nelem), dtype=np.complex_)
 273             for t in targets:
 274                 for i in range(rotN):
 275                     rotangle = 2*np.pi*i / rotN
 276                     rot =  qpms.WignerD_yy_fromvector(lMax,np.array([0,0,rotangle]))
 277                     rotinv = qpms.WignerD_yy_fromvector(lMax,np.array([0,0,-rotangle]))
 278                     TMatrix_contribs[i] = qpms.apply_matrix_left(rot,qpms.apply_matrix_left(rotinv, TMatrices[:,t], -3),-1)
 279                 TMatrices[:,t] = np.sum(TMatrix_contribs, axis=0) / rotN
 280         else:
 281             raise ValueError('\'%d\' is not an implemented symmetry operation' % op[2])
 282     elif op[1] == 'tr':
 283         mCN = reCN.match(op[2]) # Fuck van Rossum for not having assignments inside expressions
 284         if op[2] == 'σ_z':
 285             for t in targets:
 286                 TMatrices[:,t] = qpms.apply_ndmatrix_left(zfl,qpms.apply_ndmatrix_left(zfl, TMatrices[:,t], (-4,-3)),(-2,-1))
 287         elif op[2] == 'σ_y':
 288             for t in targets:
 289                 TMatrices[:,t] = qpms.apply_ndmatrix_left(yfl,qpms.apply_ndmatrix_left(yfl, TMatrices[:,t], (-4,-3)),(-2,-1))
 290         elif op[2] == 'σ_x':
 291             for t in targets:
 292                 TMatrices[:,t] = qpms.apply_ndmatrix_left(xfl,qpms.apply_ndmatrix_left(xfl, TMatrices[:,t], (-4,-3)),(-2,-1))
 293         elif op[2] == 'C2':
 294             for t in targets:
 295                 TMatrices[:,t] = qpms.apply_matrix_left(c2rot,qpms.apply_matrix_left(c2rot, TMatrices[:,t], -3),-1)
 296         elif mCN:
 297             rotN = int(mCN.group(2))
 298             power = int(mCN.group(1)) if mCN.group(1) else 1
 299             TMatrix_contribs = np.empty((rotN,TMatrices.shape[0],2,nelem,2,nelem), dtype=np.complex_)
 300             for t in targets:
 301                 rotangle = 2*np.pi*power/rotN
 302                 rot = qpms.WignerD_yy_fromvector(lMax, np.array([0,0,rotangle]))
 303                 rotinv = qpms.WignerD_yy_fromvector(lMax, np.array([0,0,-rotangle]))
 304                 TMatrices[:,t] = qpms.apply_matrix_left(rot, qpms.apply_matrix_left(rotinv, TMatrices[:,t], -3),-1)
 305         else:
 306             raise ValueError('\'%d\' is not an implemented T-matrix transformation operation' % op[2])
 307     elif op[1] == 'copy':
 308         raise # not implemented
 309     elif op[1] == 'mult':
 310         raise # not implemented
 311     elif op[1] == 'multl':
 312         incy = np.full((nelem,), False, dtype=bool)
 313         for incl in op[2][0].split(','):
 314             l = int(incl)
 315             incy += (l == ny)
 316         scaty = np.full((nelem,), False, dtype=bool)
 317         for scatl in op[2][1].split(','):
 318             l = int(scatl)
 319             scaty += (l == ny)
 320         for t in targets:
 321             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])
 322     else:
 323         raise #unknown operation; should not happen
 324
 325 TMatrices_interp = interpolate.interp1d(freqs_orig*interpfreqfactor, TMatrices, axis=0, kind='linear',fill_value="extrapolate")
 326
 327 klist_full = qpms.generate_trianglepoints(kdensity, v3d=True, include_origin=True)*3*math.pi/(3*kdensity*hexside)
 328 TMatrices_om = TMatrices_interp(freq)
 329
 330 chunkn = math.ceil(klist_full.shape[0] / chunklen)
 331
 332 if verbose:
 333     print('Evaluating %d k-points in %d chunks' % (klist_full.shape[0], chunkn), file = sys.stderr)
 334     sys.stderr.flush()
 335
 336 metadata = np.array({
 337                 'lMax' : lMax,
 338                 'maxlayer' : maxlayer,
 339                 'gaussianSigma' : gaussianSigma,
 340                 'epsilon_b' : epsilon_b,
 341                 'hexside' : hexside,
 342                 'chunkn' : chunkn,
 343                 'TMatrix_file' : TMatrix_file,
 344                 'ops' : ops,
 345                 })
 346
 347 for chunki in range(chunkn):
 348     svdout = '%s_%dnm_%.4f_c%03d.npz' % (pargs.output_prefix, hexside/1e-9, eVfreq, chunki)
 349
 350     klist = klist_full[chunki * chunklen : (chunki + 1) * chunklen]
 351
 352     svdres = qpms.hexlattice_zsym_getSVD(lMax=lMax, TMatrices_om=TMatrices_om, epsilon_b=epsilon_b, hexside=hexside, maxlayer=maxlayer,
 353             omega=freq, klist=klist, gaussianSigma=gaussianSigma, onlyNmin=False, verbose=verbose)
 354
 355     #((svUfullTElist, svSfullTElist, svVfullTElist), (svUfullTMlist, svSfullTMlist, svVfullTMlist)) = svdres
 356
 357     np.savez(svdout, omega = freq, klist = klist,
 358             metadata=metadata,
 359                      uTE = svdres[0][0],
 360                      vTE = svdres[0][2],
 361                      sTE = svdres[0][1],
 362                      uTM = svdres[1][0],
 363                      vTM = svdres[1][2],
 364                      sTM = svdres[1][1],
 365
 366     )
 367     svdres=None
 368
 369     if scp_dest:
 370         if svdout:
 371             subprocess.run(['scp', svdout, scp_dest])
 372
 373 _time_e(btime, verbose)
 374 #print(time.strftime("%H.%M:%S",time.gmtime(time.time()-begtime)))
 375 '''