<commit_before>#Load the Distributions package. Use `Pkg.install("Distributions")` to install first time.
using Distributions: TDist, cdf, ccdf

type regress_results
  coefs
  yhat
  res
  vcv
  tstat
  pval
end

# Keyword arguments are placed after semicolon.
# Symbols start with colon, e.g. `:symbol`.
function ols(y, X; corr=:none, lags=nothing)

    # more stable: β̂ = X \ y, see notes at bottom

    β̂ = X \ y
    ŷ = X * β̂
    μ̂ = y - ŷ

    T, K = size(X)
    σ̂² = dot(μ̂, μ̂) / (T - K)

    #use correction for variance covariance
    if corr == :none
        vcv = σ̂² * (X'*X)\eye(K)

    elseif corr == :white
        newey_west(Xμ̂,0)
    elseif corr == :newey_west
        vcv = newey_west(X, μ̂, lags)
    else
        error("wrong argument for correction keyword")
    end

    # T statistics for H₀: β₀ = 0
    tstat = β̂ ./ sqrt(diag(vcv))

    # absolute value and times two for double sided test
    pval  = 2 * ccdf(TDist(T-K), abs(tstat))

    return regress_results(β̂, ŷ, μ̂, vcv, tstat, pval)
end

function newey_west(X, μ̂)

  T, K = size(X)
  lags = Int(floor(T^(1/4)))
  return  newey_west(X, μ̂, lags)

end

function newey_west(X, μ̂, lags::Integer)

    T,K = size(X)
    XtXInv = (X'*X)\eye(size(K))

    if lags==0
      # In this case get the White estimator
      return XtXInv * X' * diagm(μ̂.^2) * X * XtXInv
    end

    vcv = zeros(K, K)
    for lag in 0:lags
        w = 1 - lag / (lags + 1)
        for t in (lag + 1):T
            # Calculates the off-diagonal terms
            update = w * μ̂[t] * μ̂[t-lag] * (X[t-lag,:]*X[t,:]' + X[t,:]*X[t-lag,:]')
            vcv = vcv + update
        end
    end
    vcv = XtXInv * vcv * XtXInv
    return vcv

end

function gls(y,X,Sigma)

  SigmaInv = Sigma\eye(size(Sigma))
  P=chol(SimgaInv)
  return ols(P*y,P*X)

end

function gmm(y,X,Z;corr=:none,lags=nothing)

  T,Kx = size(X)
  T,Kz = size(Z)

  if corr==:none

    # Generalized 1-step IV estimator
    W = (Z'*Z)\eye(Kz)

  elseif corr==:white | corr==:newey_west

    if corr==:white
      lags=0
    end

    gmm1_res = gmm(y,X,Z;corr=:none)
    μ̂ = gmm1_res.res

    W = zeros(Kz, Kz)
    for lag in 0:lags
        w = 1 - lag / (lags + 1)
        for t in (lag + 1):T
            # Calculates the off-diagonal terms
            update = w * μ̂[t] * μ̂[t-lag] * (Z[t-lag,:]*Z[t,:]' + Z[t,:]*Z[t-lag,:]')
            W = W + update
        end
    end

  else
      error("wrong argument for correction keyword")
  end

  ZtX = Z'*X
  XtZ = X'*Z

  XtZ_W_ZtXInv = (XtZ*W*ZtX)\eye(Kx)
  β̂  = XtZ_W_ZtXInv*(XtZ*W*Z'*y)
  ŷ = X * β̂
  μ̂ = y - ŷ
  σ̂² = dot(μ̂, μ̂) / (T - Kz)
  vcv = σ̂² * XtZ_W_ZtXInv

  # T statistics for H₀: β₀ = 0
  tstat = β̂ ./ sqrt(diag(vcv))

  # absolute value and times two for double sided test
  pval  = 2 * ccdf(TDist(T-K), abs(tstat))
<commit_msg>fix non-corrected vcv + newey_west default lags + formatting<commit_after>#Load the Distributions package. Use `Pkg.install("Distributions")` to install first time.
using Distributions: TDist, ccdf

type regress_results
    coefs
    yhat
    res
    vcv
    tstat
    pval
end

# Keyword arguments are placed after semicolon.
# Symbols start with colon, e.g. `:symbol`.
function ols(y, X; corr=:none, lags=nothing)

    # β̂ = X \ y is more stable than β̂ = inv(X'*X) * X' \ y 
    # see notes at bottom of case 1 notebook
    β̂ = X \ y
    ŷ = X * β̂
    μ̂ = y - ŷ

    T, K = size(X)
    σ̂² = dot(μ̂, μ̂) / (T - K)

    #use correction for variance covariance
    if corr == :none
        vcv = σ̂² * inv(X'*X)

    elseif corr == :white
        newey_west(X, μ̂, 0)
    elseif corr == :newey_west
        vcv = lags == nothing ? newey_west(X, μ̂) : newey_west(X, μ̂, lags)
    else
        error("wrong argument for correction keyword")
    end
    
    # T statistics for H₀: β₀ = 0
    tstat = β̂ ./ sqrt(diag(vcv))

    # absolute value and times two for double sided test
    pval  = 2 * ccdf(TDist(T-K), abs(tstat))

    regress_results(β̂, ŷ, μ̂, vcv, tstat, pval)
end

function newey_west(X, μ̂)

  T, K = size(X)
  lags = Int(floor(T^(1/4)))
  return  newey_west(X, μ̂, lags)

end

function newey_west(X, μ̂, lags::Integer)

    T,K = size(X)
    XtXInv = inv(X'*X)

    if lags==0
      # In this case get the White estimator
      return XtXInv * X' * diagm(μ̂.^2) * X * XtXInv
    end

    vcv = zeros(K, K)
    for lag in 0:lags
        w = 1 - lag / (lags + 1)
        for t in (lag + 1):T
            # Calculates the off-diagonal terms
            update = w * μ̂[t] * μ̂[t-lag] * (X[t-lag,:]*X[t,:]' + X[t,:]*X[t-lag,:]')
            vcv = vcv + update
        end
    end
    vcv = XtXInv * vcv * XtXInv
    return vcv

end

function gls(y, X, Ω)

  P = chol(inv(Ω))
  return ols(P*y, P*X)

end

function gmm(y, X, Z; corr=:none, lags=nothing)

  T, Kx = size(X)
  T, Kz = size(Z)

  if corr==:none

    # Generalized 1-step IV estimator
    W = inv(Z'*Z)

  elseif corr==:white | corr==:newey_west

    if corr==:white
      lags=0
    end

    gmm1_res = gmm(y,X,Z;corr=:none)
    μ̂ = gmm1_res.res

    W = zeros(Kz, Kz)
    for lag in 0:lags
        w = 1 - lag / (lags + 1)
        for t in (lag + 1):T
            # Calculates the off-diagonal terms
            update = w * μ̂[t] * μ̂[t-lag] * (Z[t-lag,:]*Z[t,:]' + Z[t,:]*Z[t-lag,:]')
            W = W + update
        end
    end

  else
      error("wrong argument for correction keyword")
  end

  ZtX = Z'*X
  XtZ = X'*Z

  XtZ_W_ZtXInv = inv(XtZ*W*ZtX)
  β̂  = XtZ_W_ZtXInv*(XtZ*W*Z'*y)
  ŷ = X * β̂
  μ̂ = y - ŷ
  σ̂² = dot(μ̂, μ̂) / (T - Kz)
  vcv = σ̂² * XtZ_W_ZtXInv

  # T statistics for H₀: β₀ = 0
  tstat = β̂ ./ sqrt(diag(vcv))

  # absolute value and times two for double sided test
  pval  = 2 * ccdf(TDist(T-K), abs(tstat))

{'fraction_non_alphanumeric': 0.15353920888272032, 'fraction_numerical': 0.007460097154753644, 'mean_word_length': 2.215281650864473, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '299809', 'n_tokens_mistral': 2660, 'n_tokens_neox': 2552, 'n_words': 664}