linreg.py

from __future__ import division

class Matrix:
    def __init__(self, v):
        self.value = [list(i) for i in v]
        self.rows = len(self.value)
        self.columns = len(self.value[0])
    def __iter__(self):
        return iter(self.value)
    def __repr__(self):
        return str(self.value)
    def __len__(self):
        return len(self.value)
    def __getitem__(self, n):
        return self.value[n]
    def __add__(self, other):
        assert self.rows == other.rows and self.columns == other.columns
        return Matrix([[self[i][j] + other[i][j] for j in range(self.columns)] for i in range(self.rows)])
    def __mul__(self, other):
        if not isinstance(other, Matrix):
            return Matrix([[other * j for j in i] for i in self])
        else:
            assert self.columns == other.rows
            m = self.columns
            return Matrix([[sum([self[i][k]*other[k][j] for k in range(m)]) for j in range(other.columns)]
                           for i in range(self.rows)])
    __rmul__ = __mul__
    def transpose(self):
        return Matrix(zip(*self.value))
    def inv(self):
        assert self.rows == self.columns
        if self.rows == 2:
            a, b, c, d = reduce(lambda a, b: a + b, self)
            f = 1 / (a*d - b*c)
            return Matrix([[d*f, -b*f], [-c*f, a*f]])
        else:
            raise NotImplementedError

def linreg(p):
    x_ = Matrix([[1] * len(p), [i[0] for i in p]])
    x = x_.transpose()
    y = Matrix(zip([i[1] for i in p]))
    return (x_ * x).inv() * (x_ * y)

def lambda_linreg(p):
    b, a = zip(*linreg(p))[0]
    return lambda x: a*x + b