df = SpinDataFrame(path = r'E:\Research\graphene_example\ISPIN_2\dos\sigm0_01\vasprun.xml',
bands=[3,4],elements=[[0],[1]], orbs= [(0,1),(1,2,3)],scale_data=True)
df
df_m = df.copy()
df_m[['kx','ky','kz']] = -df[['kx','ky','kz']]
df1 = df.join(df_m)
df2 = df1.masked('eu',2,0.05,n=100,band=5)
df2
ax1,ax2 = api.get_axes((3,3),ncols=2,widths=[20,1],hspace=0.1,wspace=0.1)
df2.poscar.set_bz(primitive = False)
df2.splot('kx','ky','eu',s=5,ax=ax1,color='k')
df2.splot('kx','ky','eu','su_1',arrows=['','su_1','su_1'],every=6,quiver_kws= dict(cmap='brg',scale=20),ax=ax1)
df2.colorbar(ax2)
df2.poscar.splot_bz(plane='xy',ax=ax1)
ax1.set_title('Band 5 at $E = 2 \pm 0.05$ eV')
ax2.yaxis.set_label_position("right")
ax2.set_ylabel('$S_\\uparrow(C_1-p)$')
df1['kz'] = 0
ax = df1.splot3d('kx','ky','eu','su_1',s=10,cmap='inferno')
df1.colorbar()
df3 = df1[df1['band'] == 5]
ax1 = df3.splot3d('kx','ky','kz',arrows=['','','su_1','eu'],norm=0.1,every=2,quiver_kws= dict(mutation_scale=5,arrowstyle='-|>',cmap='brg',zorder=3))
df3.poscar.splot_bz(ax=ax1).set_axis_off()
df3.colorbar()
ax.set(xlabel='$k_x $',ylabel='$k_y$',zlabel='$E_\\uparrow$')
ax1.view_init(elev=60,azim=45)
ax.set_title('$E_\\uparrow$ of Band 4,5 colored by $S_\\uparrow (C_1-p)$')
ax1.set_title('$S_\\uparrow (C_1-p)$ colored by $E_\\uparrow$ of Band 5')
ax3d = api.get_axes(axes_3d=True)
data = df2.get_data().band_5
KPTS = np.array([data.kx,data.ky,data.kz]).T
KPTS[:,2] = 0
ax3d.scatter(data.kx,data.ky,data.kz,c = data.eu,s=2,cmap='inferno')
df2.poscar.splot_bz(ax=ax3d,fill=False)
UVW = np.array([data.su_0, data.sd_1, data.su_1]).T
api.fancy_quiver3d(*KPTS[::8].T,*UVW[::8].T/5,C = np.abs(UVW)[::8]/UVW.max(),arrowstyle='-|>',ax=ax3d)
ax3d.set_axis_off()
ax3d.view_init(elev=60,azim=90)
import plotly.graph_objects as go
fig = go.Figure()
fig.add_scatter3d(x = KPTS[:,0],y = KPTS[:,1], z = KPTS[:,2],marker =dict(color= np.abs(UVW)/UVW.max(),size=2),mode='markers')
df2.poscar.iplot_bz(fig=fig)
fig.add_cone(x = KPTS[::8,0],y = KPTS[::8,1], z = KPTS[::8,2],u = UVW[::8,0],v = UVW[::8,1], w = UVW[::8,2],sizemode='absolute',sizeref=50,colorscale='magma',showscale=False)
api.iplot2html(fig, modebar= False)
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import numpy as np
import matplotlib.tri as mtri
# Creating dataset
data = df1.get_data().band_4
x = data.kx
y = data.ky
z = data.eu
tri = mtri.Triangulation(x, y)
new_x = tri.x[tri.triangles].mean(axis=1)
new_y = tri.y[tri.triangles].mean(axis=1)
new_color = griddata((x,y),data.su_1,(new_x,new_y),method='linear')
new_color = np.abs((new_color - new_color.min())/np.ptp(new_color))
verts = np.array([tri.x[tri.triangles],tri.y[tri.triangles],z[tri.triangles]]).transpose([1,2,0])
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
import pivotpy as pp
col = Poly3DCollection(verts,color=plt.get_cmap('brg')(new_color))
ax = pp.get_axes(axes_3d=True)
ax.add_collection3d(col)
ax.autoscale()
#df1.poscar.splot_bz(ax=ax,fill=False)
ax.set(zlim=(-6,0),xlim=(-0.27,0.27),ylim=(-0.24,0.24),
zlabel='$E_\\uparrow$',xlabel='$k_x$',ylabel='$k_y$',
title='Color by $S_\\uparrow$ over $E_\\uparrow$ surface')
ax.view_init(elev=45,azim=45)
cax = ax.get_figure().add_axes([0.9,0.1,0.03,0.8])
pp.add_colorbar(ax,'brg',cax = cax, ticks=np.linspace(data.su_1.min(),data.su_1.max(),5))
df1.get_data().to_json('../FermiData.json')