Constant energy contours and QPI model¶

import numpy as np
import matplotlib.pyplot as plt
import os
from tqdm import tqdm

import numpy as np
from multiprocessing import Pool
from tqdm import tqdm

parallel = False


def get_d_with_symm(E0, dE):

    # filter energy cut
    dist = abs(d[:,2]-E0)
    dd = d[dist<dE]

    # use symmetries
    d_with_symm = list(dd.copy())
    tmp = dd.copy(); tmp[:,0] *= -1
    d_with_symm += list(tmp)
    tmp = dd.copy(); tmp[:,1] *= -1
    d_with_symm += list(tmp)
    tmp = dd.copy(); tmp[:,0] *= -1; tmp[:,1] *= -1
    d_with_symm += list(tmp)
    d_with_symm = np.array(d_with_symm)

    return d_with_symm

def calculate_jdos_for_pair(A1, A2, spin=False):
    q = A2[:2] - A1[:, :2]
    
    if not spin:
        autocorr = A1[:, -1] * A2[-1]
    else:
        s1 = np.cross(A1[:, :2], [0, 0, 1])
        s2 = np.cross([A2[0], A2[1], 0], [0, 0, 1])
        
         # Normalize s1 and s2
        s1_norm = np.linalg.norm(s1, axis=1, keepdims=True)
        s2_norm = np.linalg.norm(s2)
        
        s1_normalized = s1 / s1_norm
        s2_normalized = s2 / s2_norm
        
        # |T^2| ~ 1 + cos(s1, s2)
        autocorr = A1[:, -1] * A2[-1] * (1 + np.sum(s1_normalized * s2_normalized, axis=1))
    
    result = np.column_stack((q[:, 0], q[:, 1], autocorr))
    return result


def calc_jdos_parallel(d_with_symm, num_processes, spin=False):
    jdos_results = []
    for A2 in tqdm(d_with_symm):
        current_A2_results = calculate_jdos_for_pair(d_with_symm, A2, spin=spin)
        jdos_results.extend(current_A2_results)

    jdos_array = np.array(jdos_results)
    return jdos_array

if not parallel:
    def calc_jdos_batch(d_with_symm, ibatch, Nbatch, spin=False):
        jdos_results = []
        L = len(d_with_symm)//Nbatch
        for A2 in tqdm(d_with_symm[ibatch*L:(ibatch+1)*L]):
            current_A2_results = calculate_jdos_for_pair(d_with_symm, A2, spin=spin)
            jdos_results.extend(current_A2_results)

        jdos_array = np.array(jdos_results)
        return jdos_array
else:
    from concurrent.futures import ThreadPoolExecutor

    def calc_jdos_batch(d_with_symm, ibatch, Nbatch, spin=False, num_threads=4):
        jdos_results = []
        L = len(d_with_symm) // Nbatch

        with ThreadPoolExecutor(max_workers=num_threads) as executor:
            futures = []

            for A2 in tqdm(d_with_symm[ibatch*L:(ibatch+1)*L]):
                future = executor.submit(calculate_jdos_for_pair, d_with_symm, A2, spin=spin)
                futures.append(future)

            for future in tqdm(futures):
                current_A2_results = future.result()
                jdos_results.extend(current_A2_results)

        jdos_array = np.array(jdos_results)
        return jdos_array

def get_jdos_hexbin_batch(d_with_symm, Nbatch, grid_size=100, extent=(-0.6, 0.6, -0.6, 0.6), spin=False):
    offsets_all = []
    arr_all = []

    for ibatch in range(Nbatch):
        JDOS = calc_jdos_batch(d_with_symm, ibatch, Nbatch, spin=spin)

        poly_batch = plt.hexbin(JDOS[:,0], JDOS[:,1], C=JDOS[:,2], reduce_C_function=np.sum,
                   gridsize=(int(np.sqrt(3)*grid_size), grid_size), extent=extent,
                   cmap='binary'
                  )
        offsets_batch = poly_batch.get_offsets()
        arr_batch = poly_batch.get_array()

        offsets_all.extend(offsets_batch)
        arr_all.extend(arr_batch)

    offsets_all = np.array(offsets_all)
    arr_all = np.array(arr_all)

    print(offsets_all.shape)


    poly = plt.hexbin(offsets_all[:,0], offsets_all[:,1], C=arr_all,
                            reduce_C_function=np.sum,
               gridsize=(int(np.sqrt(3)*grid_size), grid_size), extent=extent,
               cmap='binary'
              )
    
    return poly

