Figures normal state¶

# helper functions
def plot_bs(bs_BdG0, show=True, xlim=None, axhline=True, escale=-1., clim=1e3, title='Bi2Te3-xTe2', empty_ticks=False, no_yticks=False,
             colorbar=False, diffplot=True, BTfree=False,  Bi2=False,  XTe2=False, verbose=False, ylim=None,
             s0 = 1, vmin = -500, vmax = 500, smin = 200
            ):
    
    from matplotlib.colors import LinearSegmentedColormap
    
    if not bs_BdG0.is_terminated:
        if verbose:
            print('bs not done yet', title)
        return
    
    # load data
    data = load_data(bs_BdG0, verbose=verbose)
    if len(data)<7:
        e, k, klbl, dBi2Te3, dBi2Te3_free, dXTe2 = data
        dall = dBi2Te3 + dBi2Te3_free + dXTe2
    else:
        e, k, klbl, dBi2Te3, dBi2Te3_free, dXTe2, dBi2 = data
        dall = dBi2Te3 + dBi2Te3_free + dXTe2 + dBi2
    if verbose:
        print('loaded data')
        
    e = e.copy() * escale
    
    
    e, k = np.meshgrid(e, k)

    
    d1 = np.sum(dBi2Te3[:,:,4:], axis=-1)
    d1_free = np.sum(dBi2Te3_free[:,:,4:], axis=-1)
    d2 = np.sum(dXTe2[:,:,4:], axis=-1)
    if 'dBi2' in locals():
        d3 = np.sum(dBi2[:,:,4:], axis=-1)
    else:
        d3 = np.zeros_like(d1)
    
    
    # plot bandstructure
    if diffplot:
        # d1 = np.sum(dBi2Te3[:,:,4:], axis=-1)
        # d2 = np.sum(dXTe2[:,:,4:], axis=-1)
        # plt.pcolormesh(k, e, d1-d2, cmap='seismic')
        
        c = (d1-d2).transpose().reshape(-1)
        s = (d1+d2+d3).transpose().reshape(-1)
        m = np.where(abs(s)>smin); s = s[m]; s /= s.max(); s *= s0
        cmap = LinearSegmentedColormap.from_list('BlueGreyRed', ['Blue', 'Grey', 'Red'])
        plt.scatter(k.reshape(-1)[m], e.reshape(-1)[m], c=c[m], s=s, cmap=cmap, vmin=vmin, vmax=vmax, rasterized=True)
    
    elif BTfree:
        # plt.pcolormesh(k, e, np.sum(dBi2Te3_free[:,:,4:], axis=-1), cmap='binary')
        
        c = (d1_free).transpose().reshape(-1); s = c
        m = np.where(abs(s)>smin); s = s[m]; s /= s.max(); s *= s0
        plt.scatter(k.reshape(-1)[m], e.reshape(-1)[m], c=c[m], s=s, cmap='binary', vmin=0, vmax=vmax, rasterized=True)
        
    elif XTe2:
        # plt.pcolormesh(k, e, np.sum(dXTe2[:,:,4:], axis=-1), cmap='binary')
        c = (d2).transpose().reshape(-1); s = c
        m = np.where(abs(s)>smin); s = s[m]; s /= s.max(); s *= s0
        plt.scatter(k.reshape(-1)[m], e.reshape(-1)[m], c=c[m], s=s, cmap='binary', vmin=0, vmax=vmax, rasterized=True)
    elif Bi2 and 'dBi2' in locals():
        # plt.pcolormesh(k, e, np.sum(dBi2[:,:,4:], axis=-1), cmap='binary')
        c = (d3).transpose().reshape(-1); s = c
        m = np.where(abs(s)>smin); s = s[m]; s /= s.max(); s *= s0
        plt.scatter(k.reshape(-1)[m], e.reshape(-1)[m], c=c[m], s=s, cmap='binary', vmin=0, vmax=vmax, rasterized=True)
    else:
        return
    
    # labels etc
    if colorbar:
        plt.colorbar()
    plt.clim(-clim, clim)
    if not diffplot:
        plt.clim(0, clim)
    if klbl is not None:
        plt.xticks([i[0] for i in klbl], [i[1].replace('GAMMA', '$\Gamma$') for i in klbl])
        if empty_ticks:
            plt.xticks([i[0] for i in klbl], ['' for i in klbl])
    if no_yticks:
        plt.yticks([i for i in np.linspace(-1, 1, 11)], ['' for i in np.linspace(-1, 1, 11)])
    else:
        plt.yticks([i for i in np.linspace(-1, 1, 11)])
    if axhline:
        plt.axhline(0, ls='--', lw=1, color='k')
    if xlim is not None:
        plt.xlim(xlim[0], xlim[1])
    elif klbl is not None:
        plt.xlim(klbl[0][0], klbl[-1][0])
    if ylim is not None:
        plt.ylim(ylim[0], ylim[1])
    plt.title(title)

    if show:
        plt.show()
        
