Simplify logic about collinearity and fixed effects (#254)

matthieugomez · web-flow · commit 8f11f6e3def6 · 2023-12-01T09:07:34.000-06:00
* remove cholesky

* speed up

* update

* rmv find_cols

* patched version

* Update Project.toml

* Update fit.jl

* remove crossprod

* Update fit.jl

* simplify logic about fixedeffects collinerarity

* rmv getcols

* Update fit.jl

* Update fit.jl

* correct collinearity in endogeneous
diff --git a/benchmark/benchmark.jl b/benchmark/benchmark.jl
@@ -12,19 +12,21 @@ y= 3 .* x1 .+ 5 .* x2 .+ cos.(id1) .+ cos.(id2).^2 .+ randn(N)
 df = DataFrame(id1 = id1, id2 = id2, x1 = x1, x2 = x2, y = y)
 # first time
 @time reg(df, @formula(y ~ x1 + x2))
-# 3.5s
+# 1.823s
 @time reg(df, @formula(y ~ x1 + x2))
-# 0.497374 seconds (450 allocations: 691.441 MiB, 33.18% gc time)
+# 0.353469 seconds (441 allocations: 691.439 MiB, 3.65% gc time)
 @time reg(df, @formula(y ~ x1 + x2),  Vcov.cluster(:id2))
-# 1.898018 seconds (7.10 M allocations: 1.220 GiB, 8.20% gc time, 4.46% compilation time)
+# 0.763999 seconds (2.96 M allocations: 967.418 MiB, 2.29% gc time, 54.39% compilation time: 5% of which was recompilation)
 @time reg(df, @formula(y ~ x1 + x2),  Vcov.cluster(:id2))
-# 0.605172 seconds (591 allocations: 768.939 MiB, 42.38% gc time)
+# 0.401544 seconds (622 allocations: 768.943 MiB, 3.52% gc time)
 @time reg(df, @formula(y ~ x1 + x2 + fe(id1)))
 # 0.893835 seconds (1.03 k allocations: 929.130 MiB, 54.19% gc time)
+@time reg(df, @formula(y ~ x1 + x2 + fe(id1)))
+#   0.474160 seconds (1.13 k allocations: 933.340 MiB, 1.74% gc time)
 @time reg(df, @formula(y ~ x1 + x2 + fe(id1)), Vcov.cluster(:id1))
-# 1.015078 seconds (1.18 k allocations: 1008.532 MiB, 56.50% gc time)
+#   0.598816 seconds (261.08 k allocations: 1.007 GiB, 8.29% gc time, 9.21% compilation time)
 @time reg(df, @formula(y ~ x1 + x2 + fe(id1) + fe(id2)))
-# 1.835464 seconds (4.02 k allocations: 1.057 GiB, 35.59% gc time)
+# 1.584573 seconds (489.64 k allocations: 1.094 GiB, 2.10% gc time, 8.53% compilation time)
 
 # More complicated setup
 N = 800000 # number of observations
@@ -38,6 +40,10 @@ y= 3 .* x1 .+ 5 .* x2 .+ cos.(id1) .+ cos.(id2).^2 .+ randn(N)
 df = DataFrame(id1 = id1, id2 = id2, x1 = x1, x2 = x2, y = y)
 @time reg(df, @formula(y ~ x1 + x2 + fe(id1) + fe(id2)))
 # 2.504294 seconds (75.83 k allocations: 95.525 MiB, 0.23% gc time)
+# for some reason in 1.10 I now get worse time (iter 200)
+#  4.709078 seconds (108.98 k allocations: 101.417 MiB)
+
+
 
 
 +# fixest
diff --git a/src/fit.jl b/src/fit.jl
@@ -215,7 +215,7 @@ function StatsAPI.fit(::Type{FixedEffectModel},
     end
 
     # collect all variable names (outcome, exo [, endo, iv])
-    var_names_all = [response_name; coef_names]
+    var_names_all = vcat(response_name, coef_names)
 
     if has_iv
         subdf = Tables.columntable((; (x => disallowmissing(view(df[!, x], esample)) for x in endo_vars)...))
@@ -281,7 +281,9 @@ function StatsAPI.fit(::Type{FixedEffectModel},
             if i == 1
                 @info "Dependent variable $(var_names_all[1]) is probably perfectly explained by fixed effects (tol = $collinear_tol)."
             else
-                @info "RHS-variable $(var_names_all[i]) is probably collinear with the fixed effects (tol = $collinear_tol)."
+                @info "RHS-variable $(var_names_all[i]) is collinear with the fixed effects (tol = $collinear_tol)."
+                # set to zero so that removed when taking basis
+                cols[i] .= 0.0
             end
         end
 
@@ -305,86 +307,114 @@ function StatsAPI.fit(::Type{FixedEffectModel},
         end
     end
     
-
-
-
     ##############################################################################
     ##
     ## Get Linearly Independent Components of Matrix
     ##
     ##############################################################################
     # Compute linearly independent columns + create the Xhat matrix
     if has_iv    	
-        n_exo = size(Xexo, 2)
-        n_endo = size(Xendo, 2)
-        n_z = size(Z, 2)
-        perm = 1:(n_exo + n_endo)
-
-        # first pass: remove colinear variables in Xendo
-        notcollinear_fe_endo = collinear_fe[(n_exo+2):(n_exo+n_endo+1)] .== false
-    	basis_endo = basis(Xendo; has_intercept = false) .* notcollinear_fe_endo
-    	Xendo = getcols(Xendo, basis_endo)
-
-    	# second pass: remove colinear variable in Xexo, Z, and Xendo
-        notcollinear_fe_exo = collinear_fe[2:(n_exo+1)] .== false
-        notcollinear_fe_z = collinear_fe[(n_exo+n_endo+2):(n_exo+n_endo+n_z+1)] .== false
-        notcollinear_fe_endo_small = notcollinear_fe_endo[basis_endo]
-
-
-    	basis_all = basis(Xexo, Z, 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
+        perm = 1:(size(Xexo, 2) +size(Xendo, 2))
+
+        # first pass: remove collinear variables in Xendo
+        XendoXendo = Xendo' * Xendo
+    	basis_endo = basis!(Symmetric(deepcopy(XendoXendo)); has_intercept = false)
+        if !all(basis_endo)
+        	Xendo = Xendo[:, basis_endo]
+            XendoXendo = XendoXendo[basis_endo, basis_endo]
+        end
 
+    	# second pass: remove collinear variable in Xexo, Z, and Xendo
+        XexoXexo = Xexo'Xexo
+        XexoZ = Xexo'Z
+        XexoXendo = Xexo'Xendo
+        ZZ = Z'Z
+        ZXendo = Z'Xendo
+        XexoZXendo = Symmetric(hvcat(3, XexoXexo, XexoZ, XexoXendo, 
+                           zeros(size(Z, 2), size(Xexo, 2)), ZZ, ZXendo, 
+                           zeros(size(Xendo, 2), size(Xexo, 2)), zeros(size(Xendo, 2), size(Z, 2)), XendoXendo))
+    	basis_all = basis!(XexoZXendo; has_intercept = has_intercept)
+        basis_Xexo, basis_Z, basis_endo_small = basis_all[1:size(Xexo, 2)], basis_all[(size(Xexo, 2) +1):(size(Xexo, 2) + size(Z, 2))], basis_all[(size(Xexo, 2) + size(Z, 2) + 1):end]
+        # basis_endo_small has same length as number of basis_endo who are true
         if !all(basis_endo_small)
             # if adding Xexo and Z makes Xendo collinear, consider these variables are exogeneous
-            Xexo = hcat(Xexo, getcols(Xendo, .!basis_endo_small))
-            Xendo = getcols(Xendo, basis_endo_small)
+            Xexo = hcat(Xexo, Xendo[:, .!basis_endo_small])
+            Xendo = Xendo[:, basis_endo_small]
+            XexoXexo = Xexo'Xexo
+            XexoZ = Xexo'Z
+            XexoXendo = Xexo'Xendo
+            ZXendo = Z'Xendo
+            XendoXendo = Xendo'Xendo
 
             # out returns false for endo collinear with instruments
             basis_endo2 = trues(length(basis_endo))
             basis_endo2[basis_endo] = basis_endo_small
-
-            # Change coef_names and oldX
             # TODO: I should probably also change formula in this case so that predict still works 
             ans = 1:length(basis_endo)
             ans = vcat(ans[.!basis_endo2], ans[basis_endo2])
             perm = vcat(1:length(basis_Xexo), length(basis_Xexo) .+ ans)
-
+            # there are basis_endo - basis_endo_small in endo
             out = join(coefendo_names[.!basis_endo2], " ")
             @info "Endogenous vars collinear with ivs. Recategorized as exogenous: $(out)"
-                                    
+            
             # third pass
-            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))]
+            XexoZXendo = Symmetric(hvcat(3, XexoXexo, XexoZ, XexoXendo, 
+                               zeros(size(Z, 2), size(Xexo, 2)), ZZ, ZXendo, 
+                               zeros(size(Xendo, 2), size(Xexo, 2)), zeros(size(Xendo, 2), size(Z, 2)), XendoXendo))
+            basis_all = basis!(XexoZXendo; has_intercept = has_intercept)
+            basis_Xexo, basis_Z, basis_endo_small2 = basis_all[1:size(Xexo, 2)], basis_all[(size(Xexo, 2) +1):(size(Xexo, 2) + size(Z, 2))], basis_all[(size(Xexo, 2) + size(Z, 2) + 1):end]
         end
-
-    	Xexo = getcols(Xexo, basis_Xexo)
-    	Z = getcols(Z, basis_Z)
-        size(Z, 2) >= size(Xendo, 2) || throw("Model not identified. There must be at least as many ivs as endogeneous variables")
-        basis_coef = vcat(basis_Xexo, basis_endo[basis_endo_small])
+        if !all(basis_Xexo)
+        	Xexo = Xexo[:, basis_Xexo]
+            XexoXexo = XexoXexo[basis_Xexo, basis_Xexo]
+            XexoXendo = XexoXendo[basis_Xexo, :]
+        end
+        if !all(basis_Z)
+        	Z = Z[:, basis_Z]
+            ZZ = ZZ[basis_Z, basis_Z]
+            ZXendo = ZXendo[basis_Z, :]
+        end
+        XexoZ = XexoZ[basis_Xexo, basis_Z]
+        size(ZXendo, 1) >= size(ZXendo, 2) || throw("Model not identified. There must be at least as many ivs as endogeneous variables")
+        # basis_endo is true for stuff non colinear
+        # I need to have same vector but removeing the true that have been reclassified as exo and replace them by nothing.  so i need to create a vector equal to false if non endo and non basis_endo_small, which is basis_endo2
+        basis_endo2 = trues(length(basis_endo))
+        basis_endo2[basis_endo] = basis_endo_small
+        basis_coef = vcat(basis_Xexo, basis_endo[basis_endo2])
 
         # Build
         newZ = hcat(Xexo, Z)
-        Pi = ls_solve(newZ, Xendo)
-        Xhat = hcat(Xexo, newZ * Pi)
+        # now create Pi = newZ \ Xendo
+        newZnewZ = hvcat(2,  XexoXexo, XexoZ, 
+                             XexoZ', ZZ)
+        newZXendo = vcat(XexoXendo, ZXendo)
+        Pi = ls_solve!(Symmetric(hvcat(2, newZnewZ, newZXendo,
+                                zeros(size(newZXendo')), zeros(size(Xendo, 2), size(Xendo, 2)))), 
+                       size(newZnewZ, 2))
+        newnewZ = newZ * Pi
+        Xhat = hcat(Xexo, newnewZ)
+        XhatXhat = Symmetric(hvcat(2,  XexoXexo, Xexo'newnewZ, 
+                           zeros(size(newnewZ, 2), size(Xexo, 2)), newnewZ'newnewZ))
         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 = ls_solve(Xexo, Z)
+        # Now create Pi2 = Xexo \ Z
+        Pi2 = ls_solve!(Symmetric(hvcat(2, XexoXexo, XexoZ,
+                                zeros(size(Z, 2), size(Xexo, 2)), ZZ)), size(Xexo, 2))
         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[2:(n_exo+1)] .== false
-        basis_Xexo = basis(Xexo; has_intercept = has_intercept) .* notcollinear_fe_exo
-        Xexo = getcols(Xexo, basis_Xexo)
+        perm = 1:size(Xexo, 2)
+        XexoXexo = Xexo'Xexo
+        basis_Xexo = basis!(Symmetric(deepcopy(XexoXexo)); has_intercept = has_intercept)
+        if !all(basis_Xexo)
+            Xexo = Xexo[:, basis_Xexo]
+            XexoXexo = XexoXexo[basis_Xexo, basis_Xexo]
+        end
         Xhat = Xexo
+        XhatXhat = Symmetric(XexoXexo)
         X = Xexo
         basis_coef = basis_Xexo
     end
@@ -394,11 +424,10 @@ function StatsAPI.fit(::Type{FixedEffectModel},
     ## Do the regression
     ##
     ##############################################################################
-    
-    crossx = Xhat' * Xhat
-    Xy = Symmetric(hvcat(2, crossx, Xhat'y, zeros(size(Xhat, 2))', [0.0]))
+    Xy = Symmetric(hvcat(2, XhatXhat, Xhat'y, 
+                         zeros(size(Xhat, 2))', [0.0]))
     invsym!(Xy; diagonal = 1:size(Xhat, 2))
-    invcrossx = Symmetric(.- Xy[1:(end-1),1:(end-1)])
+    invXhatXhat = Symmetric(.- Xy[1:(end-1),1:(end-1)])
     coef = Xy[1:(end-1),end]
 
     ##############################################################################
@@ -419,7 +448,10 @@ function StatsAPI.fit(::Type{FixedEffectModel},
 
     augmentdf = DataFrame()
     if save_fe
-        oldX = getcols(oldX[:, perm], basis_coef)
+        oldX = oldX[:, perm]
+        if !all(basis_coef)
+            oldX = oldX[:, basis_coef]
+        end
         newfes, b, c = solve_coefficients!(oldy - oldX * coef, feM; tol = tol, maxiter = maxiter)
         for fekey in fekeys
             augmentdf[!, fekey] = df[:, fekey]
@@ -457,7 +489,7 @@ function StatsAPI.fit(::Type{FixedEffectModel},
     end
 
     # Compute standard error
-    vcov_data = Vcov.VcovData(Xhat, crossx, invcrossx, residuals, nobs - size(X, 2) - dof_fes)
+    vcov_data = Vcov.VcovData(Xhat, XhatXhat, invXhatXhat, residuals, nobs - size(X, 2) - dof_fes)
     matrix_vcov = StatsAPI.vcov(vcov_data, vcov_method)
     # Compute Fstat
     F = Fstat(coef, matrix_vcov, has_intercept)
diff --git a/src/utils/basecol.jl b/src/utils/basecol.jl
@@ -1,17 +1,3 @@
-# create the matrix X'X
-crossprod(x1) = Symmetric(x1'x1)
-
-function crossprod(x1, x2)
-    Symmetric(hvcat(2, x1'x1, x1'x2, 
-                       zeros(size(x2, 2), size(x1, 2)), x2'x2))
-end
-
-function crossprod(x1, x2, x3)
-    Symmetric(hvcat(3, x1'x1, x1'x2, x1'x3, 
-                       zeros(size(x2, 2), size(x1, 2)), x2'x2, x2'x3, 
-                       zeros(size(x3, 2), size(x1, 2)), zeros(size(x3, 2), size(x2, 2)), x3'x3))
-end
-
 # generalized 2inverse
 #actually return minus the symmetric
 function invsym!(X::Symmetric; has_intercept = false, setzeros = false, diagonal = 1:size(X, 2))
@@ -38,25 +24,23 @@ function invsym!(X::Symmetric; has_intercept = false, setzeros = false, diagonal
     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)
+## Returns base of X = [A B C ...]. Takes as input the matrix X'X (actuallyjust its right upper-triangular)
+## Important: it must be the case that colinear are first columbs in the bsae in the order of columns
+## that is [A B A] returns [true true false] not [false true true]
+function basis!(XX::Symmetric; has_intercept = false)
+    invXX = invsym!(XX; 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]
+#solve X \ y. Take as input the matrix [X'X, X'y
+#                                        y'X, y'y]
+# (but only upper matters)
+function ls_solve!(Xy::Symmetric, nx)
+    if nx > 0
+        invsym!(Xy, diagonal = 1:nx)
+        return Xy[1:nx, (nx+1):end]
+    else
+        return zeros(Float64, 0, size(Xy, 2) - nx)
+    end
 end
diff --git a/test/fit.jl b/test/fit.jl
@@ -250,6 +250,10 @@ using CUDA, Metal
 	x = reg(df, m)
 	@test iszero(coef(x)[2]) || iszero(coef(x)[3])
 
+	m = @formula Sales ~  (Price + Price2 ~ Pimin + NDI)
+	x = reg(df, m)
+	@test iszero(coef(x)[2]) || iszero(coef(x)[3])
+
 
 	## endogeneous variables collinear with instruments are reclassified
 	df.zPimin = df.Pimin
