python最小二乘法代码

备份一下解决某个问题使用到的最小二乘法的代码。

执行结果

绘图显示计算结果是正确的。效果看下图。蓝色的线是添加了随机噪声的点连成的，红色的线是计算后拟合得到的。

计算逻辑

如下。数据集P是测量得到的n个采样点，含有噪声。

然后按照下方的计算逻辑算出直线y=kx+b的斜率k和常数b。

代码

# coding=utf-8

# 这个是利用一元一次函数，生成平面上的点并添加噪声，再用最小二乘法拟合这些点的
# 测试代码。
# 使用方法：直接执行代码，重定向输出到一个csv文件，然后用Excel打开，添加
# 图表查看结果。
# 例：
# python run.py > result.csv

import random

def solve(P):
    # 格式： P=[{'x':1, 'y':2}, {'x':2, 'y':4}]
    S_x = 0
    S_y = 0
    S_x2 = 0
    S_xy = 0
    n = len(P)

    for p in P:
        S_x = S_x + p['x']
        S_y = S_y + p['y']
        S_x2 = S_x2 + p['x'] * p['x']
        S_xy = S_xy + p['x'] * p['y']

    avg_x = S_x / n
    avg_y = S_y / n

    # 算出参数k和b，对应变量是h_k和h_b：
    h_k = (S_xy - n * avg_x * avg_y) / (S_x2 - n * avg_x * avg_x)
    h_b = avg_y - h_k * avg_x
    return {'h_k': h_k, 'h_b': h_b}

if __name__ == '__main__':

    # y=kx+b
    # k是斜率，b是常数
    k = 0.1
    b = 4.0

    # 对一元一次函数按照x坐标进行采样，采样的数量
    n = 100

    # 列表，按照顺序存储生成的采样点的数据，每个元素的结构：{'x': x, 'y': y}
    P = []

    # 生成采样数据并添加噪声
    for i in range(n):
        x = i + 1
        y = k * x + b + random.gauss(0, 1.0)
        P.append({'x': x, 'y': y})

    result = solve(P) # 求解 k 和 b
    h_k = result['h_k']
    h_b = result['h_b']

    # 输出结果
    print(',,,----------------------------------------')
    print('b,k')
    print(str(h_b) + ',' + str(h_k))
    print(',,,----------------------------------------')

    # 输出数据
    print('x,y,h_y')
    for p in P:
        print(str(p['x']) + ',' + str(p['y']) + ',' + str(h_k * p['x'] + h_b))