Data analysis and plotting¶

def load_or_create_group(label):
    try:
        group = orm.load_group(label)
    except:
        group = orm.Group(label)
        group.store()
    return group

def get_builder_double_imp(impname1, impname2, add_pos):

    # set offset
    alat0 = get_alat_from_bravais(np.array(struc.cell))
    offset_imp2 = orm.Dict({'index': 50, 'r_offset': add_pos[1]/alat0})

    # imp 1 is always in layer 3
    single_imp_calc1 = [node for node in group_single_imps_3.nodes if impname1 in node.label][0]
    # imp2 in layer 3 or 4
    if add_pos[2]==3:
        single_imp_calc2 = [node for node in group_single_imps_3.nodes if impname2 in node.label][0]
    else:
        single_imp_calc2 = [node for node in group_single_imps_4.nodes if impname2 in node.label][0]


    group_label = f'new_calcs/Bi2Te3/il_3_{add_pos[2]}_Off_{add_pos[3]}'
    calc_label = f'{impname1}:{impname2}:' + group_label.split('/')[-1]
    
    
    from aiida_kkr.workflows import combine_imps_wc

    # set combine_imps inputs
    builder = combine_imps_wc.get_builder()
    builder.metadata.label = calc_label
    
    return calc_label, group_label, builder

def relabel_kpts(bs):
    """
    Relabel high-symmetry points to use convention from https://www.sciencedirect.com/science/article/pii/S0927025610002697?via%3Dihub
    """
    k = bs.called[2].inputs.kpoints
    klbls = []
    k_rename = {
        'GAMMA': r'$\Gamma$',
        'T': 'Z',
        'H_2': 'B',
        'H_0': 'B1',
        'L': 'L',
        'S_0': 'P2',
        'S_2': 'Q',
        'F': 'F',
    }
    for i, (ik, klbl) in enumerate(k.labels):
        if klbl in k_rename:
            klbl = k_rename[klbl]
        # print(i, ik, (k.labels[i+1][0]-ik)==1)
        if (i==len(k.labels)-1 or i==0) or (abs(ik-k.labels[i+1][0])>1 and abs(ik-k.labels[i-1][0])>1):
            klbls.append([ik, klbl])
        elif k.labels[i+1][0]-ik==1:
            klbl2 = k.labels[i+1][1]
            if klbl2 in k_rename:
                klbl2 = k_rename[klbl2]
            klbl += '|'+klbl2
            klbls.append([ik+0.5, klbl])
    plt.xticks([i[0] for i in klbls], [i[1] for i in klbls])

def get_add_positions(struc):
    """Get all connecting vectors and distances for impurity dimer configurations"""
    done_first = False
    struc_aux = orm.StructureData(cell=struc.cell)
    for site in struc.sites:
        symbol = site.kind_name
        if symbol in ['Te', 'Bi']:
            if symbol=='Bi' and not done_first:
                symbol = 'Sb'
                done_first = True
            struc_aux.append_atom(position=site.position, symbols=symbol)

    ps = struc_aux.get_pymatgen()
    dists = ([np.linalg.norm(i.coords - ps.sites[4].coords) for i in ps.get_all_neighbors(15)[4] if 'Bi' in i.species_string or 'Sb' in i.species_string])
    dists_34 = ([np.linalg.norm(i.coords - ps.sites[4].coords) for i in ps.get_all_neighbors(15)[4] if 'Sb' in i.species_string])
    dists_33 = ([np.linalg.norm(i.coords - ps.sites[4].coords) for i in ps.get_all_neighbors(15)[4] if 'Bi' in i.species_string])

    pos_34 = -np.array([i.coords - ps.sites[4].coords for i in ps.get_all_neighbors(15)[4] if 'Sb' in i.species_string])
    pos_33 = np.array([i.coords - ps.sites[4].coords for i in ps.get_all_neighbors(15)[4] if 'Bi' in i.species_string])
    
    num_3, num_4 = 0, -1 # counters to keep track of offset index (for consistent labeling later on)
    dists_unique = np.unique(np.round(dists, 3))
    def get_r(idist, num_3, num_4):
        try:
            pp = pos_33[abs(np.linalg.norm(pos_33, axis=1) - dists_unique[idist]) < 1e-3]
            if len(pp)==0:
                msg = f'Error in getting pos_33: {idist}, {dists_unique[idist]}, {abs(np.linalg.norm(pos_34, axis=1) - dists_unique[idist]).min()}'
                raise ValueError(msg)
            R = pp[0]
            num_3 += 1
            return num_3, num_4, [np.linalg.norm(R), R, 3, num_3]
        except:
            pp = pos_34[abs(np.linalg.norm(pos_34, axis=1) - dists_unique[idist]) < 1e-3]
            if len(pp)==0:
                msg = f'Error in getting pos_34: {idist}, {dists_unique[idist]}, {abs(np.linalg.norm(pos_34, axis=1) - dists_unique[idist]).min()}'
                raise ValueError(msg)
            R = pp[0]
            num_4 += 1
            return num_3, num_4, [np.round(np.linalg.norm(R),3), R, 4, num_4]

    add_positions = {}
    for i in range(17):
        num_3, num_4, add_positions[i] =  get_r(i, num_3, num_4)

    return add_positions

