使用cupy优化能带计算

from math import pi
import numpy as np
import cupy as cp
from numpy import exp, sqrt
from numpy.linalg import eigh
import matplotlib.pylab as plt
from numba import njit


def tridiag_block_matrix(H, c, u, d):
    # c, u, d are center, upper and lower blocks
    N, _ = H.shape
    n, _ = c.shape
    H[0:n, 0:n] = c
    for i in range(n, N, n):
        H[i:i+n, i:i+n] = c
        H[i-n:i, i:i+n] = d
        H[i:i+n, i-n:i] = u
    return H


def H0_SU4(H, k, t = 1.0):
    A =  cp.array([[sqrt(3),0,1,exp(-1j*k)],[0,sqrt(3),1,1],[1,1,sqrt(3),0],[exp(1j*k),1,0,sqrt(3)]])
    B =  cp.array([[0,0,0,0],[0,0,0,0],[1,0,0,0],[0,-1,0,0]])
    return tridiag_block_matrix(H, A, B.T, B)
    

def H0_Graphene(H, k, t = -1):
    A = np.array([[0, t*(1+exp(-1j* k))], [t*(1+ exp(1j*k)), 0]])
    B = np.array([[0, 0], [t, 0]])
    return tridiag_block_matrix(H, A, B, B.T)
    

def calculated_band(H, ks, t = 1):
    nk = ks.size
    m, _ = H.shape
    band_gpu = cp.zeros((nk, m - 1))
    Hk_gpu = cp.zeros((m, m), dtype=np.complex64)
    for i in range(nk):
        Hk_gpu = H0_SU4(Hk_gpu, ks[i])
        E, _ = cp.linalg.eigh(Hk_gpu[1:,1:])
        band_gpu[i, :] = E
    return cp.asnumpy(band_gpu)

def plot_SU4_Band():
    nk = 256
    N = 256
    ks = np.linspace(0, 2*pi, nk)
    Hk = np.zeros((4*N, 4*N), dtype=np.complex64)
    band = calculated_band(Hk, ks)
    plt.plot(band, color = "gray")
    plt.plot(band[:, N - 1], color = "red")
    plt.xticks(np.arange(0, nk, nk//3), ['0', '2/3π', '4/3π', '2π'], fontsize = 12, fontweight = 'bold')
    plt.yticks(np.arange(-0.5, 0.6, 0.5), fontsize = 12, fontweight = 'bold')
    plt.ylim(-0.5, 0.5)
    plt.xlabel("k", fontsize = 13, fontweight = 'bold')
    plt.ylabel("E", fontsize = 13, fontweight = 'bold')
    plt.show()

if __name__ == '__main__':
    plot_SU4_Band()

将numpy矩阵操作，换成对应cupy对这些函数的gpu封装。达到gpu加速的目的