regress.py 555 Bytes
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
from sklearn import linear_model
import matplotlib.pyplot as plt


def linear(data_frame, column):
    """
    Execute a simple linear regression on the given column for the passed data_frame.
    :param data_frame:
    :param column:
    :return:
    """
    x = data_frame.index.values.reshape(-1, 1)
    y = data_frame[column].values

    model = linear_model.LinearRegression()
    model.fit(x, y)
    print("R-Score:", model.score(x, y))
    plt.scatter(x, y)
    plt.grid()
    plt.plot(x, model.predict(x), color="red", linewidth=3)
    plt.show()