def get_imp_outputs(group, plot=False):
    """Extract output of single impurity calculation and plot data"""
    
    print('label           spin mom  orb. mom   rms of calc  charge neutrality of imp.')
    print(75*'-')
    out, names = [], []
    for impname in impname_all_3d4d:
        nodes = [node for node in group.nodes if impname in node.label]
        if len(nodes)==0:
            print(impname, 'not found in group')
        else:
            node = nodes[0]
            if node.is_finished_ok:
                o = node.outputs.last_calc_output_parameters.get_dict()
                o = node.outputs.last_calc_output_parameters.get_dict()
                rms =o['convergence_group']['rms']
                ms = o['magnetism_group']['spin_moment_per_atom'][0][2]
                mo = o['magnetism_group']['orbital_moment_per_atom'][0][2]
                if ms<0:
                    # flip spin and orbital moment to point in +z if necessary
                    mo *= -1
                    ms *= -1
                charge_neutr = o['total_charge_per_atom'][0] - node.inputs.impurity_info['Zimp']
                print(f'{node.label:15} {np.round(ms,3):7} {np.round(mo,3):7} {rms:15}   {charge_neutr}')
                
                out.append([ms, mo])
            else:
                print(node.label, 'not done yet. pk:', node.pk)
                out.append([np.NaN, np.NaN])
            names.append(impname)
        
                
    if plot and len(out)>0:
        plt.figure(figsize=(8,4))
        plt.subplot(1,2,1)
        out = np.array(out)
        plt.axhline(0, ls=':', color='C0')
        plt.plot(out[:,0], 'o--')
        m = np.where(abs(np.nan_to_num(out)[:,0])<1e-3)
        plt.plot(np.arange(len(out))[m], out[m][:,0], 'o', color='w', ms=3)
        plt.ylabel(r'spin moment ($\mu_B$)', color='C0')
        plt.xticks(np.arange(len(out)), [ '\n'*(i%2) + n for i, n in enumerate(names)])
        plt.axvline(9.5, color='grey', ls=':')


        plt.subplot(1,2,2)
        plt.axhline(0, ls=':', color='C1')
        plt.plot(out[:,1], 'o--', color='C1')
        plt.plot(np.arange(len(out))[m], out[m][:,0], 'o', color='w', ms=3)
        plt.ylabel(r'orbital moment ($\mu_B$)', color='C1')
        plt.xticks(np.arange(len(out)), [ '\n'*(i%2) + n for i, n in enumerate(names)])
        plt.axvline(9.5, color='grey', ls=':')
        
        plt.tight_layout()

# save figures to files?
savefigs = False

# mappings of Zimp to element names