def plot_bs_all(bs_PtTe2, bs_PdTe2, bs_NbTe2, bs_TiTe2, bs_PtTe2_Bi2, bs_PdTe2_Bi2, bs_NbTe2_Bi2, bs_TiTe2_Bi2,
                ylim = (-1,1), xlim=None, clim = 1e3,
                colorbar=False, diffplot=True, BTfree=False, XTe2=False, Bi2=False, verbose=False, suptitle=None, show=True):
    
    pkwargs = dict(colorbar=colorbar, diffplot=diffplot, BTfree=BTfree, Bi2=Bi2, XTe2=XTe2, verbose=verbose,
                   show=False, clim = clim, s0 = 3, vmin = -500, vmax = 500, smin = 300, ylim=ylim, xlim=xlim,
                  )
    
    plt.figure(figsize=(16,9))
    plt.subplot(2,4,1)
    plot_bs(bs_PtTe2, title='PtTe$_2$/Bi$_2$Te$_3$', empty_ticks=True, **pkwargs)
    
    plt.subplot(2,4,2)
    plot_bs(bs_PdTe2, title='PdTe$_2$/Bi$_2$Te$_3$', empty_ticks=True, no_yticks=True, **pkwargs)
    
    plt.subplot(2,4,3)
    plot_bs(bs_NbTe2, title='NbTe$_2$/Bi$_2$Te$_3$', empty_ticks=True, no_yticks=True, **pkwargs)
    
    plt.subplot(2,4,4)
    plot_bs(bs_TiTe2, title='TiTe$_2$/Bi$_2$Te$_3$', empty_ticks=True, no_yticks=True, **pkwargs)
    
    plt.subplot(2,4,5)
    plot_bs(bs_PtTe2_Bi2, title='PtTe$_2$/Bi$_2$/Bi$_2$Te$_3$', **pkwargs)
    
    plt.subplot(2,4,6)
    plot_bs(bs_PdTe2_Bi2, title='PdTe$_2$/Bi$_2$/Bi$_2$Te$_3$', no_yticks=True, **pkwargs)
    
    plt.subplot(2,4,7)
    plot_bs(bs_NbTe2_Bi2, title='NbTe$_2$/Bi$_2$/Bi$_2$Te$_3$', no_yticks=True, **pkwargs)
    
    plt.subplot(2,4,8)
    plot_bs(bs_TiTe2_Bi2, title='TiTe$_2$/Bi$_2$/Bi$_2$Te$_3$', no_yticks=True, **pkwargs)
    

    if suptitle is not None:
        plt.suptitle(suptitle)
    plt.tight_layout()
    
    if show:
        plt.show()

def get_data_all(calcs):
    ds_at_all= []
    for calc in calcs:
        if calc.is_finished_ok:
            denlm_at = []
            ds_at = []
            from tqdm import tqdm
            for iat in tqdm(range(len(find_parent_structure(calc).sites))):
                with calc.outputs.retrieved.open('den_lm_ns_%0.3i.npy'%(iat+1), 'rb') as _f:
                    denlm_at.append(np.load(_f))
                ds_at.append(np.sum([-np.imag(denlm_at[-1][32:, :32][:16, 16:][i,i]) for i in range(16)]))
            ds_at = np.array(ds_at)
            denlm_at = np.array(denlm_at)
            ds_at = ds_at[[2, 1, 0]+list(range(3, len(ds_at)))] # reorder left vacuum (wrong z order in structure and thus inputcard)
            ds_at_all.append(ds_at)
    return ds_at_all

