历史文章/不再维护

自己写的几个Python函数

发布时间：

最近修改时间：2023年11月6日

给出自己写的几个Python函数，方便自己或其他人有需要时直接过来调用。

目前本篇内容不常更新，功能已集成到开源项目Guan，可安装软件包直接调用。开源项目网址：https://py.guanjihuan.com。

目录

  1. Python一般的书写框架
  2. 读取txt文件数据
    2.1 读取一维数据
    2.2 读取二维数据
  3. 二维数据画图或写入txt文件
    3.1 用二维数据画图
    3.2 将二维数据写入txt文件
  4. 在运算的同时，将二维数据写入txt文件
  5. 求本征值，画能带图或写入txt文件
    5.1 求本征值，画一维能带图
    5.2 求本征值，画二维能带图
    5.3 求本征值，将一维数据写入txt文件
    5.4 求本征值，将二维数据写入txt文件

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()
1,336 次浏览

【说明：本站主要是个人的一些笔记和代码分享，内容可能会不定期修改。为了使全网显示的始终是最新版本，这里的文章未经同意请勿转载。引用请注明出处：https://www.guanjihuan.com】

头像

Published by guanjihuan

2 thoughts on “自己写的几个Python函数”

  1. 请问5.2中"eigenvalue_k[i0, j0, :] = np.sort(np.real(eigenvalue[:]))"这一块i0和j0的位置是不是写反了呢？

    回复

    1. 都可以，这里是为了和画图联系在一起，所以顺序这么排。主要看画图那边是怎么显示的。

      回复

发表评论

您的电子邮箱地址不会被公开。 必填项已用 * 标注

Captcha Code