Search code examples
pythonleast-squares

Least square fit a 2D line

I realize I could use numpy to find the line like so:

import numpy as np
import matplotlib.pyplot as plt
a = np.array([1,2,3,4,6,7])
b = np.array([5,4,3,2,-2,-1])
k,m = np.polyfit(a,b,1)
plt.scatter(a,b)
plt.plot([0,10],[m,10*k+m])
plt.show()

enter image description here

but I'd like to use raw python code instead. My math is too rusty, but if can be done in a few lines of code I'd really appreciate the help!

Solution

  • If you are looking for a simple linear regression based on minimizing the quadratic error, the pure Python implementation is pretty straightforward (check equations for α and β on the link above):

    def linear_fit(x, y):
    """For set of points `(xi, yi)`, return linear polynomial `f(x) = k*x + m` that
    minimizes the sum of quadratic errors.
    """
    meanx = sum(x) / len(x)
    meany = sum(y) / len(y)
    k = sum((xi-meanx)*(yi-meany) for xi,yi in zip(x,y)) / sum((xi-meanx)**2 for xi in x)
    m = meany - k*meanx
    return k, m

    For your sample input:

    >>> x = [1,2,3,4,6,7]
>>> y = [5,4,3,2,-2,-1]
>>> linear_fit(x, y)
(-1.1614906832298135, 6.285714285714285)