polynomial_fitting.py

from fractions import Fraction
from decimal import Decimal

class Matrix:
    def __init__(self, width, height):
        self.width = width
        self.height = height
        self.rows = [[0]*width for y in range(height)]
    def set(self, x, y, value):
        self.rows[y][x] = value
    def get(self, x, y):
        return self.rows[y][x]

    def __repr__(self):
        return "\n".join(" ".join(str(item) for item in row) for row in self.rows)

    def swap(self, row_a, row_b):
        self.rows[row_a], self.rows[row_b] = self.rows[row_b], self.rows[row_a]
    def mult(self, row, amt):
        self.rows[row] = [item * amt for item in self.rows[row]]
    #multiplies the row so that its `col`th value is 1
    def normalize(self, row, col):
        value = self.get(col, row)
        assert value != 0, "expected nonzero value in row {} column {}, got {}".format(row, col, value)
        self.mult(row, 1/value)
    def add(self, target, dest):
        assert target != dest
        for i in range(self.width):
            self.rows[dest][i] += self.rows[target][i]
    def sub(self, target, dest):
        self.mult(target, -1)
        self.add(target, dest)
        self.mult(target, -1)

    #subtracts target from dest until dest`s `col`th value is 0.
    #side effects: may multiply target.
    def reduce(self, target, dest, col):
        if self.get(col,dest) == 0: return
        assert target != dest
        assert self.get(col, target) != 0, "{} {} {}".format(target, dest, col)
        self.normalize(target, col)
        self.mult(target, self.get(col, dest))
        self.sub(target, dest)

    """manipulates the matrix so every row is all zeroes, except the rightmost column and one other column, which is 1"""
    def decompose(self):
        #triangulation step- the lower right half of the matrix must be all zeroes.
        for target in range(self.height):
            x = self.height - target - 1
            for dest in range(target+1, self.height):
                self.reduce(target, dest, x)
            self.normalize(target, x)

        #complete reduction step - iterate the rows backwards, each one being all zeroes but one, using it to eliminate that term from all rows above it.
        for target in range(self.height-1, -1, -1):
            x = self.height - target - 1
            for dest in range(target-1, -1, -1):
                self.reduce(target, dest, x)
            self.normalize(target, x)
        

def matrix_from_lists(data):
    height = len(data)
    width = len(data[0])
    m = Matrix(width, height)
    for x in range(width):
        for y in range(height):
            m.set(x,y, Fraction(data[y][x]))
    return m

def matrix_from_dict(f):
    terms = len(f)
    #construct equations
    rows = []
    for x, y in f.items():
        coeffs = [x**n for n in range(terms-1, -1, -1)]
        coeffs.append(y)
        rows.append(coeffs)
    #construct matrix
    return matrix_from_lists(rows)
    

def poly_fit(f):
    m = matrix_from_dict(f)
    m.decompose()
    return m

def printed_fit(f):
    def _format(idx, coefficient):
        if coefficient < 0:
            return "-" + _format(idx, -coefficient)
        if idx == 0:
            return str(coefficient)
        elif coefficient == 1:
            if idx == 1:
                return "x"
            else:
                return "(x^{})".format(idx)
        else:
            if idx == 1:
                return "{}*x".format(coefficient)
            else:
                return "{}*(x^{})".format(coefficient, idx)
    m = poly_fit(f)
    coefficients = [m.get(m.width-1, row) for row in range(m.height)]
    terms = [_format(idx, coefficient) for idx, coefficient in enumerate(coefficients) if coefficient != 0]
    terms.reverse()
    return " + ".join(terms).replace("+ -", "-")

#like `printed_fit`, but uses special formatting that makes the answer good on systems with markup.
def markup_fit(f):
    def _format(idx, coefficient):
        if coefficient < 0:
            return "-" + _format(idx, -coefficient)
        if coefficient.denominator != 1:
            coefficient = "\\frac{{{}}}{{{}}}".format(coefficient.numerator, coefficient.denominator)
        if idx == 0:
            return str(coefficient)
        elif coefficient == 1:
            if idx == 1:
                return "x"
            else:
                return "(x^{{{}}})".format(idx)
        else:
            if idx == 1:
                return "{}*x".format(coefficient)
            else:
                return "{}*x^{{{}}}".format(coefficient, idx)
    m = poly_fit(f)
    coefficients = [m.get(m.width-1, row) for row in range(m.height)]
    terms = [_format(idx, coefficient) for idx, coefficient in enumerate(coefficients) if coefficient != 0]
    terms.reverse()
    return "$" + (" + ".join(terms).replace("+ -", "-")) + "$"

def functional_fit(f):
    m = poly_fit(f)
    coefficients = [m.get(m.width-1, row) for row in range(m.height)]
    return lambda x: sum(c * (x**n) for n, c in enumerate(coefficients))

f = [5, 13, 24, 38, 55]
f = dict(enumerate(f))

print(printed_fit(f))
print(markup_fit(f))
func = functional_fit(f)
for n in range(len(f)+5):
    print(n, func(Fraction(n)))