Matplotlib-For-Lab-of-Physics

之前写的绘制最小二乘法代码太垃了，现在使用面向对象的技术对截距为 \(0\) 和不为 \(0\) 封装为两个类，并共同继承拥有一个父类 \(\mbox{regression}\) 在命令行执行时输入相应参数，对应两种不同拟合方式，以后有时间可以增加多项式拟合相关的类。

源码以及对应详尽注释如下：（绘制拟合直线和计算相应误差并输出为 \(\LaTeX\) 格式 ）

import argparse
from math import sqrt,log,pi
import numpy as np
from IPython import embed
import matplotlib.pyplot as plt
from TINV import TINV
# 存储xy坐标名称

xNames = ["谐频次数$\mathcal{n}$(个)", "弦长倒数$\ \mathcal{L}^{-1}(\mathcal{m}$$^{-1}$)", 
            "拉力对数$\ln\mathcal{T}}$", "线密度对数$\ \lnρ$", "振动次数 n", "振幅平方$A^2(rad^2$)"]

yNames = ["共振频率$\mathcal{f}\ \mathcal{(Hz)}$", "共振频率$\mathcal{f}\ \mathcal{(Hz)}$", 
            "基频对数$\ln\mathcal{f}}$", "基频对数$\ln\mathcal{f}}$", "振幅($\mathrm{rad}$)的对数 $\mathrm{ln}\  θ$", "周期$T(s)$"]

# 当前选择的名称，默认为列表最后一个
xName = xNames[-1]
yName = yNames[-1]

# 传入截距是否为0的参数
def getInput():

    parser = argparse.ArgumentParser(description = 'Version: v1.0. Draw the regression line for the Physics Report,' 
                        'coefficient 0 means the intercept is 0, '
                        '1 means the intercept is not 1, '
                        '01 means the intercept is 0 and save the figure, '
                        '11 means the intercept is not 0 and save the figure.'
                        '\nFor example, python LinearRegression.py 0')
    parser.add_argument('crossOrigin', help = 'Whether the intercept of the line is 0 or 1. ' 
                                        'Besides, if want to save the figure, just add 1. ')
    args = parser.parse_args()
    return args.crossOrigin

# 截距为0和不为0的父类
class regression:

    # 初始化相关数据
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.length = len(x)
        self.k = 0.0
        self.b = 0.0
        self.R_2 = 0.0
        self.Sk = 0.0
        self.Uk = 0.0
        self.Sb = 0.0
        self.Ub = 0.0
    
    # 将计算结果表达为latex格式
    def expression(self):
        lineExpression = "$\hat{y}=" + str(round(self.k, 6)) + "x"
        R2Expression = '$r^2=' + str(self.R_2) + "$"
        if self.b == 0.0:
            b = ""
        elif self.b > 0:
            b = "+" + str(round(self.b, 6))
        else:
            b = str(round(self.b, 6))
        lineExpression += (b + "$")
        return lineExpression, R2Expression

    # 绘制直线以及离散点
    def drawLineAndPoints(self):
        xdraw = np.linspace(min(self.x), max(self.x), 50)
        ydraw = self.k * xdraw + self.b
        plt.plot(xdraw, ydraw, '-' , color = 'b', linewidth = 0.5)
        plt.plot(self.x, self.y, 'ro')

    # 将拟合直线表达式和相关系数自动标定在合适位置
    def addTwoExpressions(self):
        addLimx = (max(self.x) - min(self.x)) / 10
        addLimy = (max(self.y) - min(self.y)) / 10
        plt.xlim(min(self.x) - addLimx, max(self.x) + addLimx)
        plt.ylim(min(self.y) - addLimy, max(self.y) + addLimy)

        deltaY = max(self.y ) - min(self.y) + 2 * addLimy
        deltaX = max(self.x) - min(self.x) + 2 * addLimx
        middleX = (max(self.x) + min(self.x)) / 2
        middleY = (max(self.y) + min(self.y)) / 2

        if self.k < 0:
            lineY1 = middleY + deltaY * 0.24
            lineY2 = middleY + deltaY * 0.16
            lineX1 = middleX - deltaX * 0.088
            lineX2 = middleX - deltaX * 0.088
        else:
            lineY1 = middleY - deltaY * 0.16
            lineY2 = middleY - deltaY * 0.24
            lineX1 = middleX - deltaX * 0.088
            lineX2 = middleX - deltaX * 0.088

        lineExpression, R2Expression = self.expression()
        plt.text(lineX1, lineY1, lineExpression, fontsize = 12)
        plt.text(lineX2, lineY2, R2Expression, fontsize = 12)

    # 添加xy轴标签
    def addLabels(self):
        plt.rcParams['font.sans-serif'] = 'Kaiti' # 文字使用楷体
        plt.xlabel(xName, fontsize = 12)
        plt.ylabel(yName, fontsize = 12)

    # 绘制图形
    def drawGraph(self):
        plt.rcParams['axes.unicode_minus'] = False  # 正常显示负数
        plt.rcParams['font.sans-serif'] ='Times New Roman'  # 数字使用罗马字体
        self.drawLineAndPoints()
        self.addTwoExpressions()
        self.addLabels()
        self.addLabels()

    # 输出数据
    def dataOutput(self):
        print("\quad \quad \quad 线性拟合斜率$k=", round(self.k, 10), end = "$")
        print("\quad \quad 调用Excel中TINV函数", "$t_{0.95,", str(self.length - 1) + "}=", round(float(TINV[self.length - 1]), 10), "$\par") 
        print("\quad \quad \quad 斜率标准偏差$s_k=", round(self.Sk, 10), end = "$")
        print("\quad \quad 斜率不确定度为$U_k=", round(self.Uk, 10), "$\par")

        if(self.Sb != 0.0):
            print("\quad \quad \quad 线性拟合截距$b=", round(self.b, 10), end = "$")
            print("\quad \quad 线性回归相关系数", "$R^2=", round(self.R_2, 10), "$\par") 
            print("\quad \quad \quad 截距标准偏差$s_b=", round(self.Sb, 10), end = "$")
            print("\quad \quad 截距不确定度为$U_b=", round(self.Ub,10), "$\par")

    # 预览但不保存图片
    def output(self):
        self.dataOutput()
        self.drawGraph()
        plt.show()

    # 保存且预览图片
    def save(self, saveName):
        self.dataOutput()
        self.drawGraph()
        plt.savefig(saveName)
        plt.show()