Zimp_all_3d4d = list(range(21, 31)) + list(range(39, 49))
impname_all_3d4d = 'Sc Ti V Cr Mn Fe Co Ni Cu Zn Y Zr Nb Mo Tc Ru Rh Pd Ag Cd'.split()

impname_from_Zimp = {Z: impname_all_3d4d[i] for i, Z in enumerate(Zimp_all_3d4d)}

# get impurity configurations (distance, connecting vectors etc.)  
add_positions = get_add_positions(struc)

# load groups for single imp. calculations
group_single_imps_3 = load_or_create_group('single_imps:Bi2Te3_il_3')
group_single_imps_4 = load_or_create_group('single_imps:Bi2Te3_il_4')

%%capture --no-stdout --no-display
from aiida_kkr.workflows import kkr_imp_sub_wc

for iadd in range(17):
    calc_label, group_label, builder = get_builder_double_imp('Mn', 'Mn', add_positions[iadd])
    group_double_imps = load_or_create_group(group_label)
    calc = [n for n in group_double_imps.nodes if n.label==calc_label][0]
    calc = calc.get_outgoing(node_class=kkr_imp_sub_wc).first().node
    if savefigs:
        plot_kkr(calc, strucplot=True, silent=True, static=True)

# extract single impurity moments, used later again
impnames = [n.label.split(':')[0] for n in group_single_imps_3.nodes]
moments = [n.outputs.last_calc_output_parameters['magnetism_group']['spin_moment_per_atom'][0][2] for n in group_single_imps_3.nodes]

single_imp_moments = {l: abs(moments[i]) for i,l in enumerate(impnames)}

# !verdi archive create -G 'Bi2Te3/single_imps_DOS' --compress 9 -f 'export_DOS.aiida'
# !du -sch export_DOS.aiida

def get_dat_from_node(node, reload=False):
    """Get data from node (neighbor index, radius, Zimp, ..., Jij)"""
    if 'Jij_data' in node.extras and not reload:
        # load from extra
        dat = node.extras['Jij_data']
    else:
        # load from node if not found as extra already
        imp1, imp2, off = node.label.split(':')
        # data.append([add_positions[iadd][0], imp1, imp2])
        Zimp1, Zimp2 = node.outputs.workflow_info['imps_info_in_exact_cluster']['Zimps']
        m1, m2 = np.array(node.outputs.workflow_info['magmoms'])[[0,20]]
        J = node.outputs.JijData.get_array('JijData')[0,3]

        m01 = single_imp_moments[impname_from_Zimp[Zimp1]]
        m02 = single_imp_moments[impname_from_Zimp[Zimp2]]
        
        uuid_JijData = node.outputs.JijData.uuid

        r = add_positions[iadd][1]
        dat = [iadd, add_positions[iadd][0], r[0], r[1], r[2], add_positions[iadd][2], Zimp1, Zimp2, num_neighbors, m1, m2, m01, m02, J, uuid_JijData]
        node.set_extra('Jij_data', dat)
    return dat

def get_R_num_neighbor(iadd):
    """Extract shell Radius, ilayer index and number of neighbors per shell"""
    R, _, il, _ = add_positions[iadd]
    num_neighbors = len([i for i in ps.get_neighbors_in_shell(origin=ps.sites[3].coords, r=R, dr=0.02) if i.species_string in ['Bi', 'Sb']])

    return R, il, num_neighbors

def get_all_equivalent_R(iadd):
    """Extract all equivalent positions for a shell"""
    R, _, il, _ = add_positions[iadd]
    all_neighbors =np.array( [i.coords for i in ps.get_neighbors_in_shell(origin=ps.sites[3].coords, r=R, dr=0.02) if i.species_string in ['Bi', 'Sb']])

    return all_neighbors

def get_d_part(Zimp):    
    dd = data_all[data_all[:,6]==Zimp]
    return dd

def get_m_change(Zimp, plot=True):
    dd = get_d_part(Zimp)
    m0 = dd[0,-4]
    if plot:
        plt.scatter(x=dd[:,1], y=dd[:,-6], c=dd[:,7], cmap='Spectral', )
        plt.axhline(m0)
        plt.twinx()
        plt.scatter(x=dd[:,1], y=100*(dd[:,-6]-m0)/m0, c=dd[:,9], cmap='Spectral', )

    return m0, dd

