qpms/misc/hexdisplot.py

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