# 继承父类，不强制截距为0
class notCrossOrigin(regression):

    # 计算父类中的所有数据
    def __init__(self, x, y):
        regression.__init__(self, x, y)
        self.k, self.b = self.regressK_b()
        self.R_2 = self.getR_2()
        self.Sk = sqrt((1 / self.R_2 - 1) / (self.length - 2)) * self.k
        self.Uk = self.Sk * TINV[self.length - 1]
        self.Sb = sqrt((1 - self.R_2) / (self.length - 2)) * self.b
        self.Ub = self.Sb * TINV[self.length - 1]

    # 计算线性回归斜率和截距
    def regressK_b(self):
        coef = np.polyfit(self.x, self.y, 1)
        return coef[0], coef[1]
    
    # 计算线性回归相关系数
    def getR_2(self):
        coeffs = np.polyfit(self.x, self.y, 1)
        p = np.poly1d(coeffs)
        yhat = p(self.x)
        ybar = np.sum(self.y) / len(self.y)
        ssreg = np.sum((yhat - ybar) ** 2)
        sstot = np.sum((self.y - ybar) ** 2)
        return ssreg / sstot

# 继承父类，强制截距为0
class crossOrigin(regression):

    # 计算父类中的所有数据
    def __init__(self, x, y):
        regression.__init__(self, x, y)
        self.k = np.dot(x, y) / np.dot(x, x)
        self.R_2, self.Sk= self.getR_2AndS_k()
        self.Uk = self.Sk * TINV[self.length - 1]
    
    # 计算线性回归相关系数和斜率标准偏差
    def getR_2AndS_k(self):
        yhat = [self.k * self.x[i] for i in range(self.length)]
        deltay = [self.y[i] - yhat[i] for i in range(self.length)]
        ybar = np.sum(self.y) / self.length
        newReg = np.dot(deltay, deltay)
        sstot = np.dot(self.y - ybar, yhat - ybar)
        sy = sqrt(newReg / (len(self.y) - 1))
        xrms = sqrt(np.dot(self.x, self.x))
        return (1 - newReg / sstot), sy / xrms

# 主函数定义开始
if __name__ == '__main__':

    # 原始数据，如果需要处理，添加循环语句进行计算
    x = range(1,7,1)
    y = [31.2,62.9,95.5,128.4,161.0,194.3]

    # 获取类型参数
    cross = getInput()

    # 对不同参数进行回归
    if cross == "0":
        crossOrigin(x, y).output()
    elif cross == "1":
        notCrossOrigin(x, y).output()

    elif cross == "01":
        crossOrigin(x, y).save("crossOrigin.png")
    elif cross == "11":
        notCrossOrigin(x, y).save("notCrossOrigin.png")

绘制可微的 \(Bezeir\) 曲线源代码如下

from matplotlib import pyplot as plt
from scipy.special import comb
from math import sqrt,pi,exp

# 全局变量，每次绘图前修改曲线颜色，xlabel,ylabel代表横纵坐标
colorLineWhole = "blue"
xLabel = ["Forced frequency $\omega_{f} (\mathrm{s}^{-1})$"]
yLabel = ["Amplitude θ (°)","Phase Difference φ","Amplitude θ (°)"]

