Ordinary least squares implemented with numpy.

basic_ols_numpy_example.py

import numpy as np
import sys

# lines = input.split("\n")
lines = sys.stdin.readlines()

train_header = lines[0].split()
n_train_features, n_train_observations = int(train_header[0]), int(train_header[1])
training = np.array([row.split() for row in lines[1:n_train_observations]], dtype = float)
X_train, y_train = training[:, 0:n_train_features], training[:, n_train_features:]

test_header = lines[n_train_observations + 1].split()
n_test_observations = int(test_header[0])
X_test = np.array([
    row.split() for row in 
    lines[n_train_observations + 2:n_train_observations + 2 + n_test_observations] if len(row)
], dtype = float)

class OLS:
    
    coef_ = None
    intercept = None
    fit_intercept = False
    
    def __init__(self, fit_intercept = False):        
        if fit_intercept:
            self.fit_intercept = fit_intercept
    
    def add_intercept(self, X):
        intercept = np.ones([X.shape[0], 1])
        return np.concatenate([intercept, X], axis=1)
    
    def fit(self, X, y):
        
        if self.fit_intercept:
            X = self.add_intercept(X)
            
        xt_x = np.dot(X.T, X)
        xt_x_inv = np.linalg.inv(xt_x)
        xt_x_inv_xt = np.dot(xt_x_inv, X.T)
        self.coef_ = np.dot(xt_x_inv_xt, y)
        
    def predict(self, X):
        if self.fit_intercept: 
            X = self.add_intercept(X)
        return np.dot(X, self.coef_)
    

ols = OLS(fit_intercept = True)
ols.fit(X_train, y_train)
y_hat = ols.predict(X_test)
for prediction in y_hat:
    print(prediction[0])