给出自己写的几个 Python 函数，方便自己或其他人有需要时直接过来调用。目前本篇内容不常更新，功能已集成到开源项目 Guan，可安装软件包直接调用。开源项目网址：https://py.guanjihuan.com。
1. Python一般的书写框架
import numpy as np # 导入numpy库，用来存储和处理大型矩阵，比python自带的嵌套列表更高效。numpy库还包含了许多数学函数库。python+numpy等同于matlab。
def main(): # 主函数的内容放在这里。
pass
if __name__ == '__main__': # 如果是当前文件直接运行，执行main()函数中的内容；如果是import当前文件，则不执行。同时将main()语句放在最后运行，可以避免书写的函数出现未定义的情况。
main()
2. 读取txt文件数据
补充说明：这里也可以用readline方法来写。
2.1 读取一维数据
txt文件的数据格式为（第一列为x，其他列为y）：
1 1.2 2.4
2 5.5 3.2
3 6.7 7.1
4 3.6 4.9读取txt文件的代码为：
import numpy as np
import matplotlib.pyplot as plt
# import os
# os.chdir('D:/data') # 设置路径
def main():
x, y = read_one_dimension('a.txt')
for dim0 in range(y.shape[1]):
plt.plot(x, y[:, dim0], '-k')
plt.show()
def read_one_dimension(file_name):
f = open(file_name, 'r')
text = f.read()
f.close()
row_list = np.array(text.split('\n')) # 根据“回车”分割成每一行
# print('文本格式：')
# print(text)
# print('row_list:')
# print(row_list)
# print('column:')
dim_column = np.array(row_list[0].split()).shape[0] # 列数
x = np.array([])
y = np.array([])
for row in row_list:
column = np.array(row.split()) # 每一行根据“空格”继续分割
# print(column)
if column.shape[0] != 0: # 解决最后一行空白的问题
x = np.append(x, [float(column[0])], axis=0) # 第一列为x数据
y_row = np.zeros(dim_column-1)
for dim0 in range(dim_column-1):
y_row[dim0] = float(column[dim0+1])
if np.array(y).shape[0] == 0:
y = [y_row]
else:
y = np.append(y, [y_row], axis=0)
# print('x:')
# print(x)
# print('y:')
# print(y)
return x, y
if __name__ == '__main__':
main()
2.2 读取二维数据
txt文件的数据格式为（第一个数字0不起作用，要去除。第一行是坐标x1，第一列是坐标x2，其他部分是对应的矩阵值）：
0 1 2 3 4
1 1.3 2.7 6.7 8.3
2 4.3 2.9 5.4 7.4
3 9.1 8.2 2.6 3.1读取txt文件的代码为：
import numpy as np
# import os
# os.chdir('D:/data') # 设置路径
def main():
x1, x2, matrix = read_two_dimension('b.txt')
plot_matrix(x1, x2, matrix)
def read_two_dimension(file_name):
f = open(file_name, 'r')
text = f.read()
f.close()
row_list = np.array(text.split('\n')) # 根据“回车”分割成每一行
# print('文本格式：')
# print(text)
# print('row_list:')
# print(row_list)
# print('column:')
dim_column = np.array(row_list[0].split()).shape[0] # 列数
x1 = np.array([])
x2 = np.array([])
matrix = np.array([])
for i0 in range(row_list.shape[0]):
column = np.array(row_list[i0].split()) # 每一行根据“空格”继续分割
# print(column)
if i0 == 0:
x1_str = column[1::] # x1坐标（去除第一个在角落的值）
x1 = np.zeros(x1_str.shape[0])
for i00 in range(x1_str.shape[0]):
x1[i00] = float(x1_str[i00]) # 字符串转浮点
elif column.shape[0] != 0: # 解决最后一行空白的问题
x2 = np.append(x2, [float(column[0])], axis=0) # 第一列为x数据
matrix_row = np.zeros(dim_column-1)
for dim0 in range(dim_column-1):
matrix_row[dim0] = float(column[dim0+1])
if np.array(matrix).shape[0] == 0:
matrix = [matrix_row]
else:
matrix = np.append(matrix, [matrix_row], axis=0)
# print('x1:')
# print(x1)
# print('x2:')
# print(x2)
# print('matrix:')
# print(matrix)
return x1, x2, matrix
def plot_matrix(x1, x2, matrix):
# import matplotlib as mpl # Linux下使用plt需要加上这个
# mpl.use('Agg') # Linux下使用plt需要加上这个
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter
fig = plt.figure()
ax = fig.gca(projection='3d')
x1, x2 = np.meshgrid(x1, x2)
ax.plot_surface(x1, x2, matrix, cmap=cm.coolwarm, linewidth=0, antialiased=False)
plt.xlabel('x1')
plt.ylabel('x2')
ax.set_zlabel('z')
plt.show()
if __name__ == '__main__':
main()
3. 二维数据画图或写入txt文件
3.1 用二维数据画图
import numpy as np
from math import *
# import os
# os.chdir('D:/data') # 设置路径
def main():
k1 = np.arange(-pi, pi, 0.05)
k2 = np.arange(-pi, pi, 0.05)
value = np.ones((k2.shape[0], k1.shape[0]))
plot_matrix(k1, k2, value)
def plot_matrix(k1, k2, matrix):
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter
fig = plt.figure()
ax = fig.gca(projection='3d')
k1, k2 = np.meshgrid(k1, k2)
ax.plot_surface(k1, k2, matrix, cmap=cm.coolwarm, linewidth=0, antialiased=False)
plt.xlabel('k1')
plt.ylabel('k2')
ax.set_zlabel('Z')
plt.show()
if __name__ == '__main__':
main()3.2 将二维数据写入txt文件
import numpy as np
from math import *
# import os
# os.chdir('D:/data') # 设置路径
def main():
k1 = np.arange(-pi, pi, 0.05)
k2 = np.arange(-pi, pi, 0.05)
value = np.ones((k2.shape[0], k1.shape[0]))
write_matrix(k1, k2, value)
def write_matrix(k1, k2, matrix):
with open('a.txt', 'w') as f:
# np.set_printoptions(suppress=True) # 取消输出科学记数法
f.write('0 ')
for k10 in k1:
f.write(str(k10)+' ')
f.write('\n')
i0 = 0
for k20 in k2:
f.write(str(k20))
for j0 in range(k1.shape[0]):
f.write(' '+str(matrix[i0, j0])+' ')
f.write('\n')
i0 += 1
if __name__ == '__main__':
main()4. 在运算的同时，将二维数据写入txt文件
import numpy as np
from math import *
# import os
# os.chdir('D:/data') # 设置路径
f = open('a.txt', 'w')
f.write('0 ')
for k1 in np.arange(-pi, pi, 0.05):
f.write(str(k1)+' ')
f.write('\n')
for k2 in np.arange(-pi, pi, 0.05):
f.write(str(k2)+' ')
for k1 in np.arange(-pi, pi, 0.05):
data = 1000 # 运算数据
f.write(str(data)+' ')
f.write('\n')
f.close()
5. 求本征值，画能带图或写入txt文件
5.1 求本征值，画一维能带图
import numpy as np
from math import *
# import os
# os.chdir('D:/data') # 设置路径
def hamiltonian(k):
pass
def main():
k = np.arange(-pi, pi, 0.05)
plot_bands_one_dimension(k, hamiltonian)
def plot_bands_one_dimension(k, hamiltonian):
import matplotlib.pyplot as plt
dim = hamiltonian(0).shape[0]
dim_k = k.shape[0]
eigenvalue_k = np.zeros((dim_k, dim))
i0 = 0
for k0 in k:
matrix0 = hamiltonian(k0)
eigenvalue, eigenvector = np.linalg.eig(matrix0)
eigenvalue_k[i0, :] = np.sort(np.real(eigenvalue[:]))
i0 += 1
for dim0 in range(dim):
plt.plot(k, eigenvalue_k[:, dim0], '-k')
plt.xlabel('k')
plt.ylabel('E')
plt.show()
if __name__ == '__main__':
main()5.2 求本征值，画二维能带图
import numpy as np
from math import *
# import os
# os.chdir('D:/data') # 设置路径
def hamiltonian(k1, k2):
pass
def main():
k1 = np.arange(-pi, pi, 0.05)
k2 = np.arange(-pi, pi, 0.05)
plot_bands_two_dimension(k1, k2, hamiltonian)
def plot_bands_two_dimension(k1, k2, hamiltonian):
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter
dim = hamiltonian(0, 0).shape[0]
dim1 = k1.shape[0]
dim2 = k2.shape[0]
eigenvalue_k = np.zeros((dim2, dim1, dim))
i0 = 0
for k20 in k2:
j0 = 0
for k10 in k1:
matrix0 = hamiltonian(k10, k20)
eigenvalue, eigenvector = np.linalg.eig(matrix0)
eigenvalue_k[i0, j0, :] = np.sort(np.real(eigenvalue[:]))
j0 += 1
i0 += 1
fig = plt.figure()
ax = fig.gca(projection='3d')
k1, k2 = np.meshgrid(k1, k2)
for dim0 in range(dim):
ax.plot_surface(k1, k2, eigenvalue_k[:, :, dim0], cmap=cm.coolwarm, linewidth=0, antialiased=False)
plt.xlabel('k1')
plt.ylabel('k2')
ax.set_zlabel('E')
plt.show()
if __name__ == '__main__':
main()5.3 求本征值，将一维数据写入txt文件
import numpy as np
from math import *
# import os
# os.chdir('D:/data') # 设置路径
def hamiltonian(k):
pass
def main():
k = np.arange(-pi, pi, 0.05)
write_bands_one_dimension(k, hamiltonian)
def write_bands_one_dimension(k, hamiltonian):
dim = hamiltonian(0).shape[0]
f = open('a.txt','w')
for k0 in k:
f.write(str(k0)+' ')
matrix0 = hamiltonian(k0)
eigenvalue, eigenvector = np.linalg.eig(matrix0)
eigenvalue = np.sort(np.real(eigenvalue))
for dim0 in range(dim):
f.write(str(eigenvalue[dim0])+' ')
f.write('\n')
f.close()
if __name__ == '__main__':
main()5.4 求本征值，将二维数据写入txt文件
import numpy as np
from math import *
# import os
# os.chdir('D:/data') # 设置路径
def hamiltonian(k1, k2):
pass
def main():
k1 = np.arange(-pi, pi, 0.05)
k2 = np.arange(-pi, pi, 0.05)
write_bands_two_dimension(k1, k2, hamiltonian)
def write_bands_two_dimension(k1, k2, hamiltonian):
f1 = open('a1.txt', 'w')
f2 = open('a2.txt', 'w')
f1.write('0 ')
f2.write('0 ')
for k10 in k1:
f1.write(str(k10)+' ')
f2.write(str(k10)+' ')
f1.write('\n')
f2.write('\n')
for k20 in k2:
f1.write(str(k20)+' ')
f2.write(str(k20)+' ')
for k10 in k1:
matrix0 = hamiltonian(k10, k20)
eigenvalue, eigenvector = np.linalg.eig(matrix0)
eigenvalue = np.sort(np.real(eigenvalue))
f1.write(str(eigenvalue[0])+' ')
f2.write(str(eigenvalue[1])+' ')
f1.write('\n')
f2.write('\n')
f1.close()
f2.close()
if __name__ == '__main__':
main()【说明：本站主要是个人的一些笔记和代码分享，内容可能会不定期修改。为了使全网显示的始终是最新版本，这里的文章未经同意请勿转载。引用请注明出处：https://www.guanjihuan.com】
请问5.2中"eigenvalue_k[i0, j0, :] = np.sort(np.real(eigenvalue[:]))"这一块i0和j0的位置是不是写反了呢？
都可以，这里是为了和画图联系在一起，所以顺序这么排。主要看画图那边是怎么显示的。