Least squares. Math analysis homework #2

st.py

# -*- encoding: utf-8 -*-

import sys
import os
import matplotlib.pyplot as plt
import numpy as np # instead of math


def getFullPath(fileName, fileType):
    return (os.getcwd() + '\\' + '{}.{}').format(fileName, fileType)


def savePlotAs(fileType='png'):
    order = 0
    fileName = 'plot' + str(order)
    fullFileName = getFullPath(fileName, fileType)
    
    while os.path.exists(fullFileName):
        order += 1
        fileName = 'plot' + str(order)
        fullFileName = getFullPath(fileName, fileType)

    plt.savefig(
        (os.getcwd() + '\\' + '{}.{}').format(fileName, fileType), fmt=fileType
    )


def getLine(points):
    sumx = sum(x for x, y in points)
    sumy = sum(y for x, y in points)
    sumx2 = sum(x * x for x, y in points)
    sumxy = sum(x * y for x, y in points)

    n = len(points)
    d = n * sumx2 - sumx ** 2

    dk = n * sumxy - sumx * sumy
    db = sumx2 * sumy - sumx * sumxy

    return (dk / d, db / d)


def dependence(points, k, b):
    avY = sum(y for (_, y) in points) / len(points)
    r = (
        1.0
        - (sum((y - k*x - b)**2 for (x, y) in points)
        / sum((y - avY)**2 for (_, y) in points))
    ) ** 0.5

    return r


if __name__ == '__main__':
    try:
        f = open(sys.argv[1], 'r', encoding='utf-8')
    except IndexError:
        f = open('input.txt', 'r', encoding='utf-8')

    xmin = 1e10
    xmax = -1e10

    points = []
    for line in f:
        x, y = map(float, line.split())
        points.append((x, y))

        xmin = min(xmin, x)
        xmax = max(xmax, x)

    f.close()

    x = np.linspace(xmin - 1, xmax + 1, 101)

    k, b = getLine(points)
    fx = k * x + b
    fig, ax = plt.subplots()
    ax.plot(x, fx, color='blue' , label=str(k)+' * x + ('+str(b)+')')
    ax.plot(points[0][0], points[0][1], color='white', label='r = ' + str(dependence(points, k, b)))

    for x, y in points:
        plt.scatter(x, y)

    print('dependence: r =', dependence(points, k, b))

    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.legend()

    savePlotAs('png')
    plt.show()