python最小二乘法

#文件名core.py
import numpy as np
def leastSquare(x,y):
    if len(x)==2:
    #此时x为自然序列
        sx = 0.5*(x[1]-x[0]+1)*(x[1]+x[0])
        ex = sx/(x[1]-x[0]+1)
        sx2 = ((x[1]*(x[1]+1)*(2*x[1]+1))
              -(x[0]*(x[0]-1)*(2*x[0]-1)))/6
        x = np.array(range(x[0],x[1]+1))
    else:
        sx = sum(x)
        ex = sx/len(x)
        sx2 = sum(x**2)
    
    sxy = sum(x*y)
    ey = np.mean(y)

    a = (sxy-ey*sx)/(sx2-ex*sx)
    b = (ey*sx2-sxy*ex)/(sx2-ex*sx)
    return a,b


x = np.array([125,118,135,131.5,111,80,54.5,33.35,0,
              135.167,135.33,143.33,111.17,111.17,83.33,68.33,30.99,0,
             141.4,122,35.7,50,27,10,0,
             146.5,124.286,64.29,31.71,13.929,0,
             147.5,127.5,125,98.75,50,30,15,5,8,0,
             124.99,91.67,49.99,16.67,9.99,0,
             141.7,127.5,119,92,65,30.5,16.5,0,
             232.143,199.99,133.93,128.21,124.99,96.428,57.143,32.143,6.571,0,
             272.75,55.85,131.25,60,20,5,0,
             330.99,136.67,116.67,115.83,90.83,74.99,58.33,43.33,0,
             198,193.42,99.2,96,63.2,38,21,0])   #数据 
y = np.array([812.5,1687.5,3087.5,3087.5,2988.75,2986.95,2985.55,2984.6,2983.5,
              470.833,1045.833,1533.33,1916.67,2416.67,2833.33,3124.99,3441.67,3716.67,
              505,1305,1825,2275,2442.5,2535,2590,
              428.57,1174.99,1646.43,1886.43,1993.57,2049.99,
              252.5,527.5,780,1102.5,1375,1495,1547.5,1557.5,1580,1605,
              308.33,524.99,666.67,724.99,743.33,764.99,
              117.5,185.5,270.5,376,454.5,502.25,525.75,549,
              23.21,66.43,99.82,138.86,201.78,287.49,364.2857,408.9,433.21,433.57,
              90,188.75,238.75,295,335,350,355.5,
              58.167,59.83,78.33,101.58,122.249,138.83,152.17,174.33,187.99,
              24.8,55.675,78.8,94.6,112.1,124.75,134.75,144.9])
a,b = leastSquare(x,y)
print(a,b)
plt.scatter(x,y)
plt.plot(x,a*x+b)
plt.show()