29 lines
782 B
Python
29 lines
782 B
Python
from cProfile import label
|
|
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
from sklearn.linear_model import LinearRegression
|
|
|
|
|
|
class Bind(object):
|
|
k, b = 0, 0
|
|
|
|
def __init__(self, E, n, m, data):
|
|
self.E = E
|
|
self.n, self.m = n, m
|
|
self.x, self.y = data[:, 0], data[:, 1]
|
|
|
|
def fit(self):
|
|
self.reg = LinearRegression()
|
|
self.reg.fit(self.x, self.y)
|
|
self.k = self.reg.coef_[0][0]
|
|
self.b = self.reg.intercept_[0]
|
|
|
|
def draw(self, title):
|
|
fig = plt.figure(figsize=(8, 8))
|
|
ax = fig.add_subplot(1, 1, 1)
|
|
ax.scatter(self.x, self.y, s=0.1, c='k', label='raw')
|
|
ax.plot(self.x, self.reg.predict(self.x), c='r', label='fit')
|
|
ax.legend()
|
|
fig.savefig(title)
|
|
plt.close()
|