def get_data_delta_m1(iadd = 0):
    arr = np.zeros((20, 20))
    # arr[:,:] = np.NaN
    arr2 = arr.copy()

    for i, Zimp in enumerate(Zimp_all_3d4d):
        m0, dat = get_m_change(Zimp, plot=False)
        dd = dat[dat[:,1]==add_positions[iadd][0]]
        for j, Zimp2 in enumerate(Zimp_all_3d4d):
            ddd = dd[dd[:,7]==Zimp2]
            if len(ddd)>0:
                arr[i, j] = ddd[0, -6] - ddd[0, -4]
                arr2[i, j] = ddd[0, -5] - ddd[0, -3]
                # arr[i, j] = abs(ddd[0, -6]) - abs(ddd[0, -4])
                # arr2[i, j] = abs(ddd[0, -5]) - abs(ddd[0, -3])
                
    arr22 = arr2.copy()
    for i in range(len(arr)):
        arr22[i,i] = 0
    dat = np.nan_to_num(arr) + np.nan_to_num(arr22).transpose()
    
    return dat, arr, arr2

def plot_delta_m1(iadd=0, figsize=(16,6), **kwargs):
    dat, arr, arr2  = get_data_delta_m1(iadd)
    
    moms = np.zeros((20,20))
    for i, Zimp in enumerate(Zimp_all_3d4d):
        m0, _ = get_m_change(Zimp, plot=False)
        moms[i,:] = m0
        
    plt.figure(figsize=figsize)

    plt.subplot(1,2,1)

    g = sns.heatmap(dat, cmap='seismic', center=0, square=True, vmin=-0.5, vmax=0.3, **kwargs)
    plt.xlabel('impurity 2', fontsize='large')
    plt.ylabel('impurity 1', fontsize='large')
    
    print(np.max(abs(dat), axis=1))

    plt.xticks(np.arange(len(Zimp_all_3d4d))+0.5, impname_all_3d4d)
    plt.yticks(np.arange(len(Zimp_all_3d4d))+0.5, impname_all_3d4d)

    # g.set_facecolor('lightgrey')


    plt.title(r'$\Delta_{m_1}$ ($\mu_B$)', fontsize='large')

    [f(i, color='grey') for f in [plt.axhline] for i in [1,8, 12, 16]]
    plt.annotate(xy=(0.39, 0.91), text=r'$\iota$', xycoords='figure fraction' )
    plt.annotate(xy=(0.39, 0.76), text=r'$\mu$', xycoords='figure fraction' )
    plt.annotate(xy=(0.39, 0.51), text=r'$\iota$', xycoords='figure fraction' )
    plt.annotate(xy=(0.39, 0.36), text=r'$\mu$', xycoords='figure fraction' )
    plt.annotate(xy=(0.39, 0.18), text=r'$\iota$', xycoords='figure fraction' )

    plt.subplot(1,2,2)
    arr22 = arr2.copy()
    for i in range(len(arr)):
        arr22[i,i] = 0
    dat = np.nan_to_num(arr) + np.nan_to_num(arr22).transpose()
    dat = dat / moms * 100
    dat[abs(moms)<1e-1] = np.NaN
    
    print('rel')
    print(np.max(np.nan_to_num(abs(dat)), axis=1))


    g = sns.heatmap(dat, cmap='seismic', center=0, vmin=-50, vmax=20, square=True, **kwargs)
    plt.xlabel('impurity 2', fontsize='large')
    plt.ylabel('impurity 1', fontsize='large')

    plt.xticks(np.arange(len(Zimp_all_3d4d))+0.5, impname_all_3d4d)
    plt.yticks(np.arange(len(Zimp_all_3d4d))+0.5, impname_all_3d4d)


    g.set_facecolor('lightgrey')
    plt.title(r'$\Delta_{m_1}/m_1$ (%)', fontsize='x-large')

    plt.tight_layout()
    plt.subplots_adjust(wspace=0.1)

    plt.annotate(xy=(0.02, 0.95), text="(a)", xycoords='figure fraction', fontsize=18)
    plt.annotate(xy=(0.52, 0.95), text="(b)", xycoords='figure fraction', fontsize=18)

