misc/finiterectlat-constant-driving.py

   1 #!/usr/bin/env python3
   2
   3 import math
   4 from qpms.argproc import ArgParser, make_dict_action, sslice, annotate_pdf_metadata
   5 figscale=3
   6
   7 ap = ArgParser(['rectlattice2d_finite', 'single_particle', 'single_lMax', 'single_omega'])
   8 ap.add_argument("-k", '--wavevector', nargs=2, type=float, required=True, help='"Bloch" vector, modulating phase of the driving', metavar=('KX', 'KY'), default=(0., 0.))
   9 # ap.add_argument("--kpi", action='store_true', help="Indicates that the k vector is given in natural units instead of SI, i.e. the arguments given by -k shall be automatically multiplied by pi / period (given by -p argument)")
  10 ap.add_argument("-o", "--output", type=str, required=False, help='output path (if not provided, will be generated automatically)')
  11 ap.add_argument("-O", "--plot-out", type=str, required=False, help="path to plot output (optional)")
  12 ap.add_argument("-P", "--plot", action='store_true', help="if -p not given, plot to a default path")
  13 ap.add_argument("-g", "--save-gradually", action='store_true', help="saves the partial result after computing each irrep")
  14 ap.add_argument("-S", "--symmetry-adapted", default=None, help="Use a symmetry-adapted basis of a given point group instead of individual spherical harmonics")
  15 ap.add_argument("-d", "--ccd-distance", type=float, default=math.nan, help='Far-field "CCD" distance from the sample')
  16 ap.add_argument("-D", "--ccd-size", type=float, default=math.nan, help='Far-field "CCD" width and heighth')
  17 ap.add_argument("-R", "--ccd-resolution", type=int, default=101, help='Far-field "CCD" resolution')
  18 ap.add_argument("--xslice", default={None:None}, nargs=2,
  19         action=make_dict_action(argtype=sslice, postaction='append', first_is_key=True),
  20         )
  21 ap.add_argument("--yslice", default={None:None}, nargs=2,
  22         action=make_dict_action(argtype=sslice, postaction='append', first_is_key=True),
  23         )
  24
  25
  26 #ap.add_argument("--irrep", type=str, default="none", help="Irrep subspace (irrep index from 0 to 7, irrep label, or 'none' for no irrep decomposition")
  27
  28
  29 a=ap.parse_args()
  30
  31 import logging
  32 logging.basicConfig(format='%(asctime)s %(message)s', level=logging.INFO)
  33
  34 Nx, Ny = a.size
  35 px, py = a.period
  36
  37 particlestr = ("sph" if a.height is None else "cyl") + ("_r%gnm" % (a.radius*1e9))
  38 if a.height is not None: particlestr += "_h%gnm" % (a.height * 1e9)
  39 defaultprefix = "cd_%s_p%gnmx%gnm_%dx%d_m%s_n%s_k_%g_%g_f%geV_L%d_micro-%s" % (
  40     particlestr, px*1e9, py*1e9, Nx, Ny, str(a.material), str(a.background), a.wavevector[0], a.wavevector[1], a.eV, a.lMax, "SO3" if a.symmetry_adapted is None else a.symmetry_adapted)
  41 logging.info("Default file prefix: %s" % defaultprefix)
  42
  43 import numpy as np
  44 import qpms
  45 from qpms.cybspec import BaseSpec
  46 from qpms.cytmatrices import CTMatrix, TMatrixGenerator
  47 from qpms.qpms_c import Particle, qpms_library_version
  48 from qpms.cymaterials import EpsMu, EpsMuGenerator, LorentzDrudeModel, lorentz_drude
  49 from qpms.cycommon import DebugFlags, dbgmsg_enable
  50 from qpms import FinitePointGroup, ScatteringSystem, BesselType, eV, hbar
  51 from qpms.symmetries import point_group_info
  52 eh = eV/hbar
  53
  54 # Check slice ranges and generate all corresponding combinations
  55 slicepairs = []
  56 slicelabels = set(a.xslice.keys())  | set(a.yslice.keys())
  57 for label in slicelabels:
  58     rowslices = a.xslice.get(label, None)
  59     colslices = a.yslice.get(label, None)
  60     # TODO check validity of the slices.
  61     if rowslices is None:
  62         rowslices = [slice(None, None, None)]
  63     if colslices is None:
  64         colslices = [slice(None, None, None)]
  65     for rs in rowslices:
  66         for cs in colslices:
  67             slicepairs.append((rs, cs))
  68
  69 def realdipfieldlabels(yp):
  70     if yp == 0: return 'x'
  71     if yp == 1: return 'y'
  72     if yp == 2: return 'z'
  73     raise ValueError
  74 def realdipfields(vecgrid, yp):
  75     if yp == 1:
  76         return vecgrid[...,0] + vecgrid[...,2]
  77     if yp == 0:
  78         return -1j*(vecgrid[...,0] - vecgrid[...,2])
  79     if yp == 2:
  80         return vecgrid[...,1]
  81     raise ValueError
  82
  83 def float_nicestr(x, tol=1e-5):
  84     x = float(x)
  85     if .5**2 - abs(x) < tol:
  86         return(("-" if x < 0 else '+') + "2^{-2}")
  87     else:
  88         return "%+.3g" % x
  89
  90 def cplx_nicestr(x, tol=1e-5):
  91     x = complex(x)
  92     if x == 0:
  93         return '0'
  94     ret = ""
  95     if x.real:
  96         ret = ret + float_nicestr(x.real, tol)
  97     if x.imag:
  98         ret = ret + float_nicestr(x.imag, tol) + 'i'
  99     if x.real and x.imag:
 100         return '(' + ret + ')'
 101     else:
 102         return ret
 103
 104 def cleanarray(a, atol=1e-10, copy=True):
 105     a = np.array(a, copy=copy)
 106     sieve = abs(a.real) < atol
 107     a[sieve] = 1j * a[sieve].imag
 108     sieve = abs(a.imag) < atol
 109     a[sieve] = a[sieve].real
 110     return a
 111
 112 def nicerot(a, atol=1e-10, copy=True): #gives array a "nice" phase
 113     a = np.array(a, copy=copy)
 114     i = np.argmax(abs(a))
 115     a = a / a[i] * abs(a[i])
 116     return a
 117
 118 dbgmsg_enable(DebugFlags.INTEGRATION)
 119
 120 #Particle positions
 121 orig_x = (np.arange(Nx/2) + (0 if (Nx % 2) else .5)) * px
 122 orig_y = (np.arange(Ny/2) + (0 if (Ny % 2) else .5)) * py
 123
 124 orig_xy = np.stack(np.meshgrid(orig_x, orig_y), axis = -1)
 125
 126
 127 omega = ap.omega
 128
 129 bspec = BaseSpec(lMax = a.lMax)
 130 medium = EpsMuGenerator(ap.background_epsmu)
 131 particles= [Particle(orig_xy[i], ap.tmgen, bspec) for i in np.ndindex(orig_xy.shape[:-1])]
 132
 133 sym = FinitePointGroup(point_group_info['D2h'])
 134 ss, ssw = ScatteringSystem.create(particles=particles, medium=medium, omega=omega, sym=sym)
 135
 136 wavenumber = ap.background_epsmu.k(omega) # Currently, ScatteringSystem does not "remember" frequency nor wavenumber
 137
 138 # Mapping between ss particles and grid positions
 139 positions = ss.positions
 140 xpositions = np.unique(positions[:,0])
 141 assert(len(xpositions) == Nx)
 142 ypositions = np.unique(positions[:,1])
 143 assert(len(ypositions == Ny))
 144 # particle positions as integer indices
 145 posmap = np.empty((positions.shape[0],2), dtype=int)
 146 invposmap = np.empty((Nx, Ny), dtype=int)
 147 for i, pos in enumerate(positions):
 148     posmap[i,0] = np.searchsorted(xpositions, positions[i,0])
 149     posmap[i,1] = np.searchsorted(ypositions, positions[i,1])
 150     invposmap[posmap[i,0], posmap[i, 1]] = i
 151
 152 def fullvec2grid(fullvec, swapxy=False):
 153     arr = np.empty((Nx,Ny,nelem), dtype=complex)
 154     for pi, offset in enumerate(ss.fullvec_poffsets):
 155         ix, iy = posmap[pi]
 156         arr[ix, iy] = fullvec[offset:offset+nelem]
 157     return np.swapaxes(arr, 0, 1) if swapxy else arr
 158
 159
 160 outfile_tmp = defaultprefix + ".tmp" if a.output is None else a.output + ".tmp"
 161
 162 nelem = len(bspec)
 163 phases = np.exp(1j*np.dot(ss.positions[:,:2], np.array(a.wavevector)))
 164 driving_full = np.zeros((nelem, ss.fecv_size),dtype=complex)
 165 if a.symmetry_adapted is not None:
 166     ss1, ssw1 = ScatteringSystem.create(particles=[Particle((0,0,0), ap.tmgen, bspec)], medium=medium, omega=omega,
 167                     sym=FinitePointGroup(point_group_info[a.symmetry_adapted]))
 168     fvcs1 = np.empty((nelem, nelem), dtype=complex)
 169     y = 0
 170     iris1 = []
 171     for iri1 in range(ss1.nirreps):
 172         for j in range(ss1.saecv_sizes[iri1]):
 173             pvc1 = np.zeros((ss1.saecv_sizes[iri1],), dtype=complex)
 174             pvc1[j] = 1
 175             fvcs1[y] = ss1.unpack_vector(pvc1, iri1)
 176             fvcs1[y] = cleanarray(nicerot(fvcs1[y], copy=False), copy=False)
 177             driving_full[y] = (phases[:, None] * fvcs1[y][None,:]).flatten()
 178             y += 1
 179             iris1.append(iri1)
 180     iris1 = np.array(iris1)
 181 else:
 182     for y in range(nelem):
 183         driving_full[y,y::nelem] = phases
 184
 185
 186 # Apply the driving on the specified slices only
 187 nsp = len(slicepairs)
 188 driving_full_sliced = np.zeros((nsp,) + driving_full.shape, dtype=complex)
 189 p1range = np.arange(nelem)
 190 for spi in range(nsp):
 191     xs, ys = slicepairs[spi]
 192     driven_pi = invposmap[xs, ys].flatten()
 193     driven_y = ((driven_pi * nelem)[:,None] + p1range[None,:]).flatten()
 194     driving_full_sliced[spi][:, driven_y] = driving_full[:, driven_y]
 195
 196 scattered_full = np.zeros((nsp, nelem, ss.fecv_size),dtype=complex)
 197 scattered_ir = [None for iri in range(ss.nirreps)]
 198
 199 ir_contained = np.ones((nsp, nelem, ss.nirreps), dtype=bool)
 200
 201 for iri in range(ss.nirreps):
 202     logging.info("processing irrep %d/%d" % (iri, ss.nirreps))
 203     LU = None # to trigger garbage collection before the next call
 204     translation_matrix = None
 205     LU = ssw.scatter_solver(iri)
 206     logging.info("LU solver created")
 207     #translation_matrix = ss.translation_matrix_packed(wavenumber, iri, BesselType.REGULAR) + np.eye(ss.saecv_sizes[iri])
 208     #logging.info("auxillary translation matrix created")
 209
 210     scattered_ir[iri] = np.zeros((nsp, nelem, ss.saecv_sizes[iri]), dtype=complex)
 211     scattered_ir_unpacked = np.zeros((nsp, nelem, ss.fecv_size), dtype=complex)
 212
 213     for spi in range(nsp):
 214         for y in range(nelem):
 215             ã = driving_full_sliced[spi,y]
 216             ãi = cleanarray(ss.pack_vector(ã, iri), copy=False)
 217             if np.all(ãi == 0):
 218                 ir_contained[spi, y, iri] = False
 219                 continue
 220             Tã  = ssw.apply_Tmatrices_full(ã)
 221             Tãi = ss.pack_vector(Tã, iri)
 222             fi = LU(Tãi)
 223             scattered_ir[iri][spi, y] = fi
 224             scattered_ir_unpacked[spi, y] = ss.unpack_vector(fi, iri)
 225             scattered_full[spi, y] += scattered_ir_unpacked[spi, y]
 226     if a.save_gradually:
 227         iriout = outfile_tmp + ".%d" % iri
 228         np.savez(iriout, iri=iri, meta={**vars(a), 'qpms_version' : qpms.__version__()},
 229                  omega=omega, wavenumber=wavenumber, nelem=nelem, wavevector=np.array(a.wavevector), phases=phases,
 230                  positions = ss.positions[:,:2],
 231                  scattered_ir_packed = scattered_ir[iri],
 232                  scattered_ir_full = scattered_ir_unpacked,
 233                  )
 234         logging.info("partial results saved to %s"%iriout)
 235
 236 t, l, m = bspec.tlm()
 237
 238 if not math.isnan(a.ccd_distance):
 239     logging.info("Computing the far fields")
 240     if math.isnan(a.ccd_size):
 241         a.ccd_size = (50 * a.ccd_distance / (max(Nx*px, Ny*py) *ssw.wavenumber.real))
 242     ccd_size = a.ccd_size
 243     ccd_x = np.linspace(-ccd_size/2, ccd_size/2, a.ccd_resolution)
 244     ccd_y = np.linspace(-ccd_size/2, ccd_size/2, a.ccd_resolution)
 245     ccd_grid = np.meshgrid(ccd_x, ccd_y, (a.ccd_distance,), indexing='ij')
 246     ccd_points = np.swapaxes(np.stack(ccd_grid, axis=-1).squeeze(axis=-2), 0,1) # First axis is y, second is x, because of imshow...
 247     ccd_fields = np.empty((nsp, nelem,) + ccd_points.shape, dtype=complex)
 248     for spi in range(nsp):
 249         for y in range(nelem):
 250             ccd_fields[spi, y] = ssw.scattered_E(scattered_full[spi, y], ccd_points, btyp=BesselType.HANKEL_PLUS)
 251     logging.info("Far fields done")
 252
 253 outfile = defaultprefix + ".npz" if a.output is None else a.output
 254 np.savez(outfile, meta={**vars(a), 'qpms_version' : qpms.__version__()},
 255                  omega=omega, wavenumber=wavenumber, nelem=nelem, wavevector=np.array(a.wavevector), phases=phases,
 256                  positions = ss.positions[:,:2],
 257                  scattered_ir_packed = np.array(scattered_ir, dtype=np.object),
 258                  scattered_full = scattered_full,
 259                  ir_contained = ir_contained,
 260                  t=t, l=l, m=m,
 261                  iris1 = iris1 if (a.symmetry_adapted is not None) else None,
 262                  irnames1 = ss1.irrep_names if (a.symmetry_adapted is not None) else None,
 263                  fvcs1 = fvcs1 if (a.symmetry_adapted is not None) else None,
 264                  #ccd_size = ccd_size if not math.isnan(a.ccd_distance) else None,
 265                  ccd_points = ccd_points if not math.isnan(a.ccd_distance) else None,
 266                  ccd_fields = ccd_fields if not math.isnan(a.ccd_distance) else None,
 267        )
 268 logging.info("Saved to %s" % outfile)
 269
 270
 271 if a.plot or (a.plot_out is not None):
 272
 273
 274     import matplotlib
 275     matplotlib.use('pdf')
 276     from matplotlib import pyplot as plt, cm
 277     from matplotlib.backends.backend_pdf import PdfPages
 278     t, l, m = bspec.tlm()
 279     phasecm = cm.twilight
 280     pmcm = cm.bwr
 281     abscm = cm.plasma
 282
 283     plotfile = defaultprefix + ".pdf" if a.plot_out is None else a.plot_out
 284     pp = PdfPages(plotfile)
 285
 286     for spi in range(nsp):
 287         fig, axes = plt.subplots(nelem, 12 if math.isnan(a.ccd_distance) else 16, figsize=(figscale*(12 if math.isnan(a.ccd_distance) else 16), figscale*nelem))
 288         for yp in range(0,3): # TODO xy-dipoles instead?
 289             axes[0,4*yp+0].set_title("abs / (E,1,%s)" % realdipfieldlabels(yp))
 290             axes[0,4*yp+1].set_title("arg / (E,1,%s)" % realdipfieldlabels(yp))
 291             axes[0,4*yp+2].set_title("Fabs / (E,1,%s)" % realdipfieldlabels(yp))
 292             axes[0,4*yp+3].set_title("Farg / (E,1,%s)" % realdipfieldlabels(yp))
 293         if not math.isnan(a.ccd_distance):
 294             #axes[0,12].set_title("$E_{xy}$ @ $z = %g; \phi$" % a.ccd_distance)
 295             #axes[0,13].set_title("$E_{xy}$ @ $z = %g; \phi + \pi/2$" % a.ccd_distance)
 296             axes[0,12].set_title("$|E_{x}|^2$ @ $z = %g\,\mathrm{m}$" % a.ccd_distance)
 297             axes[0,13].set_title("$|E_{y}|^2$ @ $z = %g\,\mathrm{m}$" % a.ccd_distance)
 298             axes[0,14].set_title("$|E_x + E_y|^2$ @ $z = %g\,\mathrm{m}$" % a.ccd_distance)
 299             axes[0,15].set_title("$|E_{z}|^2$ @ $z = %g\,\mathrm{m}$" % a.ccd_distance)
 300             for gg in range(12,16):
 301                 axes[-1,gg].set_xlabel("$x/\mathrm{m}$")
 302
 303
 304         for y in range(nelem):
 305             fulvec = scattered_full[spi,y]
 306             if a.symmetry_adapted is not None:
 307                 driving_nonzero_y = [j for j in range(nelem) if abs(fvcs1[y,j]) > 1e-5]
 308                 driving_descr = ss1.irrep_names[iris1[y]]+'\n'+', '.join(('$'+cplx_nicestr(fvcs1[y,j])+'$' +
 309     "(%s,%d,%+d)" % (("E" if t[j] == 2 else "M"), l[j], m[j]) for j in
 310     driving_nonzero_y)) # TODO shorten the complex number precision
 311             else:
 312                 driving_descr = "%s,%d,%+d"%('E' if t[y]==2 else 'M', l[y], m[y],)
 313             axes[y,0].set_ylabel(driving_descr)
 314             axes[y,-1].yaxis.set_label_position("right")
 315             axes[y,-1].set_ylabel("$y/\mathrm{m}$\n"+driving_descr)
 316             vecgrid = fullvec2grid(fulvec, swapxy=True)
 317             vecgrid_ff = np.fft.fftshift(np.fft.fft2(vecgrid, axes=(0,1)),axes=(0,1))
 318             lemax = np.amax(abs(vecgrid))
 319             for yp in range(0,3):
 320                 if(np.amax(abs(realdipfields(vecgrid,yp))) > lemax*1e-5):
 321                     axes[y,yp*4].imshow(abs(realdipfields(vecgrid,yp)), vmin=0, interpolation='none')
 322                     axes[y,yp*4].text(0.5, 0.5, '%g' % np.amax(abs(realdipfields(vecgrid,yp))), horizontalalignment='center', verticalalignment='center', transform=axes[y,yp*4].transAxes)
 323                     axes[y,yp*4+1].imshow(np.angle(realdipfields(vecgrid,yp)), vmin=-np.pi, vmax=np.pi, cmap=phasecm, interpolation='none')
 324                     axes[y,yp*4+2].imshow(abs(realdipfields(vecgrid_ff,yp)), vmin=0, interpolation='none')
 325                     axes[y,yp*4+3].imshow(np.angle(realdipfields(vecgrid_ff,yp)), vmin=-np.pi, vmax=np.pi, cmap=phasecm, interpolation='none')
 326                 else:
 327                     for c in range(0,4):
 328                         axes[y,yp*4+c].tick_params(bottom=False, left=False, labelbottom=False, labelleft=False)
 329             if not math.isnan(a.ccd_distance):
 330                 fxye=(-ccd_size/2, ccd_size/2, -ccd_size/2, ccd_size/2)
 331                 e2vmax = np.amax(np.linalg.norm(ccd_fields[spi,y], axis=-1)**2)
 332                 xint = abs(ccd_fields[spi,y,...,0])**2
 333                 yint = abs(ccd_fields[spi,y,...,1])**2
 334                 xyint = abs(ccd_fields[spi,y,...,0] + ccd_fields[spi,y,...,1])**2
 335                 zint = abs(ccd_fields[spi,y,...,2])**2
 336                 xintmax = np.amax(xint)
 337                 yintmax = np.amax(yint)
 338                 zintmax = np.amax(zint)
 339                 xyintmax = np.amax(xyint)
 340                 axes[y, 12].imshow(xint, origin="lower", extent=fxye, cmap=abscm, interpolation='none')
 341                 axes[y, 13].imshow(yint, origin="lower", extent=fxye, cmap=abscm, interpolation='none')
 342                 axes[y, 14].imshow(xyint, origin="lower", extent=fxye, cmap=abscm, interpolation='none')
 343                 axes[y, 15].imshow(zint, origin='lower', extent=fxye, cmap=abscm, interpolation='none')
 344                 axes[y, 12].text(0.5, 0.5, '%g\n%g' % (xintmax,xintmax/e2vmax),
 345                         horizontalalignment='center', verticalalignment='center', transform=axes[y,12].transAxes)
 346                 axes[y, 13].text(0.5, 0.5, '%g\n%g' % (yintmax,yintmax/e2vmax),
 347                         horizontalalignment='center', verticalalignment='center', transform=axes[y,13].transAxes)
 348                 axes[y, 14].text(0.5, 0.5, '%g\n%g' % (xyintmax,xyintmax/e2vmax),
 349                         horizontalalignment='center', verticalalignment='center', transform=axes[y,14].transAxes)
 350                 axes[y, 15].text(0.5, 0.5, '%g\n%g' % (zintmax,zintmax/e2vmax),
 351                         horizontalalignment='center', verticalalignment='center', transform=axes[y,15].transAxes)
 352                 for gg in range(12,16):
 353                     axes[y,gg].yaxis.tick_right()
 354                 for gg in range(12,15):
 355                     axes[y,gg].yaxis.set_major_formatter(plt.NullFormatter())
 356         fig.text(0, 0, str(slicepairs[spi]), horizontalalignment='left', verticalalignment='bottom')
 357         pp.savefig()
 358     annotate_pdf_metadata(pp, scriptname="finiterectlat-constant-driving.py")
 359     pp.close()
 360
 361 exit(0)
 362