ols python

ols.py

import numpy as np
import scipy.linalg as la

def ols(y, x, const=True):
    n = y.shape[0]

    if const:  # if we want to add a constant to the regression
        X = np.concatenate((np.ones((n, 1)), x), axis=1)
    else:
        X = x

    k = X.shape[1]

    if n < 10000:  # if the no. of obs. is small enough
        Q, R = la.qr(X, mode='economic')
        xxi = la.solve(R.T @ R, np.eye(k))  # this is (X'X)^(-1)
    else:
        xxi = la.solve(X.T @ X, np.eye(k))  # this is (X'X)^(-1)

    beta = xxi @ (X.T @ y)
    u = y - X @ beta
    sigma2 = (u.T @ u) / (n-k)
    se_beta = np.sqrt(sigma2 * np.diag(xxi))

    return beta, se_beta