头歌平台线性回归第三关

#encoding=utf8 
import numpy as np
#mse
def mse_score(y_predict,y_test):
    mse = np.mean((y_predict-y_test)**2)
    return mse
#r2
def r2_score(y_predict,y_test):
    '''
    input:y_predict(ndarray):预测值
          y_test(ndarray):真实值
    output:r2(float):r2值
    '''
    #********* Begin *********#
    y_mean = np.mean(y_test)
    r2 = 1 - np.sum((y_predict - y_test)**2) / np.sum((y_mean - y_test)**2)
    #********* End *********#
    return r2
class LinearRegression :
    def __init__(self):
        '''初始化线性回归模型'''
        self.theta = None
    def fit_normal(self,train_data,train_label):
        '''
        input:train_data(ndarray):训练样本
              train_label(ndarray):训练标签
        '''
        #********* Begin *********#
        #(xt x)-1xty
        # 先给训练数据加上一列1
        x = np.hstack([train_data, np.ones((len(train_data), 1))])
        self.theta = np.linalg.inv(x.T.dot(x)).dot(x.T).dot(train_label)
        #********* End *********#
        return self.theta
    def predict(self,test_data):
        '''
        input:test_data(ndarray):测试样本
        '''
        #********* Begin *********#
        x = np.hstack([test_data, np.ones((len(test_data), 1))])
        return x.dot(self.theta)
        #********* End *********#