OrdinaryLeastSquares
// This source code is subject to the terms of the Mozilla Public License 2.0 at https://mozilla.org/MPL/2.0/
// © lejmer

//@version=5

// @description One of the most common ways to estimate the coefficients for a linear regression is to use the Ordinary Least Squares (OLS) method.
// This library implements OLS in pine. This implementation can be used to fit a linear regression of multiple independent variables onto one dependent variable,
// as long as the assumptions behind OLS hold.
library("OrdinaryLeastSquares")

// @function Solve a linear system of equations using the Ordinary Least Squares method.
// This function returns both the estimated OLS solution and a matrix that essentially measures the model stability (linear dependence between the columns of 'x').
// NOTE: The latter is an intermediate step when estimating the OLS solution but is useful when calculating the covariance matrix and is returned here to save computation time
// so that this step doesn't have to be calculated again when things like standard errors should be calculated.
// @param x The matrix containing the independent variables. Each column is regarded by the algorithm as one independent variable. The row count of 'x' and 'y' must match.
// @param y The matrix containing the dependent variable. This matrix can only contain one dependent variable and can therefore only contain one column. The row count of 'x' and 'y' must match.
// @returns Returns both the estimated OLS solution and a matrix that essentially measures the model stability (xtx_inv is equal to (X'X)^-1).
export solve_xtx_inv(matrix<float> x, matrix<float> y)  =>
//{
    if (matrix.columns(y) > 1)
        runtime.error("OrdinaryLeastSquares.ols(): Y-matrix must contain only 1 column")
    if (matrix.rows(x) != matrix.rows(y))
        runtime.error("OrdinaryLeastSquares.ols(): Row count of X and Y does not match")
    
    matrix<float> xt = matrix.transpose(x)
    matrix<float> xtx = matrix.mult(xt, x)
    matrix<float> xtx_inv = matrix.pinv(xtx)
	matrix<float> beta_hat = matrix.mult(matrix.mult(xtx_inv, xt), y)
    
    [beta_hat, xtx_inv]
//}

// @function Solve a linear system of equations using the Ordinary Least Squares method.
// @param x The matrix containing the independent variables. Each column is regarded by the algorithm as one independent variable. The row count of 'x' and 'y' must match.
// @param y The matrix containing the dependent variable. This matrix can only contain one dependent variable and can therefore only contain one column. The row count of 'x' and 'y' must match.
// @returns Returns the estimated OLS solution.
export solve(matrix<float> x, matrix<float> y)  =>
//{
    [beta_hat, xtx_inv] = solve_xtx_inv(x, y)
	beta_hat
//}

// @function Solve a linear system of equations using the Ordinary Least Squares method.
// @param x The matrix containing the independent variables. Each column is regarded by the algorithm as one independent variable. The row count of 'x' and 'y' must match.
// @param y The array containing the dependent variable. The row count of 'x' and the size of 'y' must match.
// @returns Returns the estimated OLS solution.
export solve(matrix<float> x, array<float> y)  =>
//{
    if (matrix.rows(x) != array.size(y))
        runtime.error("OrdinaryLeastSquares.ols(): Row count of X and size of Y does not match")

    matrix<float> xt = matrix.transpose(x)
    matrix<float> xtx = matrix.mult(xt, x)
    matrix<float> xtx_inv = matrix.pinv(xtx)
	
    matrix.mult(matrix.mult(xtx_inv, xt), y)
//}

// @function Calculate the standard errors.
// @param x The matrix containing the independent variables. Each column is regarded by the algorithm as one independent variable. The row count of 'x' and 'y' must match.
// @param y The matrix containing the dependent variable. This matrix can only contain one dependent variable and can therefore only contain one column. The row count of 'x' and 'y' must match.
// @param beta_hat The Ordinary Least Squares (OLS) solution provided by solve_xtx_inv() or solve().
// @param xtx_inv This is (X'X)^-1, which means we take the transpose of the X matrix, multiply that the X matrix and then take the inverse of the result.
// This essentially measures the linear dependence between the columns of the X matrix.
// @returns The standard errors.
export standard_errors(matrix<float> x, matrix<float> y, matrix<float> beta_hat, matrix<float> xtx_inv) =>
//{
    if (matrix.columns(y) > 1)
        runtime.error("OrdinaryLeastSquares.standard_errors(): Y-matrix must contain only 1 column")
    if (matrix.rows(x) != matrix.rows(y))
        runtime.error("OrdinaryLeastSquares.standard_errors(): Row count of X and Y does not match")
    
    matrix<float> residuals = matrix.diff(y, matrix.mult(x, beta_hat))
    matrix<float> dev_score_sum_of_squares = matrix.mult(matrix.transpose(residuals), residuals)
    float df = math.max(1, matrix.rows(x) - matrix.columns(x))
    float sigma2 = matrix.get(matrix.mult(dev_score_sum_of_squares, 1 / df), 0, 0)
    matrix<float> covar = matrix.mult(xtx_inv, sigma2)
    
    int size = matrix.rows(covar)
    array<float> errors = array.new<float>(size)
    
    for i = 0 to size - 1
        array.set(errors, i, math.pow(matrix.get(covar, i, i), 0.5))
    
    errors
//}

// @function Estimate the next step of a linear model.
// @param x The matrix containing the independent variables. Each column is regarded by the algorithm as one independent variable. The row count of 'x' and 'y' must match.
// @param beta_hat The Ordinary Least Squares (OLS) solution provided by solve_xtx_inv() or solve().
// @returns Returns the new estimate of Y based on the linear model.
export estimate(matrix<float> x, matrix<float> beta_hat) =>
//{
    float y_hat = 0
    
    for column = 0 to matrix.columns(x) - 1
        y_hat += matrix.get(beta_hat, column, 0) * matrix.get(x, 0, column)
//}



// ----------------------------------------------------------------------------------------------------------------------------
//    Test code for the library. This basically reproduces the result from ta.linreg() with the offset parameter set to 0.
//    Essentially it takes and a constant (a series of 1's) as one independent variable and the bar index (counting from 1) as
//    another independent variable and takes the close price as the dependent variable and produces of estimate of the close
//    price. This is not very useful but works as a demo of how the OLS method can be used to calculate a linear regression
//    in a way that can be tested against something that already exists, to verify that the calculations work as expected.
// ----------------------------------------------------------------------------------------------------------------------------

var int n = input.int(50, "Lookback Window [1 to 4999]", minval = 1, maxval = 4999)

var matrix<float> x = matrix.new<float>(n, 2, 0)
var matrix<float> y = matrix.new<float>(n, 1, 0)

if (barstate.isfirst)
//{
    for row = 0 to matrix.rows(x) - 1
    //{
        matrix.set(x, row, 0, 1)       // X1: Constant
        matrix.set(x, row, 1, row + 1) // X2: Time rank
    //}
//}

for row = 0 to matrix.rows(y) - 1
    matrix.set(y, row, 0, close[row]) // Y: Close price

[beta_hat, xtx_inv] = solve_xtx_inv(x, y) // Solve using solve_xtx_inv() in order to get access to the xtx_inv matrix
float y_hat = estimate(x, beta_hat)
float[] se = standard_errors(x, y, beta_hat, xtx_inv) // Here the xtx_inv matrix is used

color se_color = color.new(color.blue, transp = 80)

u_se = plot(y_hat + array.get(se, 0), "SE (Upper)", color = se_color)
l_se = plot(y_hat - array.get(se, 0), "SE (Lower)", color = se_color)
fill(u_se, l_se, title = "SE Fill Color", color = se_color)

plot(y_hat, color = color.white, linewidth = 2)