def reshape_data_to_arr(data, plot_J=True):
    data = np.array(data)

    Zimp_all_3d4d = list(range(21, 31)) + list(range(39, 49))
    impname_all_3d4d = 'Sc Ti V Cr Mn Fe Co Ni Cu Zn Y Zr Nb Mo Tc Ru Rh Pd Ag Cd'.split()

    arr = np.zeros((len(Zimp_all_3d4d), len(Zimp_all_3d4d)))
    arr[:,:] = np.NaN

    for i1, Z1 in enumerate(Zimp_all_3d4d):
        for i2, Z2 in enumerate(Zimp_all_3d4d):
            d = data[data[:,6]==Z1]
            # print(len(d))
            if len(d)>0:
                dd = d[d[:,7]==Z2]
                if len(dd)>0:
                    if plot_J:
                        arr[i1, i2] = dd[0][-2]
                    else:
                        arr[i1, i2] = dd[0][-4] + dd[0][-3]
                    if i1 != i2:
                        arr[i2, i1] = arr[i1, i2]
    return arr


def get_and_plot_data(iadd = 0, plot_J=True, reload=False, vlim=None, hide_threshold=None, **kwargs):
    data = data_all[data_all[:,0]==iadd]
    

    if len(data)>0:
        arr = reshape_data_to_arr(data)
        if hide_threshold is not None:
            arr[abs(arr)<hide_threshold] = 0 #np.NaN

        import seaborn as sns

        # plt.figure(figsize=(18,14))
        ivm = 1
        cmap='seismic'
        if vlim is None:
            vlim = abs(np.nan_to_num(arr)).max()
        if not plot_J:
            ivm = 0
            cmap='viridis'
        p = sns.heatmap(arr, cmap=cmap,  vmin=ivm * -vlim, vmax=vlim, **kwargs)
        R, _, il1, ioff = add_positions[iadd]
        plt.title(f'3-{il1}:{ioff} (R={R:.3f} Ang.)'.replace('3-3', '1-1').replace('3-4', '1-2'))
        plt.xticks(np.arange(len(Zimp_all_3d4d))+0.5, impname_all_3d4d)
        plt.yticks(np.arange(len(Zimp_all_3d4d))+0.5, impname_all_3d4d)

kB = 8.6173324e-5 # eV / K
# prefac = 1e-3 * 2/3 / kB
prefac = 1e-3 * 1/3 / kB # 1/3 instead of 2/3 due to definition of Heisenberg Hamiltonian

