def tilts_show(tilts, mp2, iBs, iZs, **kwargs):
''' display beams for a sequence of tilted simulations
tilts,mp2 : tilts array and multislice objects
- iBs,iZs : beam and slice indices (integer allowed)
'''
bopt, zopt, nz = isinstance(iBs, int), isinstance(iZs, int), 1
if bopt: iBs = np.array([iBs])
if zopt: iZs = np.array([iZs])
if isinstance(iZs, slice):
stop = (mp2[0].z.size + 1) * (iZs.stop < 0) + iZs.stop #;print(stop)
nz = int(np.ceil((stop - iZs.start) / iZs.step)) #;print(nz)
else:
nz = np.array(iZs).size
Itbz = np.zeros((tilts.size, iBs.size, nz))
for i, t in enumerate(tilts):
Itbz[i, :, :] = mp2[i].getB(iBs)[iZs, :].T
if zopt and bopt:
plts = [tilts, Itbz[:, 0, iZs], 'bo-']
dsp.stddisp(plts, labs=[r'$\theta(deg)$', '$I$'], **kwargs)
if bopt and not zopt:
csZ = dsp.getCs('Reds', nz)
plts1 = [[tilts, Itbz[:, 0, iZ], [csZ[iZ], 'o-'],
'%dnm' % (z0 / 10)] for iZ, z0 in enumerate(mp2[0].z[iZs])]
dsp.stddisp(plts1, labs=[r'$\theta(deg)$', '$I$'], **kwargs)
if zopt and not bopt:
cs = dsp.getCs('Spectral', iBs.size)
# for iZ in iZs:Iz[:,iZ]/=Iz[:,iZ].max()
plts2 = [[tilts, Itbz[:, i], [cs[i], 'o-'],
'iB=%d' % (iB)] for i, iB in enumerate(iBs)]
dsp.stddisp(plts2, labs=[r'$\theta(deg)$', '$I$'], **kwargs)
def _test_get_arrow():
#N = 1
U = np.array([[0,1,1],[1,1,1],[1,0,1]])
x0 = [0,0,1]
cs,lw = dsp.getCs('viridis',U.shape[0]),2
plots,surfs = [],[]
for u,cu in zip(U,cs) :
(x,y,z),(xl,yl,zl) = get_arrow_3d(u,x0,rc=0.1,h=0.2)
plots += [[xl,yl,zl,cu]]
surfs += [[x,y,z,cu,None,lw,cu]]
dsp.stddisp(rc='3d',plots=plots,surfs=surfs,lw=lw,pOpt='e')#,texts=txts)
def Xz_show(self, iZs=1, cmap='jet', **kwargs):
'''Show wave propagation for slices iZs '''
if isinstance(iZs, int): iZs = slice(0, self.z.size, iZs)
z = self.z[iZs]
Px = self.psi_xz[iZs, :]
cs = dsp.getCs(cmap, z.size)
plts = [[self.x, Px[i, :], cs[i]] for i in range(z.size)]
return dsp.stddisp(plts,
labs=[r'$x(\AA)$', r'$|\Psi(x)|^2$'],
imOpt='hc',
caxis=[z.min(), z.max()],
cmap=cmap,
**kwargs)
def plot_beam_thickness(beams, rip='I', linespec='-', cm='Greens', **kwargs):
''' plot the beams as function of thickness
- beams : beams info from import_beams or file containing them
- rip flags : I(Intens),r(real),i(imag)
'''
hk, t, re, im, Ib = beams
nbs = len(hk)
csp, csr, csi, plts = dsp.getCs(cm, nbs), dsp.getCs('Blues',
nbs), dsp.getCs(
'Reds', nbs), []
for i in range(nbs):
if 'I' in rip:
plts += [[
t, Ib[i], [csp[i], '%sx' % linespec],
'$I_{%s}$' % (hk[i])
]]
if 'r' in rip: plts += [[t, re[i], csr[i], '$re$']]
if 'i' in rip: plts += [[t, im[i], csi[i], '$im$']]
return dsp.stddisp(plts,
lw=2,
labs=['$thickness(\AA)$', '$I_{hk}$'],
**kwargs) #,xylims=[0,t.max(),0,5])
def Tz_show(self,
iSz=slice(0, None, 1),
opts='',
Vopt='VT',
cmaps=['Greens', 'Blues', 'Reds'],
**kwargs):
'''Show Transmission function
- opts or Vopt : 'V(potential) 'T'(Transmission) 'l'(slice legend)
if 'l' is selected
'''
if opts: Vopt = opts
if isinstance(iSz, int): iSz = [iSz]
if isinstance(iSz, slice): iSz = list(np.arange(self.ns)[iSz])
if isinstance(iSz, list): N = len(iSz)
if isinstance(cmaps, str): cmaps = [cmaps] * 3
cs1, cs2, cs3 = dsp.getCs(cmaps[0],
N), dsp.getCs(cmaps[1],
N), dsp.getCs(cmaps[2], N)
plts, legElt = [], {}
if 'V' in Vopt:
plts += [[self.x, self.Vz[iSz[i], :].T, cs1[i]] for i in range(N)]
legElt[r'$V_z(kV\AA)$'] = [cs1[int(N / 2)], '-']
if 'T' in Vopt:
plts += [[self.x, self.T.real[iSz[i], :].T, cs2[i]]
for i in range(N)]
plts += [[self.x, self.T.imag[iSz[i], :].T, cs3[i]]
for i in range(N)]
legElt['$re(T)$'] = [cs2[int(N / 2)], '-']
legElt['$im(T)$'] = [cs3[int(N / 2)], '-']
return dsp.stddisp(
plts,
labs=[
r'$x(\AA)$', ''
], #title='Projected potential $V_z$, Transmission function $T$',
legElt=legElt,
**kwargs)
def Qz_show(self, iZs=1, opts='', cmap='jet', **kwargs):
'''Show wave propagation in q space for slices iZs
- opts : O(include Origin), N(Normalize with max), S(Shape normalize) l(show label)
'''
# if isinstance(iZs,int):iZs=slice(0,iZs,self.z.size)
if isinstance(iZs, int): iZs = slice(0, None, iZs)
# if isinstance(iZs,list):iZs=slice(0,iZs,self.z.size)
Pqs = np.zeros((self.psi_qz[iZs, :].shape))
z = self.z[iZs]
Pqs = self.psi_qz[iZs, :].copy()
if 'S' in opts:
if 'O' not in opts: Pqs[:, 0] = 0 # do not show central beam
for i in range(z.size):
Pqs[i, :] /= np.sum(Pqs[i, :])
else:
if 'O' not in opts: Pqs[:, 0] = 0 # do not show central beam
if 'N' in opts:
for i in range(z.size):
Pqs[i, :] /= Pqs[i, :].max()
q = fft.fftshift(self.q.copy())
Pqs = [fft.fftshift(Pqs[i, :]) for i in range(z.size)]
cs = dsp.getCs(cmap, z.size)
if 'l' in opts:
plts = [[q, Pqs[i], cs[i], '$%.1fA$' % z[i]]
for i in range(z.size)]
return dsp.stddisp(plts,
labs=[r'$q(\AA^{-1})$', r'$|\Psi(q)|^2$'],
imOpt='c',
caxis=[z.min(), z.max()],
cmap=cmap,
axPos='V',
**kwargs)
else:
plts = [[q, Pqs[i], cs[i]] for i in range(z.size)]
return dsp.stddisp(plts,
labs=[r'$q(\AA^{-1})$', r'$|\Psi(q)|^2$'],
imOpt='c',
caxis=[z.min(), z.max()],
cmap=cmap,
axPos='V',
**kwargs)
def Bz_show(self,
iBs='O',
tol=1e-3,
cmap='jet',
sym_opt=False,
plts=[],
**kwargs):
'''Show selected beam iBs as function of thickness(see getB)
- sym_opt : special display for symmetry pairs
'''
# print(iBs)
iBs, Ib = self.getB(iBs, tol, v=1) #;print(iBs)
h = ['%d_{%d}' % (i / self.Nx, i % self.Nx) for i in iBs]
cs = dsp.getCs(cmap, iBs.size)
plts = [[self.z, Ib[:, i], [cs[i], '-'],
'$%s$' % h[i]] for i, iB in enumerate(iBs)]
if sym_opt:
iBs, Ib = self.getB(-iBs, tol, v=1)
h = ['%d_{%d}' % (i / self.Nx, i % self.Nx) for i in iBs]
plts += [[self.z, Ib[:, i], [cs[i], '--'],
'$%s$' % h[i]] for i, iB in enumerate(iBs)]
return dsp.stddisp(plts, labs=[r'$z(\AA)$', r'$I_b$'], **kwargs)
def Li_wobble(name,nvals=np.inf,**kwargs):
wobbles = [0.001,0.01,0.1]
nvals = min(nvals,len(wobbles))
wobbles = np.array(wobbles)[:nvals]
plts,cs = [],dsp.getCs('coolwarm',nvals)
tol,Iopt = 1e-5,'Incsl',
for i,val in zip(range(nvals),wobbles):
tail = str(i).zfill(int(wobbles.size/10)+1)
data = Li_xyz(name,n=[0,0,1],dopt=1,wobble=val,tail='_'+tail)
multi=mupy.Multislice(name,tail=tail,data=data,
mulslice=0,keV=100,TDS=True,n_TDS=20,
NxNy=1024,slice_thick=1.5,Nhk=5,repeat=[8,8,100],
**kwargs)
q,I = multi.azim_avg(out=1,opt='',Iopt=Iopt,tol=tol)
plts+=[[q,I,cs[i],'w=%.3f' %val]]
multi.pattern(rings=[0.3,0.6,1,2],Iopt=Iopt,tol=tol,
imOpt='ch',cmap='binary',axPos=[0.17,0.13,0.75,0.75],lw=2,xylims=[0,3,0,3],
opt='sc',name='docs_fig/wobble_effect%d.png' %i,pOpt='tpX',
caxis = [-5,0])
dsp.stddisp(plts,labs=[r'$q(\AA^{-1})$','$I_q$'],lw=2,fonts={'leg':20},
xylims=['x',0,4,-5,0],opt='sc',name='docs_fig/wobble_effectIavg.svg')