Validate fixed-effect slope interactions and reduce latency

CHANGELOG.md

@@ -0,0 +1,7 @@
+# Changelog
+
+## 2.0.0
+
+### Breaking changes
+
+- A standalone RHS slope that is exactly spanned by a continuous-slope fixed effect is now omitted from coefficient output. For example, `fe(id)*x` still includes group intercepts and group-specific slopes, but no longer reports the unidentified common `x` coefficient as a dropped `0` with `NaN` inference statistics. Code that relies on coefficient positions or names should account for the shorter coefficient vector.

Project.toml

@@ -1,6 +1,6 @@
 name = "FixedEffectModels"
 uuid = "9d5cd8c9-2029-5cab-9928-427838db53e3"
-version = "1.13.3"
+version = "2.0.0"
 
 [deps]
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
@@ -36,9 +36,10 @@ julia = "1.10"
 CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597"
 CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["CategoricalArrays", "CSV", "CUDA", "Metal", "Random", "Test"]
+test = ["CategoricalArrays", "CSV", "CUDA", "Logging", "Metal", "Random", "Test"]

README.md

@@ -40,7 +40,7 @@ reg(df, @formula(Sales ~ NDI + fe(State) + fe(Year)), Vcov.cluster(:State), weig
 
 	High-dimensional fixed effect variables are indicated with the function `fe`.  You can add an arbitrary number of high dimensional fixed effects, separated with `+`. You can also interact fixed effects using `&` or `*`.
 
-	For instance, to add state fixed effects use `fe(State)`. To add both state and year fixed effects, use `fe(State) + fe(Year)`. To add state-year fixed effects, use `fe(State)&fe(Year)`. To add state specific slopes for year, use `fe(State)&Year`. To add both state fixed-effects and state specific slopes for year use `fe(State)*Year`.
+	For instance, to add state fixed effects use `fe(State)`. To add both state and year fixed effects, use `fe(State) + fe(Year)`. To add state-year fixed effects, use `fe(State)&fe(Year)`. To add state specific slopes for year, use `fe(State)&Year`. To add both state fixed-effects and state specific slopes for year use `fe(State)*Year`. Since the state-specific slopes span the common year slope, `Year` is omitted from the reported coefficients in the last specification.
 
 	```julia
 	reg(df, @formula(Sales ~ Price + fe(State) + fe(Year)))

benchmark/benchmark.jl

@@ -40,11 +40,16 @@ 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)))
 #   1.546023 seconds (19.89 k allocations: 119.673 MiB, 1.70% gc time)
+@time reg(df, @formula(y ~ x1 + fe(id1) + fe(id1)&x2 + fe(id2) + fe(id2)&x2))
+# 1.679007 seconds (3.61 M allocations: 334.110 MiB, 2.31% gc time, 45.11% compilation time: <1% of which was recompilation)
+@time reg(df, @formula(y ~ fe(id1)*x1 + fe(id2)*x2))
+# 4.245766 seconds (3.18 M allocations: 336.433 MiB, 0.46% gc time, 16.92% compilation time)
 
 
 
 
-+# fixest
+
+# fixest
 n = 10_000_000
 nb_dum = [div(n,20), floor(Int, sqrt(n)), floor(Int, n^.33)]
 N = nb_dum.^3

src/fit.jl

@@ -24,6 +24,8 @@ Estimate a linear model with high dimensional categorical variables / instrument
 Models with instruments variables are estimated using 2SLS. `reg` tests for weak instruments by computing the Kleibergen-Paap rk Wald F statistic, a generalization of the Cragg-Donald Wald F statistic for non i.i.d. errors. The statistic is similar to the one returned by the Stata command `ivreg2`.
 
 Regressors that are collinear with other regressors (or with the fixed effects) are dropped from the estimation. A dropped coefficient is reported as `0` with a `NaN` standard error (and `NaN` t-statistic, p-value, and confidence interval).
+An RHS term exactly spanned by a continuous-slope fixed effect, such as `x` in
+`fe(id)*x`, is structurally unidentified and is omitted from the coefficient output entirely.
 
 ### Examples
 ```julia
@@ -64,7 +66,7 @@ function reg(df,
     first_stage::Bool = true)
     StatsAPI.fit(FixedEffectModel, formula, df, vcov; contrasts = contrasts, weights = weights, save = save, method = method, nthreads = nthreads, double_precision = double_precision, tol = tol, maxiter = maxiter, drop_singletons = drop_singletons, progress_bar = progress_bar, subset = subset, first_stage = first_stage)
 end
-    
+
 function StatsAPI.fit(::Type{FixedEffectModel},     
     @nospecialize(formula::FormulaTerm),
     @nospecialize(df),
@@ -90,8 +92,8 @@ function StatsAPI.fit(::Type{FixedEffectModel},
     ========================================================#
 
     if method == :gpu
-        @info "method = :gpu is deprecated. Use method = :CUDA or method = :Metal"
-        method = :CUDA
+        @info "method = :gpu is deprecated and falls back to CPU with the existing precision setting. Use method = :CUDA after `using CUDA`, or method = :Metal after `using Metal`, to select a GPU backend."
+        method = :cpu
     end
     if nthreads !== nothing
         @info "The keyword argument nthreads is deprecated. Multiple threads are now used by default."
@@ -129,8 +131,8 @@ function StatsAPI.fit(::Type{FixedEffectModel},
     save_fes = save ∈ (:fe, :all) && has_fes
     has_weights = weights !== nothing
 
-    # Compute feM, an AbstractFixedEffectSolver
-    fes, feids, fekeys = parse_fixedeffect(df, formula_fes)
+    # Avoid FE parser on no-FE regressions; it touches table/formula dispatch.
+    fes, feids, fekeys = has_fes ? parse_fixedeffect(df, formula_fes) : (FixedEffect[], Symbol[], Symbol[])
     has_fe_intercept = any(fe.interaction isa UnitWeights for fe in fes)
 
     # remove intercept if absorbed by fixed effects
@@ -164,7 +166,7 @@ function StatsAPI.fit(::Type{FixedEffectModel},
     esample .&= Vcov.completecases(df, vcov)
 
     n_singletons = 0
-    if drop_singletons
+    if drop_singletons && has_fes
         n_singletons = drop_singletons!(esample, fes)
     end
 
@@ -287,8 +289,7 @@ function StatsAPI.fit(::Type{FixedEffectModel},
     (see https://github.com/FixedEffects/FixedEffectModels.jl/issues/249).
     ========================================================#
 
-    Xy = Symmetric(hvcat(2, XhatXhat, Xhat'reshape(y, length(y), 1),
-                         zeros(size(Xhat, 2))', [0.0]))
+    Xy = sweep_rhs_matrix(XhatXhat, Xhat, y)
     invsym!(Xy; diagonal = 1:size(Xhat, 2))
     invXhatXhat = Symmetric(.- Xy[1:(end-1),1:(end-1)])
     coef = Xy[1:(end-1),end]

src/partial_out.jl

@@ -51,8 +51,8 @@ function partial_out(
     df = DataFrame(df; copycols = false)
 
     if method == :gpu
-        @info "method = :gpu is deprecated. Use method = :CUDA or method = :Metal"
-        method = :CUDA
+        @info "method = :gpu is deprecated and falls back to CPU with the existing precision setting. Use method = :CUDA after `using CUDA`, or method = :Metal after `using Metal`, to select a GPU backend."
+        method = :cpu
     end
 
     if  (ConstantTerm(0) ∉ eachterm(f.rhs)) && (ConstantTerm(1) ∉ eachterm(f.rhs))

src/utils/basecol.jl

@@ -78,15 +78,60 @@ Build a Symmetric matrix from upper-triangular blocks, filling the lower triangl
 """
 function upper_block_symmetric(A11, A12, A13, A22, A23, A33)
     n1, n2, n3 = size(A11, 1), size(A22, 1), size(A33, 1)
-    Symmetric(hvcat(3, A11, A12, A13,
-                       zeros(n2, n1), A22, A23,
-                       zeros(n3, n1), zeros(n3, n2), A33))
+    # Avoid hvcat here: broad concatenation methods are easy for unrelated packages to invalidate.
+    out = Matrix{Float64}(undef, n1 + n2 + n3, n1 + n2 + n3)
+    r1 = 1:n1
+    r2 = (n1 + 1):(n1 + n2)
+    r3 = (n1 + n2 + 1):(n1 + n2 + n3)
+    @views out[r1, r1] .= A11
+    @views out[r1, r2] .= A12
+    @views out[r1, r3] .= A13
+    @views out[r2, r1] .= 0.0
+    @views out[r2, r2] .= A22
+    @views out[r2, r3] .= A23
+    @views out[r3, r1] .= 0.0
+    @views out[r3, r2] .= 0.0
+    @views out[r3, r3] .= A33
+    return Symmetric(out)
 end
 
 function upper_block_symmetric(A11, A12, A22)
-    n2 = size(A22, 1)
-    Symmetric(hvcat(2, A11, A12,
-                       zeros(n2, size(A11, 1)), A22))
+    n1, n2 = size(A11, 1), size(A22, 1)
+    # Avoid hvcat here: broad concatenation methods are easy for unrelated packages to invalidate.
+    out = Matrix{Float64}(undef, n1 + n2, n1 + n2)
+    r1 = 1:n1
+    r2 = (n1 + 1):(n1 + n2)
+    @views out[r1, r1] .= A11
+    @views out[r1, r2] .= A12
+    @views out[r2, r1] .= 0.0
+    @views out[r2, r2] .= A22
+    return Symmetric(out)
+end
+
+function block_matrix(A11, A12, A21, A22)
+    # Keep IV block assembly on concrete Matrix{Float64} code, away from generic hvcat.
+    n1, n2 = size(A11, 1), size(A22, 1)
+    out = Matrix{Float64}(undef, n1 + n2, n1 + n2)
+    r1 = 1:n1
+    r2 = (n1 + 1):(n1 + n2)
+    @views out[r1, r1] .= A11
+    @views out[r1, r2] .= A12
+    @views out[r2, r1] .= A21
+    @views out[r2, r2] .= A22
+    return out
+end
+
+function sweep_rhs_matrix(XtX::Symmetric, X::Matrix{Float64}, y::Vector{Float64})
+    # Build the sweep system explicitly to avoid hvcat invalidations on the hot regression core.
+    k = size(X, 2)
+    out = Matrix{Float64}(undef, k + 1, k + 1)
+    if k > 0
+        @views out[1:k, 1:k] .= XtX
+        @views mul!(out[1:k, k + 1], transpose(X), y)
+        @views out[k + 1, 1:k] .= 0.0
+    end
+    out[k + 1, k + 1] = 0.0
+    return Symmetric(out)
 end
 
 """
@@ -166,20 +211,19 @@ function collinearity!(
 
         # Build Xhat via 2SLS
         newZ = hcat(Xexo, Z)
-        newZnewZ = hvcat(2, XexoXexo, XexoZ, XexoZ', ZZ)
+        newZnewZ = block_matrix(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))
+        Pi = ls_solve!(
+            upper_block_symmetric(newZnewZ, newZXendo, zeros(Float64, size(Xendo, 2), size(Xendo, 2))),
+            size(newZnewZ, 2))
         newnewZ = newZ * Pi
         Xhat = hcat(Xexo, newnewZ)
         XhatXhat = upper_block_symmetric(XexoXexo, Xexo'newnewZ, newnewZ'newnewZ)
         X = hcat(Xexo, Xendo)
 
         # prepare residuals for first stage F statistic
         Xendo_res = BLAS.gemm!('N', 'N', -1.0, newZ, Pi, 1.0, Xendo)
-        Pi2 = ls_solve!(Symmetric(hvcat(2, XexoXexo, XexoZ,
-                                zeros(size(Z, 2), size(Xexo, 2)), ZZ)), size(Xexo, 2))
+        Pi2 = ls_solve!(upper_block_symmetric(XexoXexo, XexoZ, ZZ), size(Xexo, 2))
         Z_res = BLAS.gemm!('N', 'N', -1.0, Xexo, Pi2, 1.0, Z)
 
         return Xexo, Xendo, Z, X, Xhat, XhatXhat, basis_coef, perm, Xendo_res, Z_res, Pi

src/utils/formula.jl

@@ -50,7 +50,8 @@ Mark a variable as a high-dimensional fixed effect (a categorical variable to be
 inside a `@formula` passed to [`reg`](@ref) or [`partial_out`](@ref), e.g.
 `@formula(y ~ x + fe(id))`. Several fixed effects are added with `+`, and they can be
 interacted with `&`/`*`: `fe(id)&fe(year)` for interacted fixed effects, `fe(id)&x` for
-group-specific slopes on a continuous variable `x`.
+group-specific slopes on a continuous variable `x`. Because those group-specific slopes
+span the common slope on `x`, `fe(id)*x` omits `x` from the reported coefficients.
 
 When building a formula programmatically, `fe` also accepts a `Symbol`: `fe(:id)`.
 """
@@ -65,14 +66,30 @@ has_fe(@nospecialize(t::FormulaTerm)) = any(has_fe(x) for x in eachterm(t.rhs))
 
 function parse_fe(@nospecialize(f::FormulaTerm))
     if has_fe(f)
-        formula_main = FormulaTerm(f.lhs, Tuple(term for term in eachterm(f.rhs) if !has_fe(term)))
-        formula_fe = FormulaTerm(ConstantTerm(0), Tuple(term for term in eachterm(f.rhs) if has_fe(term)))
+        rhs_terms = eachterm(f.rhs)
+        fe_terms = Tuple(term for term in rhs_terms if has_fe(term))
+        # A continuous-slope FE spans its standalone slope exactly. Remove that
+        # unidentified main-effect column before schema/model-matrix construction.
+        formula_main = FormulaTerm(f.lhs, Tuple(term for term in rhs_terms
+            if !has_fe(term) && !any(fe_term -> _is_absorbed_fe_slope(term, fe_term), fe_terms)))
+        formula_fe = FormulaTerm(ConstantTerm(0), fe_terms)
         return formula_main, formula_fe
     else
         return f, FormulaTerm(ConstantTerm(0), ConstantTerm(0))
     end
 end
 
+function _is_absorbed_fe_slope(@nospecialize(main_term::AbstractTerm),
+                               @nospecialize(fe_term::AbstractTerm))
+    fe_term isa InteractionTerm || return false
+    slope_terms = Tuple(term for term in fe_term.terms if !has_fe(term))
+    isempty(slope_terms) && return false
+    if length(slope_terms) == 1
+        return main_term == slope_terms[1]
+    end
+    return main_term isa InteractionTerm && main_term.terms == slope_terms
+end
+
 
 fesymbol(t::FixedEffectTerm) = t.x
 fesymbol(t::FunctionTerm{typeof(fe)}) = Symbol(t.args[1])
@@ -140,17 +157,34 @@ end
 
 # Construct FixedEffect from an InteractionTerm
 function _parse_fixedeffect(data, @nospecialize(t::InteractionTerm))
-    fes = (x for x in t.terms if has_fe(x))
-    interactions = (x for x in t.terms if !has_fe(x))
-    if !isempty(fes)
-        # x1&x2 from (x1&x2)*id
-        fe_names = [fesymbol(x) for x in fes]
-        v1 = _multiply(data, Symbol.(interactions))
-        fe = FixedEffect((Tables.getcolumn(data, fe_name) for fe_name in fe_names)...; interaction = v1)
-        interactions = string.(interactions)
-        s = vcat(["fe_" * string(fe_name) for fe_name in fe_names], interactions)
-        return fe, Symbol(reduce((x1, x2) -> x1*"&"*x2, s)), fe_names
+    fe_terms = [x for x in t.terms if has_fe(x)]
+    slope_terms = [x for x in t.terms if !has_fe(x)]
+    isempty(fe_terms) && return nothing
+
+    # x1&x2 from (x1&x2)*id
+    fe_names = [fesymbol(x) for x in fe_terms]
+    slope_names = _validate_fe_interactions(data, slope_terms)
+    interaction = _multiply(data, slope_names)
+    fe = FixedEffect((Tables.getcolumn(data, fe_name) for fe_name in fe_names)...; interaction = interaction)
+    pieces = vcat(["fe_" * string(fe_name) for fe_name in fe_names], string.(slope_terms))
+    return fe, Symbol(join(pieces, "&")), fe_names
+end
+
+function _validate_fe_interactions(data, @nospecialize(interactions))
+    out = Symbol[]
+    for t in interactions
+        if !(t isa Term)
+            throw(ArgumentError("Fixed-effect slope interactions only support plain numeric columns. The term `$t` is transformed or expanded; create a numeric column first, then use `fe(...)&new_column`."))
+        end
+        name = Symbol(t)
+        col = Tables.getcolumn(data, name)
+        nonmissing_type = nonmissingtype(eltype(col))
+        if !(nonmissing_type <: Number)
+            throw(ArgumentError("Fixed-effect slope interactions only support numeric columns. Column `$name` has element type $(eltype(col)); convert it to numeric dummy/slope columns first, then use `fe(...)&new_column`."))
+        end
+        push!(out, name)
     end
+    return out
 end
 
 function _multiply(data, ss::AbstractVector)
@@ -163,4 +197,4 @@ function _multiply(data, ss::AbstractVector)
     else
         return convert(AbstractVector{Float64}, replace!(.*((Tables.getcolumn(data, x) for x in ss)...), missing => 0))
     end
-end
+end

test/collinearity.jl

@@ -89,3 +89,24 @@ end
   @test isnan(stderror(m)[3])
   @test !isnan(stderror(m)[2])
 end
+
+@testset "fixed-effect slopes omit matching RHS terms" begin
+  df = DataFrame(
+    id1 = repeat(1:3, inner = 4),
+    id2 = repeat(1:2, inner = 6),
+    x1 = collect(1.0:12.0),
+    x2 = [3.0, 1.0, 4.0, 1.0, 5.0, 9.0, 2.0, 6.0, 5.0, 3.0, 5.0, 8.0],
+  )
+  df.y = df.x1 .* repeat([1.0, 2.0, 3.0], inner = 4) .+
+         df.x2 .* repeat([0.5, -0.5], inner = 6) .+ sin.(df.x1)
+
+  m = @test_logs min_level=Base.CoreLogging.Info reg(
+    df, @formula(y ~ fe(id1)*x1 + fe(id2)*x2); progress_bar = false)
+  @test isempty(coefnames(m))
+  @test isempty(coef(m))
+
+  # Unrelated, data-driven collinearity with an intercept FE is still reported.
+  df.group_constant = Float64.(df.id1)
+  @test_logs (:info, r"RHS-variable group_constant is collinear with the fixed effects") reg(
+    df, @formula(y ~ group_constant + fe(id1)); progress_bar = false)
+end

test/fit.jl

@@ -37,10 +37,10 @@ using CUDA, Metal
 	@test coef(x)  ≈ [-1.08471] atol = 1e-4
 	m = @formula Sales ~ Price + fe(State) + fe(State)*Year
 	x = reg(df, m)
-	@test coef(x) ≈  [-0.53470, 0.0] atol = 1e-4
+	@test coef(x) ≈  [-0.53470] atol = 1e-4
 	m = @formula Sales ~ Price + fe(State)*Year
 	x = reg(df, m)
-	@test coef(x) ≈  [-0.53470, 0.0] atol = 1e-4
+	@test coef(x) ≈  [-0.53470] atol = 1e-4
 
 	#@test isempty(coef(reg(df, @formula(Sales ~ 0), @fe(State*Price))))
 	df.mState = div.(df.State, 10)
@@ -76,7 +76,7 @@ using CUDA, Metal
 	# SSR does not work well here
 	m = @formula Sales ~ Pimin + (Price&NDI)*fe(State)
 	x = reg(df, m)
-	@test coef(x) ≈   [0.421406, 0.0] atol = 1e-4
+	@test coef(x) ≈   [0.421406] atol = 1e-4
 
 	# only one intercept
 	m = @formula Sales ~ 1 + fe(State) + fe(Year)
@@ -813,9 +813,10 @@ end
 end
 
 @testset "keyword arguments" begin
-	df = DataFrame(y = [1.0, 2.0, 2.5], x = [0.0, 1.0, 2.0])
+	df = DataFrame(y = [1.0, 2.0, 2.5], x = [0.0, 1.0, 2.0], id = [1, 1, 2])
 	@test_logs (:info, r"The keyword argument nthreads is deprecated") reg(df, @formula(y ~ x), nthreads = 1)
-	@test_logs (:info, r"method = :gpu is deprecated") reg(df, @formula(y ~ x), method = :gpu)
+	@test_logs (:info, r"method = :gpu is deprecated and falls back to CPU") reg(df,
+		@formula(y ~ fe(id)), method = :gpu)
 end
 
 @testset "error handling" begin