kx_max = dKGM[:,0].max()

def fold_back_x(ddd):
    for i, tmp in enumerate(ddd):
        if tmp[0]>2*kx_max/3:
            ddd[i, 0] *= -1
            ddd[i, 0] += 4*kx_max/3
        if tmp[0]<-2*kx_max/3:
            ddd[i, 0] *= -1
            ddd[i, 0] -= 4*kx_max/3
        if tmp[0]>kx_max/3:
            ddd[i, 0] *= -1
            ddd[i, 0] += 2*kx_max/3
        if tmp[0]<-kx_max/3:
            ddd[i, 0] *= -1
            ddd[i, 0] -= 2*kx_max/3
    return ddd


def plot_CEC(E0 = -0.5, dE = 2e-2,
             verbose=False, nofig=False, noshow=False,
             show_dispersion=True, show_CEC=True,  s0=1, cmap='binary',
             interpol=False, xrange=(-0.5, 0.5), yrange=(-0.6, 0.6), Nx=100, Ny=100, method='linear',
             export_name=None, save_fig=None,
             joint_density_of_states=False, grid_size=100, extent=(-0.6, 0.6, -0.6, 0.6), recalc_JDOS=False,
             num_processes=12,
             fft_corr=False,
             Nbatch=10,
             fold_back = False,
             log_intensity=False, clim=None,
             spin=False, subplot = (1,2,0)
            ):
    
    if E0 is not None and dE is not None:
        fname_jdos = f'save_jdos_E0={E0:.3f}_dE={dE:.3f}.npz'
        if fold_back:
            fname_jdos = fname_jdos.replace('jdos', 'jdos_backfolded')
        if spin:
            fname_jdos = fname_jdos.replace('.npz', '_spin.npz')
            
    if not nofig:
        if show_dispersion and show_CEC:
            plt.figure(figsize=(16,9))
        elif show_dispersion or show_CEC:
            plt.figure(figsize=(4,3))

    if show_dispersion and show_CEC:
        plt.subplot(subplot[0], subplot[1], subplot[2]+1)
    if show_dispersion:
        ddd = dKGM.copy()
        if fold_back:
            ddd = fold_back_x(ddd)
        plt.scatter(ddd[:,0]-ddd[:,1], ddd[:,2], c=ddd[:,3], s=ddd[:,3], cmap=cmap, rasterized=True)
        plt.ylabel(r'$E-E_F$ (eV)')
        plt.xlabel(r'$k$')
        plt.ylim(-2,1)
        
        if E0 is not None and dE is not None:
            ymin, ymax = E0 - dE, E0 + dE
            plt.fill_between(x=[-dKGM[:,1].max(), dKGM[:,0].max()],
                             y1=[ymin, ymin], y2=[ymax, ymax], color='grey', alpha=0.3)
        
    if (show_CEC or
        (joint_density_of_states and (fname_jdos not in os.listdir('data/CEC_data/') or recalc_JDOS)) or
        (export_name is not None)
       ):

        if verbose:
            print('get CEC data')
            
        d_with_symm = get_d_with_symm(E0, dE)
        if fold_back:
            d_with_symm = fold_back_x(d_with_symm)
    
    if export_name is not None:
        if verbose:
            print('export CEC data')
        from pandas import DataFrame
        df = DataFrame.from_dict({'# kx': d_with_symm[:,0], 'ky': d_with_symm[:,1], 'weight': d_with_symm[:,-1]})
        with open(export_name, 'w') as _f:
            _f.writelines(df.to_csv(header=True, index=False))

    if show_dispersion and show_CEC:
        plt.subplot(subplot[0], subplot[1], subplot[2]+2)
    
    if show_CEC:
        if verbose:
            print('plot CEC')
        if not interpol:
            plt.scatter(d_with_symm[:,0], d_with_symm[:,1], c=d_with_symm[:,-1], s=s0*d_with_symm[:,-1], cmap=cmap)

        else:
            from scipy.interpolate import griddata
            grid_x, grid_y = np.mgrid[-0.5:0.5:100j, -0.6:0.6:200j]
            interpol = griddata(points=d_with_symm[:,0:2], values=d_with_symm[:,-1], xi=(grid_x, grid_y), method=method)
            plt.pcolormesh(grid_x, grid_y, np.nan_to_num(interpol), cmap=cmap)
            
        if spin:
            s1 = np.cross(d_with_symm[:, :2], [0, 0, 1])
            s1_norm = np.linalg.norm(s1, axis=1, keepdims=True)
            s1_normalized = s1 / s1_norm
            plt.quiver(d_with_symm[:,0], d_with_symm[:,1], s1_normalized[:,0], s1_normalized[:,1])

        plt.axis('equal')
        plt.xlabel(r'$k_x$')
        plt.ylabel(r'$k_y$')
        plt.title(rf'$E_0={E0:.2f}\,$eV, $\Delta E={dE:.2f}\,$eV')
    if (show_dispersion or show_CEC) and save_fig is not None:
        plt.savefig(save_fig)
    if not noshow and (show_dispersion or show_CEC):
        plt.show()
    
    if joint_density_of_states:
        if fname_jdos not in os.listdir('data/CEC_data/') or recalc_JDOS:
            print('load JDOS and save to ', fname_jdos)

            poly = get_jdos_hexbin_batch(d_with_symm, Nbatch, spin=spin, grid_size=grid_size, extent=extent)

            offsets = poly.get_offsets()
            arr = poly.get_array()
        
            np.savez('data/CEC_data/'+fname_jdos, x=offsets, y=arr)
        else:
            print('load JDOS from ', fname_jdos)
            tmp = np.load('data/CEC_data/'+fname_jdos)
            offsets = tmp.get('x')
            arr = tmp.get('y')

        if verbose:
            print('plot JDOS')
        if not nofig:
            plt.figure()
        
        # plot the hexbins
        if log_intensity:
            arr = np.log10(arr)
        plt.hexbin(offsets[:, 0], offsets[:, 1], C=arr,
                   gridsize=(int(np.sqrt(3)*grid_size), grid_size), extent=extent,
                   cmap=cmap
                  )
            
        plt.axis('equal')
        plt.xlabel(r'$q_x$')
        plt.ylabel(r'$q_y$')
        plt.title(rf'$E_0={E0:.2f}\,$eV, $\Delta E={dE:.2f}\,$eV')
        if clim is not None:
            try:
                plt.clim(clim[0], clim[1])
            except:
                plt.clim(clim)
        if not noshow:
            plt.show()

