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_2.nodes if impname1 in node.label][0]
    # imp2 in layer 2 or 6
    if add_pos[2]==2:
        single_imp_calc2 = [node for node in group_single_imps_2.nodes if impname2 in node.label][0]
    else:
        single_imp_calc2 = [node for node in group_single_imps_6.nodes if impname2 in node.label][0]


    group_label = f'new_calcs/Bi2Se3/il_2_{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 ['Se', 'Bi']:
            if symbol=='Bi' and not done_first:
                symbol = 'Sb'
                done_first = True
            struc_aux.append_atom(position=site.position, symbols=symbol)
    # print(struc_aux.sites)
    
    ps = struc_aux.get_pymatgen()
    dists = ([np.linalg.norm(i.coords - ps.sites[1].coords) for i in ps.get_all_neighbors(15)[1] if 'Bi' in i.species_string or 'Sb' in i.species_string])
    dists_34 = ([np.linalg.norm(i.coords - ps.sites[1].coords) for i in ps.get_all_neighbors(15)[1] if 'Bi' in i.species_string])
    dists_33 = ([np.linalg.norm(i.coords - ps.sites[1].coords) for i in ps.get_all_neighbors(15)[1] if 'Sb' in i.species_string])
    
    pos_34 = np.array([i.coords - ps.sites[1].coords for i in ps.get_all_neighbors(15)[1] if 'Bi' in i.species_string])
    pos_33 = np.array([i.coords - ps.sites[1].coords for i in ps.get_all_neighbors(15)[1] if 'Sb' in i.species_string])
    
    dists_unique = np.unique(np.round(dists, 3))
    num_3, num_4 = 0, -1 # counters to keep track of offset index (for consistent labeling later on)
    def get_r(idist, num_3, num_4):
        # global num_3, num_4
        try:
            R = pos_33[abs(np.linalg.norm(pos_33, axis=1)-dists_unique[idist])<1e-3][0]
            num_3 += 1
            return num_3, num_4, [np.linalg.norm(R), R, 2, num_3]
        except:
            R = pos_34[abs(np.linalg.norm(pos_34, axis=1)-dists_unique[idist])<1e-3][0]
            num_4 += 1
            return num_3, num_4, [np.round(np.linalg.norm(R),3), R, 6, num_4]
    
    # new: add all positions
    add_positions = {}
    for i in range(17):
        num_3, num_4, add_positions[i] =  get_r(i, num_3, num_4)
    
    # add_positions
    

    return add_positions

# get_add_positions(struc)

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_2 = load_or_create_group('single_imps:Bi2Se3_il_2')
group_single_imps_6 = load_or_create_group('single_imps:Bi2Se3_il_6')

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

if savefigs:
    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)
        # print(iadd, group_label, len(group_double_imps.nodes))
        if len(group_double_imps.nodes)==0:
            raise ValueError(f'no calculation in 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
        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_2.nodes]
moments = [n.outputs.last_calc_output_parameters['magnetism_group']['spin_moment_per_atom'][0][2] for n in group_single_imps_2.nodes]

single_imp_moments = {l: abs(moments[i]) for i,l in enumerate(impnames) if 'Pd-mag' not in l}

def get_dat_from_node(node, reload=False, verbose=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,24]]
        if verbose:
            print(np.array(node.outputs.workflow_info['magmoms']), m1, m2)
        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]
        R, il, num_neighbors = get_R_num_neighbor(iadd)

        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[2].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[2].coords, r=R, dr=0.02) if i.species_string in ['Bi', 'Sb']])

    return all_neighbors

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[1:8,1:8], 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-6', '1-2'))
        plt.xticks(np.arange(len(Zimp_all_3d4d[1:8]))+0.5, impname_all_3d4d[1:8])
        plt.yticks(np.arange(len(Zimp_all_3d4d[1:8]))+0.5, impname_all_3d4d[1:8])

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 len(dd)<17:
        return np.NaN, np.NaN, np.NaN, np.NaN
    
    
    if plot:
        
        plt.plot(dd[:,1], dd[:,-1], '--', color='C0')
        ddd = dd[dd[:,5]==2]
        plt.plot(ddd[:,1], ddd[:,-1], 'o', color='C0', label='1-1')
        ddd = dd[dd[:,5]==6]
        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]==2], Tcs[dd[:,5]==2], 'o', color='C1')
        plt.plot(dd[:,1][dd[:,5]==6], Tcs[dd[:,5]==6], '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

# np.savetxt('Tcs_Bi2Se3.txt', np.nan_to_num(Tcs))

df_BT = read_csv('Jij_table.csv', usecols=range(1,16))
data_all_BT = df_BT.to_numpy()
data_all_BT = np.array(data_all_BT[:,:-1], dtype=float)

isort_BT = df_BT['i_shell'] + 100*df_BT['Zimp1'] + 10000*df_BT['Zimp2']
isort_BS = df['i_shell'] + 100*df['Zimp1'] + 10000*df['Zimp2']
isort_both = set(isort_BS).intersection(set(isort_BS))

jij_compare = []
for iBS, isort in enumerate(isort_BS):
    iBT = np.where(abs(isort_BT-isort)<1e-1)[0]
    if len(iBT)>0:
        if np.linalg.norm(data_all_BT[iBT,[0,6,7]]-data_all[iBS,[0,6,7]])<1e-3:
            jij_compare.append(list(data_all_BT[iBT,[0,6,7]]) + [data_all_BT[iBT[0],-1]] + [data_all[iBS,-2]])
jij_compare = np.array(jij_compare, dtype=float)

Tcs_BT = np.loadtxt('Tcs_Bi2Te3.txt')[2:7,2:7]
Tcs_BS = np.loadtxt('Tcs_Bi2Se3.txt')[2:7,2:7]