三对角矩阵

import numpy as np
import cmath

def ssh_hamiltonian(e,v,w,N):
    H = [[ 0 for i in range(2*N)] for j in range(2*N)]
    for i in range(2*N):
        H[i][i] = e
    for i in range(1,2*N,2):
        H[i-1][i] = v
        H[i][i-1] = v.conjugate()

        if i+1 < 2*N:
            H[i][i+1] = w
            H[i+1][i] = w.conjugate()
    a = np.linalg.eigvals(H)
    print(sorted(a))

    res = []
    for i in range(1,N+1):
        k = i*np.pi/(N+1)
        H01 = v + w.conjugate()*np.exp(-1j*k)
        H10 = v.conjugate() + w*np.exp(1j*k)
        res.append(e+cmath.sqrt(H01*H10))
        res.append(e-cmath.sqrt(H01*H10))
    print(res)

    return H

ssh_hamiltonian(e=0,v=1,w=2,N=5)



# only complex number
def hamiltonian_tri_complex_number_n(A=1,B1=2,B2=3,N=4):
    import time
    start = time.time()
    res = []
    for i in range(1,N+1):
        # print(theta(B))
        k = i*np.pi/(N+1)# - np.pi/2
        # A + 2*sqrt(B*B')* cos k
        res.append(A+2*cmath.sqrt(B1*B2)*np.cos(i*np.pi/(N+1))) ## 这个结果也是对的
        # A + B*e^{ik} + B*e^{-ik}
        # res.append(A+B1*np.exp(1j*k)+B2*np.exp(-1j*k))
    end = time.time()
    print('公式求解值时间为: {} s.'.format(end-start))
    print(sorted(np.array(res)))

    start = time.time()
    H = np.zeros((N,N),dtype="complex")
    for i in range(N):
        if i+1<N:
            H[i][i+1] = B1
            H[i+1][i] = B2
    for i in range(N):
        H[i][i] = A
    a = np.linalg.eigvals(H)
    end = time.time()
    print('严格求解值时间为: {} s.'.format(end-start))
    print(sorted(a))
    return H



def hamiltonian(x1,x2,x3,x4,x5,x6,Nx=4,Ny=4):
    # H = np.zeros((Nx*Ny,Nx*Ny),dtype="complex")
    H = [[ 0 for i in range(Nx*Ny)] for i in range(Nx*Ny)]
    for i in range(Nx):
        for j in range(Ny):
            index = Nx*j+i
            ## x1
            H[index][index] = x1

            ## x4
            ## (i, j) --> (i, j+1)
            if j + 1 < Ny:
                H[index][index+Nx] = x4
                H[index+Nx][index] = x4.conjugate()
                
            ## x2 matrix
            ## (i, j) --> (i+1, j)
            if i + 1 < Nx:
                H[index][index+1] = x2
                H[index+1][index] = x3

            ## x5
            ## (i, j) --> (i+1, j+1)
            if i + 1 < Nx and j + 1 < Ny:
                H[index][index+Nx+1] = x5
                H[index+Nx+1][index] = x5.conjugate()

            ## x6
            ## (i, j) --> (i+1, j-1)
            if i > 0 and j + 1 < Ny:
                H[index][index+Nx-1] = x6
                H[index+Nx-1][index] = x6.conjugate()
    return H


def main():
    x1 = 1
    x2 = 2
    x3 = 3
    x4 = 4 + 1j
    x5 = 0
    x6 = 0
    Nx = 4
    Ny = 4
    Ham = hamiltonian(x1=x1,x2=x2,x3=x3,x4=x4,x5=x5,x6=x6,Nx=Nx,Ny=Ny)
    a = np.linalg.eigvals(Ham)
    print(sorted(a))
    res = []
    for i in range(1,Nx+1):
        for j in range(1,Ny+1):
            kx = i*np.pi/(Nx+1)
            ky = j*np.pi/(Ny+1)

            H00 = x1 + x4*np.exp(1j*ky) + x4.conjugate()*np.exp(-1j*ky)
            H01 = x2 + x5*np.exp(1j*ky) + x6.conjugate()*np.exp(-1j*ky)
            H10 = x3 + x6*np.exp(1j*ky) + x5.conjugate()*np.exp(-1j*ky)

            tmp = H00 + 2*cmath.sqrt(H01*H10)*np.cos(kx)
            res.append(tmp)
    print(sorted(np.array(res)))
# main()

# x1 = 1
# x2 = 2
# x3 = 3
# x4 = 4
# x5 = 5
# x6 = 5
# Nx = 4
# Ny = 4

# ky = 0
# H00 = x1 + x4*np.exp(1j*ky) + x4.conjugate()*np.exp(-1j*ky)
# H01 = x2 + x5*np.exp(1j*ky) + x6.conjugate()*np.exp(-1j*ky)
# H10 = x3 + x6*np.exp(1j*ky) + x5.conjugate()*np.exp(-1j*ky)
# # hamiltonian_tri_complex_number_n(A=H00,B1=H01,B2=H10,N=4)