qpms/misc/hexdisplot.py

224 lines
8.4 KiB
Python
Raw Permalink Normal View History

#!/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()