Python专题, 语言

自己写的几个Python函数

给出自己写的几个 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()
1,412 次浏览

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

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

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

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

发表评论

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

Captcha Code