def get_adens_data(calcs, iBi2_list, iBi2):
    ds_at_all = get_data_all(calcs)
    dmean = []
    for i, ds in enumerate(ds_at_all):
        if iBi2_list is not None:
            iBi2 = iBi2_list[i]
        if iBi2<1:
            dmean.append([np.mean(ds[3:6]), 0, np.mean(ds[6:6+5]), np.mean(ds[6+5:6+10]), np.mean(ds[6+10:6+15])])
        else:
            dmean.append([np.mean(ds[3:6]), np.mean(ds[6:8]), np.mean(ds[8:8+5]), np.mean(ds[8+5:8+10]), np.mean(ds[8+10:8+15])])
    dmean = np.array(dmean)
    return ds_at_all, dmean


def plot_adens_decay(calcs, iBi2 = 0, iBi2_list=None, scale = [1., 1., 1., 1.], all_layers=False, normalize=True, verbose=False, show=True, names=None, average_line=False):
    ds_at_all, dmean = get_adens_data(calcs, iBi2_list, iBi2)
    
    if verbose:
        print(ds_at_all, dmean)

    if all_layers:
        plt.figure()
        plt.title('Decay of anomalous density')
        for i, ds in enumerate(ds_at_all):
            if normalize:
                plt.plot(scale[i]*ds/np.sum(ds[:6]), 'o-')
            else:
                plt.plot(scale[i]*ds, 'o-')
        # plt.ylim(0, 1e-5)

        plt.ylabel('$\overline{\chi}$ (arb. units)')

        plt.twinx()
        for i, ds in enumerate(ds_at_all):
            if normalize:
                plt.plot(scale[i]*ds/np.sum(ds[:6]), ls='--')
            else:
                plt.plot(scale[i]*ds, ls='--')
        plt.yscale('log')
        # plt.ylim(1e-9)

        plt.axvline(2.5, color='grey', ls=':')
        plt.axvline(21.5, color='grey', ls=':')
        plt.ylabel('$\overline{\chi}$ (arb. units)')
        plt.xlabel('layer index')
        if show:
            plt.show()
    else:
        plt.figure()
        plt.title('Normalized decay of anomalous density, region-integrated')

        if names is None:
            names = ['PtTe$_2$', 'PdTe$_2$', 'NbTe$_2$', 'TiTe$_2$']
        for i, d in enumerate(dmean):
            if normalize:
                d /= d[0]
            d *= scale[i]
            plt.bar(np.arange(len(d))-0.15+i*0.1, d, width=0.2, label=names[i])
            if average_line:
                plt.axhline(d[-3:].sum(), color=f'C{i}', ls='--')
            print(d, d[-3:].sum())
        plt.yscale('log')
        plt.legend()
        plt.ylabel('$\overline{\chi}$ (arb. units)')
        plt.xticks(range(5), ['XTe$_2$', 'Bi$_2$', '(Bi$_2$Te$_3$)$^{(1)}$', '(Bi$_2$Te$_3$)$^{(2)}$', '(Bi$_2$Te$_3$)$^{(3)}$'])
        if show:
            plt.show()

BdG0_PtTe2 = orm.load_node('5d1c6b7c-8063-4c7c-8303-b019375441ac')
BdG0_PtTe2_nTI = orm.load_node('b819aa78-e0d8-4d34-9c51-e443bc935a6b')
BdG0_PdTe2 = orm.load_node('4625f596-713a-4bf0-91ac-135dfadfcdce')
BdG0_PdTe2_nTI = orm.load_node('687c84d4-b95d-4e47-911a-7a4e75439df1')
BdG0_NbTe2 = orm.load_node('bcd6f4bf-cd5e-44cb-a92f-7ab43950895b')
BdG0_NbTe2_nTI = orm.load_node('3d6ba41c-59e6-407d-b77b-d80184fbe956')
BdG0_TiTe2 = orm.load_node('119c6cf4-f405-44ba-a7ca-04bea0ac0a3f')
BdG0_TiTe2_nTI = orm.load_node('86145f11-305a-4641-84e9-fdc1fa13eba1')