from scipy.interpolate import griddata
from tqdm import tqdm

def get_d_with_zeros(dd, dmin0 = 2e-2):
    d0 = []
    x = np.linspace(xmin, xmax, 64+1)
    y = np.linspace(ymin, ymax, 300)
    for ix in x:
        for iy in y:
            d0.append([ix, 0, iy, 1e-7])
    d0 = np.array(d0)
    print(d0.shape)

    d00 = []
    for temp in tqdm(d0):
        diffmin = min(np.sqrt((temp[0]-dd[:,0])**2+(temp[2]-dd[:,2])**2))
        if diffmin>dmin0:
            d00.append(temp)

    ddd = np.array(list(dd) + d00)
    d00 = np.array(d00)

    return ddd, d00


def get_hexbin_data(ddd):
    extent = (xmin, xmax, ymin, ymax)
    grid_size = 300
    poly_batch = plt.hexbin(ddd[:,0], ddd[:,2], C=ddd[:,-1], reduce_C_function=np.sum,
                            gridsize=(int(np.sqrt(3)*grid_size), grid_size), extent=extent,
                            cmap='binary')
    offsets_batch = poly_batch.get_offsets()
    arr_batch = poly_batch.get_array()
    return offsets_batch, arr_batch

def interpolate(offsets_batch, arr_batch):
    x = offsets_batch[:,0]
    y = offsets_batch[:,1]
    scalar_values = arr_batch

    # interpolation grid
    Nx, Ny = 200, 500
    xi = np.linspace(xmin, xmax, Nx)
    yi = np.linspace(ymin, ymax, Ny)

    # Create a meshgrid from the regular grid
    XI, YI = np.meshgrid(xi, yi)

    interpol = griddata((x, y), scalar_values, (XI, YI), method='linear')
    
    return interpol