224 lines
8.4 KiB
Python
Executable File
224 lines
8.4 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
import argparse, re, random, string
|
|
from scipy.constants import hbar, e as eV, pi, c
|
|
|
|
parser = argparse.ArgumentParser()
|
|
#TODO? použít type=argparse.FileType('r') ?
|
|
parser.add_argument('--output', '-o', action='store', help='Path to output PDF')
|
|
parser.add_argument('--nSV', action='store', metavar='N', type=int, default=1, help='Draw N minimum singular values')
|
|
parser.add_argument('--bitmap', action='store_true', help='Create an interpolated bitmap rather than a scatter plot.')
|
|
#parser.add_argument('--eVfreq', action='store', required=True, type=float, help='Frequency in eV')
|
|
parser.add_argument('inputfile', nargs='+', help='Npz file(s) generated by dispersion_chunks.py or other script')
|
|
pargs=parser.parse_args()
|
|
print(pargs)
|
|
|
|
#freq = eVfreq*eV/hbar
|
|
|
|
pdfout = pargs.output if pargs.output else '%s.pdf' % pargs.inputfile[-1]
|
|
print(pdfout)
|
|
|
|
svn = pargs.nSV
|
|
|
|
# -----------------finished basic CLI parsing (except for op arguments) ------------------
|
|
import time
|
|
begtime=time.time()
|
|
|
|
import qpms
|
|
from matplotlib.path import Path
|
|
import matplotlib.patches as patches
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
import os, sys, warnings, math
|
|
from matplotlib import pyplot as plt
|
|
from matplotlib.backends.backend_pdf import PdfPages
|
|
from scipy import interpolate
|
|
|
|
# We do not want to import whole qpms, so copy and modify this only fun needed
|
|
def nelem2lMax(nelem):
|
|
lMax = round(math.sqrt(1+nelem) - 1)
|
|
if ((lMax < 1) or ((lMax + 2) * lMax != nelem)):
|
|
raise
|
|
else:
|
|
return lMax
|
|
|
|
|
|
|
|
nx = None
|
|
s3 = math.sqrt(3)
|
|
|
|
|
|
# read data from files
|
|
lMax = None
|
|
epsilon_b = None
|
|
hexside = None
|
|
karrlist = list()
|
|
svTElist = list()
|
|
svTMlist = list()
|
|
omegalist = list()
|
|
records = 0
|
|
for filename in pargs.inputfile:
|
|
npz = np.load(filename)
|
|
lMaxRead = npz['metadata'][()]['lMax'] if 'lMax' in npz['metadata'][()] else nelem2lMax(npz['sTE'].shape[1] / 2)
|
|
if lMax is None: lMax = lMaxRead
|
|
elif lMax != lMaxRead: raise
|
|
if epsilon_b is None: epsilon_b = npz['metadata'][()]['epsilon_b']
|
|
elif epsilon_b != npz['metadata'][()]['epsilon_b'] : raise
|
|
if hexside is None: hexside = npz['metadata'][()]['hexside']
|
|
elif hexside != npz['metadata'][()]['hexside'] : raise
|
|
omegalist.append(npz['omega'][()])
|
|
karrlist.append(np.array(npz['klist']))
|
|
svTElist.append(np.array(npz['sTE'][:,-svn:]))
|
|
svTMlist.append(np.array(npz['sTM'][:,-svn:]))
|
|
records += 1
|
|
npz.close()
|
|
|
|
# sort by frequencies
|
|
omegas = set(omegalist)
|
|
print(omegas)
|
|
k = dict()
|
|
svTE = dict()
|
|
svTM = dict()
|
|
for omega in omegas:
|
|
k[omega] = list()
|
|
svTE[omega] = list()
|
|
svTM[omega] = list()
|
|
for i in range(records):
|
|
omega = omegalist[i]
|
|
k[omega].append(karrlist[i])
|
|
svTE[omega].append(svTElist[i])
|
|
svTM[omega].append(svTMlist[i])
|
|
# concatenate arrays for each frequency
|
|
for omega in omegas:
|
|
k[omega] = np.concatenate(k[omega])
|
|
svTE[omega] = np.concatenate(svTE[omega])
|
|
svTM[omega] = np.concatenate(svTM[omega])
|
|
|
|
# ... that was for the slices. TODO fill also the righternmost plot with the calculated (which?) modes.
|
|
|
|
pdf = PdfPages(pdfout)
|
|
|
|
# In[3]:
|
|
|
|
cdn = c/ math.sqrt(epsilon_b)
|
|
#my, ny = qpms.get_mn_y(lMax)
|
|
#nelem = len(my)
|
|
nelem = lMax * (lMax + 2)
|
|
|
|
|
|
''' The new pretty diffracted order drawing '''
|
|
maxlayer_reciprocal=4
|
|
cdn = c/ math.sqrt(epsilon_b)
|
|
bz_0 = np.array((0,0,))
|
|
bz_K1 = np.array((1.,0))*4*np.pi/3/hexside/s3
|
|
bz_K2 = np.array((1./2.,s3/2))*4*np.pi/3/hexside/s3
|
|
bz_M = np.array((3./4, s3/4))*4*np.pi/3/hexside/s3
|
|
|
|
# reciprocal lattice basis
|
|
B1 = 2* bz_K1 - bz_K2
|
|
B2 = 2* bz_K2 - bz_K1
|
|
|
|
k2density = 100
|
|
k0Mlist = bz_0 + (bz_M-bz_0) * np.linspace(0,1,k2density)[:,nx]
|
|
kMK1list = bz_M + (bz_K1-bz_M) * np.linspace(0,1,k2density)[:,nx]
|
|
kK10list = bz_K1 + (bz_0-bz_K1) * np.linspace(0,1,k2density)[:,nx]
|
|
k0K2list = bz_0 + (bz_K2-bz_0) * np.linspace(0,1,k2density)[:,nx]
|
|
kK2Mlist = bz_K2 + (bz_M-bz_K2) * np.linspace(0,1,k2density)[:,nx]
|
|
k2list = np.concatenate((k0Mlist,kMK1list,kK10list,k0K2list,kK2Mlist), axis=0)
|
|
kxmaplist = np.concatenate((np.array([0]),np.cumsum(np.linalg.norm(np.diff(k2list, axis=0), axis=-1))))
|
|
|
|
centers2=qpms.generate_trianglepoints(maxlayer_reciprocal, v3d = False, include_origin= True)*4*np.pi/3/hexside
|
|
rot90 = np.array([[0,-1],[1,0]])
|
|
centers2=np.dot(centers2,rot90)
|
|
|
|
import matplotlib.pyplot as plt
|
|
import matplotlib
|
|
from matplotlib.path import Path
|
|
import matplotlib.patches as patches
|
|
cmap = matplotlib.cm.prism
|
|
colormax = np.amax(np.linalg.norm(centers2,axis=0))
|
|
|
|
|
|
for omega in sorted(omegas):
|
|
klist = k[omega]
|
|
if pargs.bitmap:
|
|
minx = np.amin(klist[:,0])
|
|
maxx = np.amax(klist[:,0])
|
|
miny = np.amin(klist[:,1])
|
|
maxy = np.amax(klist[:,1])
|
|
l = klist.shape[0]
|
|
meshstep = math.sqrt((maxy - miny) * (maxx - minx) / l) / 9
|
|
x = np.linspace(minx, maxx, (maxx-minx) / meshstep)
|
|
y = np.linspace(miny, maxy, (maxy-miny) / meshstep)
|
|
fullshape = np.broadcast(x[:,nx],y[nx,:]).shape
|
|
flatx = np.broadcast_to(x[:,nx], fullshape).flatten
|
|
flaty = np.broadcast_to(y[nx,:], fullshape).flatten
|
|
minsvTElist = svTE[omega]
|
|
minsvTMlist = svTM[omega]
|
|
for minN in reversed(range(svn)):
|
|
f, axes = plt.subplots(1,3, figsize=(20,4.8))
|
|
ax = axes[0]
|
|
if pargs.bitmap:
|
|
interpolator = interpolate.interp2d(klist[:,0], klist[:,1], np.abs(minsvTElist[:,minN]))
|
|
z = interpolator(flatx, flaty)
|
|
z.reshape(fullshape)
|
|
sc = ax.pcolormesh(x[:,nx],y[nx,:],z)
|
|
else:
|
|
sc = ax.scatter(klist[:,0], klist[:,1], c = np.clip(np.abs(minsvTElist[:,minN]),0,1), lw=0)
|
|
for center in centers2:
|
|
circle=plt.Circle((center[0],center[1]),omega/cdn, facecolor='none', edgecolor=cmap(np.linalg.norm(center)/colormax),lw=0.5)
|
|
ax.add_artist(circle)
|
|
verts = [(math.cos(math.pi*i/3)*4*np.pi/3/hexside/s3,math.sin(math.pi*i/3)*4*np.pi/3/hexside/s3) for i in range(6 +1)]
|
|
codes = [Path.MOVETO,Path.LINETO,Path.LINETO,Path.LINETO,Path.LINETO,Path.LINETO,Path.CLOSEPOLY,]
|
|
path = Path(verts, codes)
|
|
patch = patches.PathPatch(path, facecolor='none', edgecolor='black', lw=1)
|
|
ax.add_patch(patch)
|
|
ax.set_xticks([])
|
|
ax.set_yticks([])
|
|
ax.title.set_text('E in-plane ("TE"), %d. lowest SV' % minN)
|
|
f.colorbar(sc,ax=ax)
|
|
|
|
|
|
ax = axes[1]
|
|
if pargs.bitmap:
|
|
interpolator = interpolate.interp2d(klist[:,0], klist[:,1], np.abs(minsvTMlist[:,minN]))
|
|
meshstep = math.sqrt((maxy - miny) * (maxx - minx) / l) / 9
|
|
z = interpolator(flatx, flaty)
|
|
z.reshape(fullshape)
|
|
sc = ax.pcolormesh(x[:,nx],y[nx,:],z)
|
|
else:
|
|
sc = ax.scatter(klist[:,0], klist[:,1], c = np.clip(np.abs(minsvTMlist[:,minN]),0,1), lw=0)
|
|
for center in centers2:
|
|
circle=plt.Circle((center[0],center[1]),omega/cdn, facecolor='none', edgecolor=cmap(np.linalg.norm(center)/colormax),lw=0.5)
|
|
ax.add_artist(circle)
|
|
verts = [(math.cos(math.pi*i/3)*4*np.pi/3/hexside/s3,math.sin(math.pi*i/3)*4*np.pi/3/hexside/s3) for i in range(6 +1)]
|
|
codes = [Path.MOVETO,Path.LINETO,Path.LINETO,Path.LINETO,Path.LINETO,Path.LINETO,Path.CLOSEPOLY,]
|
|
path = Path(verts, codes)
|
|
patch = patches.PathPatch(path, facecolor='none', edgecolor='black', lw=1)
|
|
ax.add_patch(patch)
|
|
ax.set_xticks([])
|
|
ax.set_yticks([])
|
|
ax.title.set_text('E perpendicular ("TM"), %d. lowest SV' % minN)
|
|
f.colorbar(sc,ax=ax)
|
|
|
|
ax = axes[2]
|
|
for center in centers2:
|
|
ax.plot(kxmaplist, np.linalg.norm(k2list-center,axis=-1)*cdn, '-', color=cmap(np.linalg.norm(center)/colormax))
|
|
|
|
#ax.set_xlim([np.min(kxmlarr),np.max(kxmlarr)])
|
|
#ax.set_ylim([np.min(omegalist),np.max(omegalist)])
|
|
xticklist = [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]]
|
|
ax.set_xticks(xticklist)
|
|
for xt in xticklist:
|
|
ax.axvline(xt, ls='dotted', lw=0.3,c='k')
|
|
ax.set_xticklabels(['Γ', 'M', 'K', 'Γ', 'K\'','M'])
|
|
ax.axhline(omega, c='black')
|
|
ax.set_ylim([0,5e15])
|
|
ax2 = ax.twinx()
|
|
ax2.set_ylim([ax.get_ylim()[0]/eV*hbar,ax.get_ylim()[1]/eV*hbar])
|
|
|
|
pdf.savefig(f)
|
|
|
|
pdf.close()
|
|
|