def get_Tc(imp1, imp2, plot=False, mark_dists=False):
    d = data_all[data_all[:,6]==imp1]
    dd = d[d[:,7]==imp2]
    
    
    if plot:
        
        plt.plot(dd[:,1], dd[:,-1], '--', color='C0')
        ddd = dd[dd[:,5]==3]
        plt.plot(ddd[:,1], ddd[:,-1], 'o', color='C0', label='1-1')
        ddd = dd[dd[:,5]==4]
        plt.plot(ddd[:,1], ddd[:,-1], 's', color='C0', label='1-2')
        plt.legend(loc=5)
        plt.ylabel(r'$N\cdot J_{ij}$ (meV)', color='C0')
        plt.axhline(0, ls=':', color='C0')
        plt.xlabel(r'$R_{ij}$ ($\AA$)')
        plt.twinx()
        Tcs = np.array([np.sum(dd[:1+i,-2]*dd[:1+i,5]) * prefac for i in range(len(dd))])
        plt.plot(dd[:,1], Tcs, '-', color='C1')
        plt.plot(dd[:,1][dd[:,5]==3], Tcs[dd[:,5]==3], 'o', color='C1')
        plt.plot(dd[:,1][dd[:,5]==4], Tcs[dd[:,5]==4], 's', color='C1')
        plt.ylabel('$T_c$ (K/c)', color='C1')
        # plt.axhline(0, ls=':', color='C1')
        plt.xlim(4, 12.5)
        plt.tight_layout()
        
        if mark_dists:            
            plt.axvline(np.min(np.round(dists_33, 3)), ls=':', color='r', label='33')
            [plt.axvline(i, ls=':', color='r') for i in np.unique(np.round(dists_33, 3))[1:]]
            plt.axvline(np.min(np.round(dists_34, 3)), ls=':', color='m', label='34')
            [plt.axvline(i, ls=':', color='m') for i in np.unique(np.round(dists_34, 3))[1:]]

            
    Tc = np.sum(dd[:,-2]*dd[:,8]) * prefac
    Tc1 = np.sum(dd[1:,-2]*dd[1:,8]) * prefac
    ddd_z, ddd_xy = [], []
    for i, iadd in enumerate(dd[:,0]):
        z = add_positions[iadd][1][2]
        if abs(z)>5:
            ddd_z.append(dd[i])
        else:
            ddd_xy.append(dd[i])
    ddd_z, ddd_xy = np.array(ddd_z), np.array(ddd_xy)
    Tc_xy = np.sum(ddd_xy[:,-2]*ddd_xy[:,8]) * prefac
    Tc_z = np.sum(ddd_z[:,-2]*ddd_z[:,8]) * prefac
    return Tc, Tc1, Tc_xy, Tc_z

Zimp_all_3d4d = list(range(21, 31)) + list(range(39, 49))
impname_all_3d4d = 'Sc Ti V Cr Mn Fe Co Ni Cu Zn Y Zr Nb Mo Tc Ru Rh Pd Ag Cd'.split()

Tcs = np.zeros((len(impname_all_3d4d), len(impname_all_3d4d)))
Tcs[:,:] = np.NaN
Tcs1 = Tcs.copy()
Tcs_xy = Tcs.copy()
Tcs_z = Tcs.copy()
for imp1 in range(len(impname_all_3d4d)):
    for imp2 in range(imp1, len(impname_all_3d4d)):
        Tcs[imp1, imp2], Tcs1[imp1, imp2], Tcs_xy[imp1, imp2], Tcs_z[imp1, imp2] = get_Tc(Zimp_all_3d4d[imp1], Zimp_all_3d4d[imp2])
for imp1 in range(len(impname_all_3d4d)):
    for imp2 in range(imp1, len(impname_all_3d4d)):
        if imp1 != imp2:
            Tcs[imp1, imp2] = (Tcs[imp1, imp2] + Tcs[imp1, imp1] + Tcs[imp2, imp2])/3
            Tcs[imp2, imp1] = Tcs[imp1, imp2]
            # same for Tc without the first neighbor
            Tcs1[imp1, imp2] = (Tcs1[imp1, imp2] + Tcs1[imp1, imp1] + Tcs1[imp2, imp2])/3
            Tcs1[imp2, imp1] = Tcs1[imp1, imp2]
            Tcs_xy[imp1, imp2] = (Tcs_xy[imp1, imp2] + Tcs_xy[imp1, imp1] + Tcs_xy[imp2, imp2])/3
            Tcs_xy[imp2, imp1] = Tcs_xy[imp1, imp2]
            Tcs_z[imp1, imp2] = (Tcs_z[imp1, imp2] + Tcs_z[imp1, imp1] + Tcs_z[imp2, imp2])/3
            Tcs_z[imp2, imp1] = Tcs_z[imp1, imp2]
            

m = np.array(list(single_imp_moments.values()))<1e-3
Tcs[m, :] = np.NaN
Tcs[:, m] = np.NaN
Tcs1[m, :] = np.NaN
Tcs1[:, m] = np.NaN
Tcs_xy[m, :] = np.NaN
Tcs_xy[:, m] = np.NaN
Tcs_z[m, :] = np.NaN
Tcs_z[:, m] = np.NaN