remove cholesky (#252)

* remove cholesky

* speed up

Author: matthieugomez

Project.toml

@@ -1,6 +1,6 @@
 name = "FixedEffectModels"
 uuid = "9d5cd8c9-2029-5cab-9928-427838db53e3"
-version = "1.9.5"
+version = "1.10.0"
 
 [deps]
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
@@ -28,5 +28,5 @@ StatsBase = "0.33, 0.34"
 StatsFuns = "0.9, 1"
 StatsModels = "0.7"
 Tables = "1"
-Vcov = "0.7"
+Vcov = "0.8"
 julia = "1.6"

src/fit.jl

@@ -330,7 +330,7 @@ function StatsAPI.fit(::Type{FixedEffectModel},
         notcollinear_fe_z = collinear_fe[find_cols_z(n_exo, n_endo, n_z)] .== false
         notcollinear_fe_endo_small = notcollinear_fe_endo[basis_endo]
 
-    	basis_all = basis(eachcol(Xexo)..., eachcol(Z)..., eachcol(Xendo)...; has_intercept = has_intercept)
+    	basis_all = basis(Xexo, Z, eachcol(Xendo)...; has_intercept = has_intercept)
         basis_Xexo = basis_all[1:size(Xexo, 2)] .* notcollinear_fe_exo
         basis_Z = basis_all[(size(Xexo, 2) +1):(size(Xexo, 2) + size(Z, 2))] .* notcollinear_fe_z
         basis_endo_small = basis_all[(size(Xexo, 2) + size(Z, 2) + 1):end] .* notcollinear_fe_endo_small
@@ -354,7 +354,7 @@ function StatsAPI.fit(::Type{FixedEffectModel},
             @info "Endogenous vars collinear with ivs. Recategorized as exogenous: $(out)"
 
             # third pass
-            basis_all = basis(eachcol(Xexo)..., eachcol(Z)..., eachcol(Xendo)...; has_intercept = has_intercept)
+            basis_all = basis(Xexo, Z, Xendo; has_intercept = has_intercept)
             basis_Xexo = basis_all[1:size(Xexo, 2)]
             basis_Z = basis_all[(size(Xexo, 2) +1):(size(Xexo, 2) + size(Z, 2))]
         end
@@ -366,22 +366,22 @@ function StatsAPI.fit(::Type{FixedEffectModel},
 
         # Build
         newZ = hcat(Xexo, Z)
-        Pi = ldiv!(cholesky!(Symmetric(newZ'newZ)), newZ'Xendo)
+        Pi = ls_solve(newZ, Xendo)
         Xhat = hcat(Xexo, newZ * Pi)
         X = hcat(Xexo, Xendo)
 
         # prepare residuals used for first stage F statistic
         ## partial out Xendo in place wrt (Xexo, Z)
         Xendo_res = BLAS.gemm!('N', 'N', -1.0, newZ, Pi, 1.0, Xendo)
         ## partial out Z in place wrt Xexo
-        Pi2 = ldiv!(cholesky!(Symmetric(Xexo'Xexo)), Xexo'Z)
+        Pi2 = ls_solve(Xexo, Z)
         Z_res = BLAS.gemm!('N', 'N', -1.0, Xexo, Pi2, 1.0, Z)
     else
         # get linearly independent columns
         n_exo = size(Xexo, 2)
         perm = 1:n_exo
         notcollinear_fe_exo = collinear_fe[find_cols_exo(n_exo)] .== false
-        basis_Xexo = basis(eachcol(Xexo)...; has_intercept = has_intercept) .* notcollinear_fe_exo
+        basis_Xexo = basis(Xexo; has_intercept = has_intercept) .* notcollinear_fe_exo
         Xexo = getcols(Xexo, basis_Xexo)
         Xhat = Xexo
         X = Xexo
@@ -393,9 +393,11 @@ function StatsAPI.fit(::Type{FixedEffectModel},
     ## Do the regression
     ##
     ##############################################################################
-
-    crossx = cholesky!(Symmetric(Xhat'Xhat))
-    coef = ldiv!(crossx, Xhat'y)
+    crossx = Xhat' * Xhat
+    Xy = Symmetric(hvcat(2, crossx, Xhat'y, zeros(size(Xhat, 2))', [0.0]))
+    invsym!(Xy; diagonal = 1:size(Xhat, 2))
+    invcrossx = Symmetric(.- Xy[1:(end-1),1:(end-1)])
+    coef = Xy[1:(end-1),end]
 
     ##############################################################################
     ##
@@ -453,9 +455,8 @@ function StatsAPI.fit(::Type{FixedEffectModel},
     end
 
     # Compute standard error
-    vcov_data = Vcov.VcovData(Xhat, crossx, residuals, nobs - size(X, 2) - dof_fes)
+    vcov_data = Vcov.VcovData(Xhat, crossx, invcrossx, residuals, nobs - size(X, 2) - dof_fes)
     matrix_vcov = StatsAPI.vcov(vcov_data, vcov_method)
-
     # Compute Fstat
     F = Fstat(coef, matrix_vcov, has_intercept)
     # dof_ is the number of estimated coefficients beyond the intercept.
@@ -465,15 +466,15 @@ function StatsAPI.fit(::Type{FixedEffectModel},
     # Compute Fstat of First Stage
     if has_iv && first_stage
         Pip = Pi[(size(Pi, 1) - size(Z_res, 2) + 1):end, :]
-        try 
+        #try 
             r_kp = Vcov.ranktest!(Xendo_res, Z_res, Pip,
                               vcov_method, size(X, 2), dof_fes)
             p_kp = chisqccdf(size(Z_res, 2) - size(Xendo_res, 2) + 1, r_kp)
             F_kp = r_kp / size(Z_res, 2)
-        catch
-            @info "ranktest failed ; first-stage statistics not estimated"
-            p_kp, F_kp = NaN, NaN
-        end
+       # catch
+       #     @info "ranktest failed ; first-stage statistics not estimated"
+       #     p_kp, F_kp = NaN, NaN
+       # end
     end
 
     # Compute rss, tss

src/utils/basecol.jl

@@ -1,64 +1,65 @@
-##############################################################################
-##
-## Returns base of [A B C ...]
-## Important: it must be the case that it returns the in order, that is [A B A] returns [true true false] not [false true true]
-## 
-##
-##############################################################################
-function basis(@nospecialize(xs::AbstractVector...); has_intercept = false)
-    invXX = invsym!(crossprod(collect(xs)); has_intercept = has_intercept)
-    return diag(invXX) .> 0
-end
-
-function crossprod(xs::Vector{<:AbstractVector})
-    XX = zeros(length(xs), length(xs))
-    for i in 1:length(xs)
-        for j in 1:i
-            XX[i, j] = xs[i]' * xs[j]
-        end  
-    end
+# create the matrix X'X
+function crossprod(xs...)
+    idx = vcat(0, cumsum(size(x, 2) for x in xs))
+    XX = zeros(idx[end], idx[end])
     for i in 1:length(xs)
-        for j in (i+1):length(xs)
-            XX[i, j] = XX[j, i]
+        for j in i:length(xs)
+            XX[(idx[i]+1):idx[i+1], (idx[j]+1):idx[j+1]] .= xs[i]' * xs[j]
         end  
     end
-    return XX
+    return Symmetric(XX, :U)
 end
 
 # generalized 2inverse
-function invsym!(X::AbstractMatrix; has_intercept = false)
-    # Options from SAS
-    # The C value adjusts the check to the relative scale of the variable. 
-    # The C value is equal to the corrected sum of squares for the variable, unless the corrected sum of squares is 0, in which case C is 1. 
-    # If you specify the NOINT option but not the ABSORB statement, PROC GLM uses the uncorrected sum of squares instead. 
-    # The default value of the SINGULAR= option, 107, might be too small, but this value is necessary in order to handle the high-degree polynomials used in the literature to compare regression routines
+#actually return minus the symmetric
+function invsym!(X::Symmetric; has_intercept = false, setzeros = false, diagonal = 1:size(X, 2))
     tols = max.(diag(X), 1)
-    for j in 1:size(X, 1)
+    buffer = zeros(size(X, 1))
+    for j in diagonal
         d = X[j,j]
-        if abs(d) < tols[j] * sqrt(eps())
-            X[j,:] .= 0
-            X[:,j] .= 0
+        if setzeros && abs(d) < tols[j] * sqrt(eps())
+            X.data[1:j,j] .= 0
+            X.data[j,(j+1):end] .= 0
         else
-            X[j,:] = X[j,:] ./ d
-            for i in 1:size(X, 1)
-                if i != j
-                    X[i,:] .= X[i,:] .- X[i,j] .* X[j,:]
-                    X[i,j] = -X[i,j] / d
-                end
-            end
-            X[j,j] = 1 / d
+            # used to mimic SAS; now similar to SweepOperators
+            copy!(buffer, view(X, :, j))
+            Symmetric(BLAS.syrk!('U', 'N', -1/d, buffer, one(eltype(X)), X.data))
+            rmul!(buffer, 1 / d)
+            @views copy!(X.data[1:j-1,j], buffer[1:j-1])        
+            @views copy!(X.data[j, j+1:end], buffer[j+1:end])   
+            X[j,j] = - 1 / d
         end
-        if has_intercept && j == 1
+        if setzeros && has_intercept && j == 1
             tols = max.(diag(X), 1)
         end
     end
     return X
 end
 
+## Returns base of [A B C ...]
+## Important: it must be the case that it returns the in order, that is [A B A] returns [true true false] not [false true true]
+function basis(xs...; has_intercept = false)
+    invXX = invsym!(crossprod(xs...); has_intercept = has_intercept, setzeros = true)
+    return diag(invXX) .< 0
+end
+
 function getcols(X::AbstractMatrix,  basecolX::AbstractVector)
     sum(basecolX) == size(X, 2) ? X : X[:, basecolX]
 end
 
+
+function ls_solve(X, y::AbstractVector)
+    Xy = crossprod(X, y)
+    invsym!(Xy, diagonal = 1:size(X, 2))
+    return Xy[1:size(X, 2),end]
+end
+
+function ls_solve(X, Y::AbstractMatrix)
+    XY = crossprod(X, Y)
+    invsym!(XY, diagonal = 1:size(X, 2))
+    return XY[1:size(X, 2),(end-size(Y, 2)+1):end]
+end
+
 ##############################################################################
 # Auxiliary functions to find columns of exogeneous, endogenous and IV variables
 ##############################################################################

test/fit.jl

@@ -1,6 +1,6 @@
-using CUDA, Metal, FixedEffectModels, CategoricalArrays, CSV, DataFrames, Test, LinearAlgebra
+using FixedEffectModels, CategoricalArrays, CSV, DataFrames, Test, LinearAlgebra
 using FixedEffectModels: nullloglikelihood_within
-
+using CUDA, Metal
 
 ##############################################################################
 ##
@@ -723,7 +723,6 @@ end
 
 @testset "missings" begin
 
-	# Missing
 	df1 = DataFrame(a=[1.0, 2.0, 3.0, 4.0], b=[5.0, 7.0, 11.0, 13.0])
 	df2 = DataFrame(a=[1.0, missing, 3.0, 4.0], b=[5.0, 7.0, 11.0, 13.0])
 	x = reg(df1, @formula(a ~ b))
@@ -742,7 +741,3 @@ end
 	x = reg(df1, @formula(a ~ b + fe(c)))
 	@test coef(x) ≈ [0.5] atol = 1e-4
 end
-
-
-
-