# 使用 python 中面向对象特性，编写贝塞尔曲线类
class BezierCurve:

    # 初始化点坐标和控制点个数
    def __init__(self, list_of_points: list[tuple[float, float]]):
        self.list_of_points = list_of_points
        self.degree = len(list_of_points) - 1

    # Bezier 曲线基函数对应多项式系数数组 B_{i,n}(t) = C_{n}^{i} * (1 - t)^{n - i} * t^{i}
    def basis_function(self, t: float) -> list[float]:
        assert 0 <= t <= 1, "Time t must be between 0 and 1."
        output_values: list[float] = []
        for i in range(len(self.list_of_points)):
            output_values.append(comb(self.degree, i) * ((1 - t) ** (self.degree - i)) * (t**i))
        return output_values

    # 计算 Bezier 曲线函数在 t∈[0,1] 处的值
    def bezier_curve_function(self, t: float) -> tuple[float, float]:
        assert 0 <= t <= 1, "Time t must be between 0 and 1."
        basis_function = self.basis_function(t)
        x = 0.0
        y = 0.0
        for i in range(len(self.list_of_points)):
            x += basis_function[i] * self.list_of_points[i][0]
            y += basis_function[i] * self.list_of_points[i][1]
        return (x, y)

    # 绘制给定 n 个控制点下的 Bezier 曲线
    def plot_curve(self, step_size: float = 0.001):
        to_plot_x: list[float] = []
        to_plot_y: list[float] = []
        t = 0.0
        while t <= 1:
            value = self.bezier_curve_function(t)
            to_plot_x.append(value[0])
            to_plot_y.append(value[1])
            t += step_size
        x = [i[0] for i in self.list_of_points]
        y = [i[1] for i in self.list_of_points]
        plt.plot(to_plot_x, to_plot_y, color = colorLineWhole, linewidth = 0.88)

#  x 排序从大到小，并保证数据对应
def sortXY(x, y):
    data = []
    for i in range(len(x)):
        data.append([x[i],y[i]])
    data.sort(key = lambda x: x[0])
    x = [i[0] for i in data]
    y = [i[1] for i in data]
    return x, y

# 计算两个点的欧式距离
def distance(p1,p2):
    return sqrt((p1[0] - p2[0]) * (p1[0] - p2[0]) + (p1[1] - p2[1]) * (p1[1] - p2[1]))

# 计算四个点对应两条直线的交点
def intersectPoints(p1, p2, p3, p4):
    k1 = (p2[1] - p1[1]) / (p2[0] - p1[0])
    b1 = p1[1] - k1 * p1[0]
    k2 = (p4[1] - p3[1]) / (p4[0] - p3[0])
    b2 = p3[1] - k2 * p3[0]
    return [(b1 - b2) / (k2 - k1), (k2 * b1 - k1 * b2) / (k2 - k1)]

# 计算由上述所生成三个点所构成三角形的内心
def angularBisector(p1,p2,p3,p4):
    interP = intersectPoints(p1,p2,p3,p4)
    a = distance(p2, p3)
    b = distance(p3, interP)
    c = distance(p2, interP)
	return [(a * interP[0] + b * p2[0] + c * p3[0]) / (a + b + c),
             (a * interP[1] + b * p2[1] + c * p3[1]) / (a + b + c) ]

# 计算给定离散点生成的 Bezier 曲线
def pointdBezierGraph(x,y,colorPoint):
    for i in range(len(x)):
        plt.scatter(x[i], y[i], color = colorPoint, marker = "+", s = 24, linewidth = 0.88)
    averageX = (x[-1] - x[0]) / len(x)
    x.insert(0, x[0] - averageX)
    x.append(x[-1] + averageX)
    y.insert(0, y[0])
    y.append(y[-1])
    for i in range(len(x) - 3):
        p1 = [x[i], y[i]]
        p2 = [x[i + 1], y[i + 1]]
        p3 = [x[i + 2], y[i + 2]]
        p4 = [x[i + 3], y[i + 3]]
        I = angularBisector(p1, p2, p3, p4)
        BezierCurve([p2, I, p3]).plot_curve()

# 自动处理数据，传入两个列表，线条颜色，标记点颜色， 绘制 Bezier 曲线
def drawBezier(x,y,colorLine,colorPoint):
    x,y = sortXY(x,y)
    global colorLineWhole
    colorLineWhole = colorLine
    pointdBezierGraph(x,y,colorPoint)

# 主函数定义开始
if __name__ == "__main__":
    x = [1.575,1.553,1.531,1.516,1.501,1.487,1.479,1.472,1.465,1.451,1.443,1.427,1.413,1.391,1.369]
    for i in range(len(x)):
        x[i] = 2 * pi / x[i]

    y = [15,25,35,44,58,76,87,99,112,130,138,147,152,156,160]
    drawBezier(x,y,"green","blue")

    plt.xlabel(xLabel[1])
    plt.ylabel(yLabel[0])
    # plt.savefig("图片.png",dpi = 2000)
    plt.show()