DFT plots of the bandstructure for Bi$_2$Te$_3$, Bi$_2$Se$_3$ and Sb$_2$Te$_3$¶

from aiida import orm, load_profile
from aiida_kkr.tools import plot_kkr, find_parent_structure
import numpy as np
import matplotlib.pyplot as plt
_ = load_profile()

from aiida_kkr.calculations import KkrCalculation
from masci_tools.util.constants import RY_TO_EV

def load_data_and_store_in_extra(bs, reload=False, ef=None):
    from tqdm import tqdm
    if 'e_k_dall_uuid' not in bs.extras or reload:
        print('load data')
        try:
            ret = bs.get_outgoing(node_class=KkrCalculation).first().node.outputs.retrieved
        except:
            ret = bs.outputs.retrieved
        dall = []
        for name in tqdm(np.sort([i for i in ret.list_object_names() if 'qdos.' in i and 'dat_' not in i])):
            if 'npy' not in name:
                # print(name)
                with ret.open(name, 'r') as _f:
                    d = np.loadtxt(_f)
                dall.append(d)
        dall = np.array(dall)
        ne = len(set(dall[0][:,0]))
        nk = len(dall[0])//ne
        try:
            Nsites = len(find_parent_structure(bs.inputs.remote_data).sites)
        except:
            Nsites = len(find_parent_structure(bs).sites)
        dall = dall.reshape(2*Nsites, ne, nk, -1)

        if ef is None:
            try:
                ef = bs.inputs.remote_data.get_incoming(node_class=KkrCalculation).first().node.outputs.output_parameters['fermi_energy']
            except:
                ef = bs.inputs.parent_folder.get_incoming(node_class=KkrCalculation).first().node.outputs.output_parameters['fermi_energy']

        e = (dall[0, :, 0, 0]-ef)*RY_TO_EV
        try:
            k = bs.get_outgoing(node_class=KkrCalculation).first().node.inputs.kpoints.get_kpoints(cartesian=True)
            klbl = bs.get_outgoing(node_class=KkrCalculation).first().node.inputs.kpoints.labels
        except:
            k = bs.inputs.kpoints.get_kpoints(cartesian=True)
            klbl = bs.inputs.kpoints.labels
        if klbl is None:
            klbl = []
            
        print('save data to extra')
        data_node = orm.ArrayData()
        data_node.set_array('e', e)
        data_node.set_array('k', k)
        data_node.set_array('klbl_id', np.array([i[0] for i in klbl]))
        data_node.set_array('klbl_txt', np.array([i[1] for i in klbl]))
        data_node.set_array('dall', dall)
        data_node.store()
        bs.set_extra('e_k_dall_uuid', data_node.uuid)

    print('load data from extra')
    data_node = orm.load_node(bs.extras['e_k_dall_uuid'])
    (e, k, klbl_id, klbl_txt, dall) = data_node.get_array('e'), data_node.get_array('k'), data_node.get_array('klbl_id'), data_node.get_array('klbl_txt'), data_node.get_array('dall')

  
    return e, k, dall, klbl_id, klbl_txt

def plot_fitting_data(filename):
    from pandas import DataFrame, read_csv

    with open(filename) as _f:
        df = read_csv(_f)

    df_x = df.where(df[df.columns[0]]>0)
    plt.plot(df_x[df_x.columns[0]], df_x.E, color='C0', marker='o', ms=2)

    df_y = df.where(df[df.columns[1]]>0)
    plt.plot(df_y.ky+0.3, df_y.E, color='C1', marker='o', ms=2)

    df_z = df.where(df[df.columns[2]]>0)
    plt.plot(df_z.kz+0.6, df_z.E, color='C2', marker='o', ms=2)

    plt.axvline(0.3, color='grey')
    plt.axvline(0.6, color='grey')
    plt.xticks([0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9],
               [r'$\Gamma$', 0.1, 0.2, r'$\rightarrow k_x$|$\Gamma$'+6*' ', 0.1, 0.2, r'$\rightarrow k_y$|$\Gamma$'+6*' ', 0.1, 0.2, r'$\rightarrow k_z$']
              )
    plt.xlim(0, 0.9)
    plt.ylim(-1, 1)
    plt.ylabel(r'$E-E_F$ (eV)')
    plt.xlabel(r'$k$ ($\AA^{-1}$)')
    plt.title('Fitting data')
    plt.show()

def plot_overview_bandstruc(uuid, title, **kwargs):
    bs_zoom = orm.load_node(uuid)

    plot_kkr(bs_zoom, silent=True, cmap='binary', noshow=True, show_cbar=False, **kwargs)
    plt.clim(0,6)
    plt.title(title)
    plt.ylabel(r'$E-E_F$ (eV)')
    plt.show()

def plot_sbz_projection(uuid, title, eshift=0, clims=(0, 1e4)):
    bs_zoom_kx_ky_kz_int = orm.load_node(uuid)                                     

    e, k, dall, klbl_id, klbl_txt = load_data_and_store_in_extra(bs_zoom_kx_ky_kz_int, reload=False)
    e += eshift
    dall = dall.reshape(dall.shape[0], dall.shape[1], 20, -1, dall.shape[-1])
    dall = np.sum(dall, axis=0)
    k = np.sum(k.reshape(20,-1,3)[0][:,:2], axis=1)

    plt.figure()
    plt.pcolormesh(k, e, np.sum(dall, axis=1)[:,:,4], cmap='binary')
    plt.xlabel(r'$k_x \leftarrow \qquad$ $k$ ($\AA^{-1}$) $\qquad \rightarrow k_y$')
    plt.ylabel('$E-E_\mathrm{F}$ (eV)')
    plt.title(title)
    plt.clim(clims)
    plt.axhline(0, lw=0.5, ls='--', color='r')
    plt.show()