最小二乘法拟合能带

from math import *
import math
import sys
from ReadEIGENVAL import *
from numpy import *
from scipy import optimize
import matplotlib.pyplot as pl

num = 4
knum = 239

lat_constant=1

Ham = [[0 for i in range(num)] for j in range(num)]

def H(kx, ky, lmd, t_bot, t_parallel):
    # lattice
    Ham[0][0] = lmd
    Ham[0][1] = 0
    Ham[0][2] = t_parallel*exp(complex(0, (kx*lat_constant*1/2+ky*lat_constant*sqrt(3)/6)))*(3/4)+t_parallel*exp(complex(0, (-kx*lat_constant*1/2+ky*lat_constant*sqrt(3)/6)))*(3/4)+t_bot*exp(complex(0, (kx*lat_constant*1/2+ky*lat_constant*sqrt(3)/6)))*(1/4)+t_bot*exp(complex(0, (-kx*lat_constant*1/2+ky*lat_constant*sqrt(3)/6)))*(1/4)+t_bot*exp(complex(0, (-ky*lat_constant/sqrt(3))))
    Ham[0][3] = t_parallel*exp(complex(0, (kx*lat_constant*1/2+ky*lat_constant*sqrt(3)/6)))*(sqrt(3)/4)-t_parallel*exp(complex(0,(-kx*lat_constant*1/2+ky*lat_constant*sqrt(3)/6)))*(sqrt(3)/4)-t_bot*exp(complex(0, (kx*lat_constant*1/2+ky*lat_constant*sqrt(3)/6)))*(sqrt(3)/4)+t_bot*exp(complex(0, (-kx*lat_constant*1/2+ky*lat_constant*sqrt(3)/6)))*(sqrt(3)/4)

    Ham[1][0] = 0
    Ham[1][1] = lmd
    Ham[1][2] = Ham[0][3]
    Ham[1][3] = t_parallel*exp(complex(0, (kx*lat_constant*1/2+ky*lat_constant*sqrt(3)/6)))*(1/4)+t_parallel*exp(complex(0, (-kx*lat_constant*1/2+ky*lat_constant*sqrt(3)/6)))*(1/4)+t_bot*exp(complex(0, (kx*lat_constant*1/2+ky*lat_constant*sqrt(3)/6)))*(3/4)+t_bot*exp(complex(0, (-kx*lat_constant*1/2+ky*lat_constant*sqrt(3)/6)))*(3/4)+t_parallel*exp(complex(0, (-ky*lat_constant/sqrt(3))))

    Ham[2][0] = Ham[0][2].conjugate()
    Ham[2][1] = Ham[1][2].conjugate()
    Ham[3][0] = Ham[0][3].conjugate()
    Ham[3][1] = Ham[1][3].conjugate()

    Ham[2][2] = -lmd
    Ham[2][3] = 0
    Ham[3][2] = 0
    Ham[3][3] = -lmd

    e, w = linalg.eig(Ham)
    e = [real(e[y]) for y in range(num)]
    e = sort(e)
    return e

nKpts, nBnds, kpts, ens1, enFermi = readEIGENVAL(1)

def fit(p):

    lmd, t_bot, t_parallel =  p

    ens = zeros((num, knum), dtype=float64)
    a1 = open("kpoints1", 'r+')
    for m in range(knum):
        Kq = a1.readline().strip().split()
        kx,ky=[float(Kq[n]) for n in range(2)]
        e=H(kx,ky, lmd, t_bot, t_parallel)
        for y in range(num):
            ens[y, m] = real(e[y])

    Y = np.concatenate(
        [ens[0, :],ens[1, :],ens[2, :],ens[3, :]
            ])

    F = np.concatenate(
        [ens1[7, :],ens1[8, :],ens1[9, :],ens1[10, :]
            ])

    diff = np.sum((Y - F) **2)

    print(p)
    print(diff)
    return diff


lmd = 100
t_bot = 200
t_parallel = -10

p =  lmd, t_bot, t_parallel
r = optimize.minimize(fit, p, method='SLSQP', options={
                        'disp': True, 'maxiter': 20000})


print(r.x)
lmd, t_bot, t_parallel = r.x
ens = zeros((num, knum), dtype=float64)
a1 = open("kpoints1", 'r+')
for m in range(knum):
    Kq = a1.readline().strip().split()
    kx,ky=[float(Kq[n]) for n in range(2)]
    e=H(kx,ky,lmd, t_bot, t_parallel)
    for y in range(num):
        ens[y, m] = real(e[y])



pl.figure(figsize=(16.0, 12.4))
pl.subplots_adjust(wspace=0.35)
pl.ylim(-4,4)
pl.xlim(0, knum)
pl.grid(True)

# print(nOrbits)
for i in range(num):
    pl.plot(ens[i], 'r-', ms=5, mew=1)



for i in range(nBnds):
    print (i)
    pl.plot(ens1[i], 'g-', ms=5, mew=1)

pl.ylabel("Energy")


pl.show()

print('Done.\n')