""" This code is supported by the website: https://www.guanjihuan.com The newest version of this code is on the web page: https://www.guanjihuan.com/archives/17984 """ import numpy as np import cmath from math import * import functools def hamiltonian(kx, ky): # BBH model # label of atoms in a unit cell # (2) —— (0) # | | # (1) —— (3) gamma_x = 0.5 # hopping inside one unit cell lambda_x = 1 # hopping between unit cells gamma_y = gamma_x lambda_y = lambda_x h = np.zeros((4, 4), dtype=complex) h[0, 2] = gamma_x+lambda_x*cmath.exp(1j*kx) h[1, 3] = gamma_x+lambda_x*cmath.exp(-1j*kx) h[0, 3] = gamma_y+lambda_y*cmath.exp(1j*ky) h[1, 2] = -gamma_y-lambda_y*cmath.exp(-1j*ky) h[2, 0] = np.conj(h[0, 2]) h[3, 1] = np.conj(h[1, 3]) h[3, 0] = np.conj(h[0, 3]) h[2, 1] = np.conj(h[1, 2]) return h def main(): kx = np.arange(-pi, pi, 0.05) ky = np.arange(-pi, pi, 0.05) eigenvalue_array = calculate_eigenvalue_with_two_parameters(kx, ky, hamiltonian) plot_3d_surface(kx, ky, eigenvalue_array, xlabel='kx', ylabel='ky', zlabel='E', title='BBH bands') hamiltonian0 = functools.partial(hamiltonian, ky=0) eigenvalue_array = calculate_eigenvalue_with_one_parameter(kx, hamiltonian0) plot(kx, eigenvalue_array, xlabel='kx', ylabel='E', title='BBH bands ky=0') # import guan # eigenvalue_array = guan.calculate_eigenvalue_with_two_parameters(kx, ky, hamiltonian) # guan.plot_3d_surface(kx, ky, eigenvalue_array, xlabel='kx', ylabel='ky', zlabel='E', title='BBH bands') # hamiltonian0 = functools.partial(hamiltonian, ky=0) # eigenvalue_array = guan.calculate_eigenvalue_with_one_parameter(kx, hamiltonian0) # guan.plot(kx, eigenvalue_array, xlabel='kx', ylabel='E', title='BBH bands ky=0') def calculate_eigenvalue_with_one_parameter(x_array, hamiltonian_function, print_show=0): dim_x = np.array(x_array).shape[0] i0 = 0 if np.array(hamiltonian_function(0)).shape==(): eigenvalue_array = np.zeros((dim_x, 1)) for x0 in x_array: hamiltonian = hamiltonian_function(x0) eigenvalue_array[i0, 0] = np.real(hamiltonian) i0 += 1 else: dim = np.array(hamiltonian_function(0)).shape[0] eigenvalue_array = np.zeros((dim_x, dim)) for x0 in x_array: if print_show==1: print(x0) hamiltonian = hamiltonian_function(x0) eigenvalue, eigenvector = np.linalg.eigh(hamiltonian) eigenvalue_array[i0, :] = eigenvalue i0 += 1 return eigenvalue_array def calculate_eigenvalue_with_two_parameters(x_array, y_array, hamiltonian_function, print_show=0, print_show_more=0): dim_x = np.array(x_array).shape[0] dim_y = np.array(y_array).shape[0] if np.array(hamiltonian_function(0,0)).shape==(): eigenvalue_array = np.zeros((dim_y, dim_x, 1)) i0 = 0 for y0 in y_array: j0 = 0 for x0 in x_array: hamiltonian = hamiltonian_function(x0, y0) eigenvalue_array[i0, j0, 0] = np.real(hamiltonian) j0 += 1 i0 += 1 else: dim = np.array(hamiltonian_function(0, 0)).shape[0] eigenvalue_array = np.zeros((dim_y, dim_x, dim)) i0 = 0 for y0 in y_array: j0 = 0 if print_show==1: print(y0) for x0 in x_array: if print_show_more==1: print(x0) hamiltonian = hamiltonian_function(x0, y0) eigenvalue, eigenvector = np.linalg.eigh(hamiltonian) eigenvalue_array[i0, j0, :] = eigenvalue j0 += 1 i0 += 1 return eigenvalue_array def plot(x_array, y_array, xlabel='x', ylabel='y', title='', fontsize=20, labelsize=20, show=1, save=0, filename='a', file_format='.jpg', dpi=300, style='', y_min=None, y_max=None, linewidth=None, markersize=None, adjust_bottom=0.2, adjust_left=0.2): import matplotlib.pyplot as plt fig, ax = plt.subplots() plt.subplots_adjust(bottom=adjust_bottom, left=adjust_left) ax.plot(x_array, y_array, style, linewidth=linewidth, markersize=markersize) ax.grid() ax.set_title(title, fontsize=fontsize, fontfamily='Times New Roman') ax.set_xlabel(xlabel, fontsize=fontsize, fontfamily='Times New Roman') ax.set_ylabel(ylabel, fontsize=fontsize, fontfamily='Times New Roman') if y_min!=None or y_max!=None: if y_min==None: y_min=min(y_array) if y_max==None: y_max=max(y_array) ax.set_ylim(y_min, y_max) ax.tick_params(labelsize=labelsize) labels = ax.get_xticklabels() + ax.get_yticklabels() [label.set_fontname('Times New Roman') for label in labels] if save == 1: plt.savefig(filename+file_format, dpi=dpi) if show == 1: plt.show() plt.close('all') def plot_3d_surface(x_array, y_array, matrix, xlabel='x', ylabel='y', zlabel='z', title='', fontsize=20, labelsize=15, show=1, save=0, filename='a', file_format='.jpg', dpi=300, z_min=None, z_max=None, rcount=100, ccount=100): import matplotlib.pyplot as plt from matplotlib import cm from matplotlib.ticker import LinearLocator matrix = np.array(matrix) fig, ax = plt.subplots(subplot_kw={"projection": "3d"}) plt.subplots_adjust(bottom=0.1, right=0.65) x_array, y_array = np.meshgrid(x_array, y_array) if len(matrix.shape) == 2: surf = ax.plot_surface(x_array, y_array, matrix, rcount=rcount, ccount=ccount, cmap=cm.coolwarm, linewidth=0, antialiased=False) elif len(matrix.shape) == 3: for i0 in range(matrix.shape[2]): surf = ax.plot_surface(x_array, y_array, matrix[:,:,i0], rcount=rcount, ccount=ccount, cmap=cm.coolwarm, linewidth=0, antialiased=False) ax.set_title(title, fontsize=fontsize, fontfamily='Times New Roman') ax.set_xlabel(xlabel, fontsize=fontsize, fontfamily='Times New Roman') ax.set_ylabel(ylabel, fontsize=fontsize, fontfamily='Times New Roman') ax.set_zlabel(zlabel, fontsize=fontsize, fontfamily='Times New Roman') ax.zaxis.set_major_locator(LinearLocator(5)) ax.zaxis.set_major_formatter('{x:.2f}') if z_min!=None or z_max!=None: if z_min==None: z_min=matrix.min() if z_max==None: z_max=matrix.max() ax.set_zlim(z_min, z_max) ax.tick_params(labelsize=labelsize) labels = ax.get_xticklabels() + ax.get_yticklabels() + ax.get_zticklabels() [label.set_fontname('Times New Roman') for label in labels] cax = plt.axes([0.8, 0.1, 0.05, 0.8]) cbar = fig.colorbar(surf, cax=cax) cbar.ax.tick_params(labelsize=labelsize) for l in cbar.ax.yaxis.get_ticklabels(): l.set_family('Times New Roman') if save == 1: plt.savefig(filename+file_format, dpi=dpi) if show == 1: plt.show() plt.close('all') if __name__ == '__main__': main()