BdG0_PtTe2_Bi2 = orm.load_node('1e2b9d7e-ee41-43e4-a898-8ac68b2c378f')
BdG0_PtTe2_Bi2_nTI = BdG0_PtTe2_Bi2 # same calculation
BdG0_PdTe2_Bi2 = orm.load_node('c2360b63-a9ea-4cc8-86e6-ea0a7d4e8682')
BdG0_PdTe2_Bi2_nTI = BdG0_PdTe2_Bi2 # same calculation
BdG0_NbTe2_Bi2 = orm.load_node('4ef2a45f-0c76-480e-8758-fb349760f0c1')
BdG0_NbTe2_Bi2_nTI = orm.load_node('1f4df7aa-189a-43b3-a2c2-9a2b11831882')
BdG0_TiTe2_Bi2 = orm.load_node('0995c730-564a-46cb-a17a-343918272b1c')
BdG0_TiTe2_Bi2_nTI = orm.load_node('d95c59ed-e239-4f58-bc7e-3e78b26fafe1')

def load_dos_data(dos_BdG0):
    from tqdm import tqdm
    from masci_tools.util.constants import RY_TO_EV

    calc = dos_BdG0.get_outgoing(node_class=orm.CalcJobNode).first().node
    ef = calc.inputs.parent_folder.get_incoming(node_class=orm.CalcJobNode).first().node.outputs.output_parameters['fermi_energy']
    ret = calc.outputs.retrieved
    dosnames = [i for i in ret.list_object_names() if 'dos.atom' in i and '_' not in i]
    # dosnames
    dall = []
    for iatom in tqdm(range(1, 25)):
        fname = 'dos.atom'+str(iatom)
        with ret.open(fname) as _f:
            tmp = np.loadtxt(_f)
        # print(tmp.sum(axis=0))
        # add hole contribution
        fname = 'dos.atom'+str(iatom)+'_hole'
        with ret.open(fname) as _f:
            tmp[:,1:] += np.loadtxt(_f)[:,1:]
        # print(tmp.sum(axis=0))
        dall.append(tmp)
    dall = np.array(dall)

    e = dall[0][:,0]
    e = (e[len(e)//2:]-ef)*RY_TO_EV
    # e /= Delta0

    dall_eh = []
    for iatom in tqdm(range(1, 25)):
        fname = 'dos.atom'+str(iatom)+'_eh'
        with ret.open(fname) as _f:
            tmp = np.loadtxt(_f)
        dall_eh.append(tmp)
    dall_eh = np.array(dall_eh)

    return e, dall, dall_eh

# load data nodes
dos_PtTe2_nTI = orm.load_node('402aa788-569a-4102-953d-051efb2210ff')
dos_PdTe2_nTI = orm.load_node('139f9031-28f2-4226-8895-cb000b28954c')
dos_NbTe2_Bi2_nTI = orm.load_node('2b3d73e8-23be-4cd3-ad29-15dacd90231e')
dos_TiTe2_Bi2_nTI = orm.load_node('837e0f21-7304-40e3-9615-d3dedcdeaf08')

def plot_dos_BdG(dos_BdG0, icolor=0, color=None, show_eh=False, show_ee=True, show_layers=False, norm=True, label=None, iBi2 = 1, BTonly=False, xfac=1., lw=3):
    
    e, dall, dall_eh = load_dos_data(dos_BdG0)
    e *= xfac

    if color is None:
        color=f'C{icolor}'
    if show_ee:
        dd = np.sum(dall[:, len(e):, 1], axis=0)/RY_TO_EV
        if norm:
            dd /= dd.max()
        if label is None:
            label = 'ee'
        # print(dd)
        plt.plot(e, dd, label=label, color=color, ls='-', lw=lw)
    if show_eh:
        dd = -np.sum(dall_eh[:, len(e):, 1], axis=0)/RY_TO_EV
        if norm:
            dd /= dd.max()
        if label is None:
            label = 'eh'
        plt.plot(e, dd, label=label, color=color, ls='--', lw=lw)
    # plt.legend()
    
    if show_layers:

        # norm = True
        second_y = False
        Delta0 = 1

        lw = 2

        plt.figure(figsize=(16,6))

        plt.subplot(1,2,1)

        dd = np.sum(dall[:, len(e):, 1], axis=0)/RY_TO_EV
        if norm: dd /= np.sum(dd)
        if not BTonly: plt.plot(e, dd, label='sum', color='C0', lw=3)
        dd = np.sum(dall[:6, len(e):, 1], axis=0)/RY_TO_EV
        if norm: dd /= np.sum(dd)
        if not BTonly: plt.plot(e, dd, label='XTe2', color='C1', lw=lw)
        if second_y:
            plt.twinx()
        if iBi2>0:
            dd = np.sum(dall[6:8, len(e):, 1], axis=0)/RY_TO_EV
            if norm: dd /= np.sum(dd)
            if not BTonly: plt.plot(e, dd, label='Bi2', color='C2', lw=lw)
        dd = np.sum(dall[6+2*iBi2:11+2*iBi2, len(e):, 1], axis=0)/RY_TO_EV
        if norm: dd /= np.sum(dd)
        plt.plot(e, dd, label='BT1', color='C3', lw=lw)
        dd = np.sum(dall[11+2*iBi2:16+2*iBi2, len(e):, 1], axis=0)/RY_TO_EV
        if norm: dd /= np.sum(dd)
        plt.plot(e, dd, label='BT2', color='C4', lw=lw)
        dd = np.sum(dall[16+2*iBi2:, len(e):, 1], axis=0)/RY_TO_EV
        if norm: dd /= np.sum(dd)
        plt.plot(e, dd, label='BT3', color='C5', lw=lw)

        plt.xlim(0, 0.005/Delta0)
        plt.xlabel('$(E-E_\mathrm{F})/\Delta_0$')

        if norm:
            plt.ylabel('normalized DOS (arb. u.)')
        else:
            plt.ylabel('DOS (1/eV)')

        plt.axvline(0.0024/Delta0, color='k', ls='--')

        plt.ylim(0)
        plt.legend()


        plt.subplot(1,2,2)

        dd = -np.sum(dall_eh[:, len(e):, 1], axis=0)/RY_TO_EV
        if norm: dd /= np.sum(dd)
        if not BTonly: plt.plot(e, dd, label='sum', color='C0', lw=3)
        dd = -np.sum(dall_eh[:6, len(e):, 1], axis=0)/RY_TO_EV
        if norm: dd /= np.sum(dd)
        if not BTonly: plt.plot(e, dd, label='XTe2', color='C1', lw=lw)
        if second_y:
            plt.twinx()
        if iBi2>0:
            dd = -np.sum(dall_eh[6:8, len(e):, 1], axis=0)/RY_TO_EV
            if norm: dd /= np.sum(dd)
            if not BTonly: plt.plot(e, dd, label='Bi2', color='C2', lw=lw)
        dd = -np.sum(dall_eh[6+2*iBi2:11+2*iBi2, len(e):, 1], axis=0)/RY_TO_EV
        if norm: dd /= np.sum(dd)
        plt.plot(e, dd, label='BT1', color='C3', lw=lw)
        dd = -np.sum(dall_eh[11+2*iBi2:16+2*iBi2, len(e):, 1], axis=0)/RY_TO_EV
        if norm: dd /= np.sum(dd)
        plt.plot(e, dd, label='BT2', color='C4', lw=lw)
        dd = -np.sum(dall_eh[16+2*iBi2:, len(e):, 1], axis=0)/RY_TO_EV
        if norm: dd /= np.sum(dd)
        plt.plot(e, dd, label='BT3', color='C5', lw=lw)

        plt.xlim(0, 0.005/Delta0)
        plt.xlabel('$(E-E_\mathrm{F})/\Delta_0$')

        if norm:
            plt.ylabel('normalized DOS (arb. u.)')
        else:
            plt.ylabel('DOS (1/eV)')

        plt.axvline(0.0024/Delta0, color='k', ls='--')
        plt.ylim(0)