diff --git a/Project.toml b/Project.toml index 3f2ea62c..e4f039a1 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "FixedEffectModels" uuid = "9d5cd8c9-2029-5cab-9928-427838db53e3" -version = "2.0.0" +version = "2.1.0" [deps] DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" diff --git a/README.md b/README.md index fa3c9528..d77f8c53 100755 --- a/README.md +++ b/README.md @@ -6,7 +6,7 @@ This package estimates linear models with high dimensional categorical variables The package is registered in the [`General`](https://github.com/JuliaRegistries/General) registry and so can be installed at the REPL with `] add FixedEffectModels`. ## Benchmarks -The objective of the package is similar to the Stata command [`reghdfe`](https://github.com/sergiocorreia/reghdfe) and the R packages [`lfe`](https://cran.r-project.org/web/packages/lfe/lfe.pdf) and [`fixest`](https://lrberge.github.io/fixest/). The package is much faster than `reghdfe` (Stata) or `lfe` (R). It also tends to be a bit faster than the more recent `fixest` (R), especially when computing clustered standard errors. `FixedEffectModels` can also run on NVIDIA or Apple GPUs for even faster performances (see below) +The objective of the package is similar to the Stata command [`reghdfe`](https://github.com/sergiocorreia/reghdfe) and the R packages [`lfe`](https://cran.r-project.org/web/packages/lfe/lfe.pdf) and [`fixest`](https://lrberge.github.io/fixest/). The package is much faster than `reghdfe` (Stata) or `lfe` (R). It also tends to be a bit faster than the more recent `fixest` (R). `FixedEffectModels` can also run on NVIDIA or Apple GPUs for even faster performances (see below) ![benchmark](http://www.matthieugomez.com/files/fixedeffectmodels_benchmark.png) diff --git a/benchmark/benchmark.csv b/benchmark/benchmark.csv index b72424db..a87e9682 100644 --- a/benchmark/benchmark.csv +++ b/benchmark/benchmark.csv @@ -1 +1,6 @@ -Order,Command,FixedEffectModels.jl (Julia),fixest (R),lfe (R),reghdfe (Stata) 1,simple,0.35,0.317,1.843, 0.61 2,1 hd fe,0.463 ,0.704 ,14.831, 4.64 3,2 hd fe,1.00,1.297 ,10.626, 22.99 4, 1 cluster se,0.38058,0.700 ,9.255, 8.28 5, 2 clusters se,0.765,1.803,96.958, 70.44 \ No newline at end of file +Order,Command,FixedEffectModels.jl (Julia),fixest (R),lfe (R),reg / reghdfe (Stata) +1,simple,0.134,0.354,0.914,0.46 +2,1 hd fe,0.204,0.444,5.770,2.67 +3,2 hd fe,0.490,0.770,3.504,15.14 +4,1 cluster se,0.118,0.498,3.271,6.44 +5,2 clusters se,0.385,1.449,49.466,35.71 diff --git a/benchmark/benchmark.jl b/benchmark/benchmark.jl index af7e9cbb..518c773b 100755 --- a/benchmark/benchmark.jl +++ b/benchmark/benchmark.jl @@ -1,6 +1,6 @@ -using DataFrames, Random, CategoricalArrays +using DataFrames, Random @time using FixedEffectModels -# 0.418712 seconds (742.49 k allocations: 45.500 MiB, 4.07% gc time, 1.07% compilation time) +# 0.442555 seconds (624.29 k allocations: 41.117 MiB, 2.59% gc time, 1.01% compilation time) # Very simple setup N = 10000000 K = 100 @@ -12,21 +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)) -# 1.810739 seconds (14.05 M allocations: 1.583 GiB, 3.76% gc time, 84.28% compilation time: 91% of which was recompilation) +# 0.175934 seconds (3.43 k allocations: 386.445 MiB, 5.27% gc time, 21.75% compilation time) @time reg(df, @formula(y ~ x1 + x2)) -# 0.288344 seconds (712 allocations: 920.405 MiB, 18.71% gc time) +# 0.198124 seconds (310 allocations: 386.310 MiB, 16.34% gc time) @time reg(df, @formula(y ~ x1 + x2), Vcov.cluster(:id2)) -# 0.528590 seconds (1.33 M allocations: 1.039 GiB, 8.76% gc time, 55.33% compilation time) +# 0.221376 seconds (163.66 k allocations: 471.871 MiB, 19.93% gc time, 43.55% compilation time) @time reg(df, @formula(y ~ x1 + x2), Vcov.cluster(:id2)) -# 0.268619 seconds (879 allocations: 997.916 MiB, 15.08% gc time) +# 0.124714 seconds (478 allocations: 463.822 MiB, 6.91% gc time) @time reg(df, @formula(y ~ x1 + x2 + fe(id1))) -# 0.824536 seconds (3.09 M allocations: 1.426 GiB, 5.37% gc time, 61.39% compilation time: 3% of which was recompilation) +# 0.251411 seconds (187.01 k allocations: 673.778 MiB, 15.87% gc time, 52.26% compilation time) @time reg(df, @formula(y ~ x1 + x2 + fe(id1))) -# 0.356793 seconds (1.41 k allocations: 1.276 GiB, 19.31% gc time) +# 0.209155 seconds (2.77 k allocations: 702.849 MiB, 32.16% gc time) @time reg(df, @formula(y ~ x1 + x2 + fe(id1)), Vcov.cluster(:id1)) -# 0.435500 seconds (495.96 k allocations: 1.381 GiB, 15.97% gc time, 10.72% compilation time) +# 0.242202 seconds (148.74 k allocations: 752.135 MiB, 14.90% gc time, 12.01% compilation time) @time reg(df, @formula(y ~ x1 + x2 + fe(id1) + fe(id2))) -# 1.367264 seconds (1.91 M allocations: 1.592 GiB, 5.74% gc time, 23.85% compilation time: 20% of which was recompilation) +# 0.594494 seconds (404.37 k allocations: 913.358 MiB, 11.85% gc time, 26.83% compilation time) # More complicated setup N = 800000 # number of observations @@ -39,17 +39,17 @@ x2 = cos.(id1) + sin.(id2) + randn(N) 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) +# 0.727029 seconds (84.78 k allocations: 69.409 MiB, 4.53% 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) +# 1.262707 seconds (1.85 M allocations: 204.607 MiB, 3.03% gc time, 43.07% compilation time) @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) - +# 0.525260 seconds (669.35 k allocations: 141.178 MiB, 7.81% gc time, 28.31% compilation time) # fixest +using CategoricalArrays n = 10_000_000 nb_dum = [div(n,20), floor(Int, sqrt(n)), floor(Int, n^.33)] N = nb_dum.^3 @@ -57,11 +57,18 @@ id1 = categorical(rand(1:nb_dum[1], n)) id2 = categorical(rand(1:nb_dum[2], n)) id3 = categorical(rand(1:nb_dum[3], n)) X1 = rand(n) -ln_y = 3 .* X1 .+ rand(n) +ln_y = 3 .* X1 .+ rand(n) df = DataFrame(X1 = X1, ln_y = ln_y, id1 = id1, id2 = id2, id3 = id3) +# the first call recompiles code invalidated by loading CategoricalArrays +@time reg(df, @formula(ln_y ~ X1 + fe(id1)), Vcov.cluster(:id1)) +# 2.247757 seconds (14.46 M allocations: 1.324 GiB, 3.71% gc time, 91.03% compilation time: 67% of which was recompilation) @time reg(df, @formula(ln_y ~ X1 + fe(id1)), Vcov.cluster(:id1)) -# 0.311420 seconds (1.35 k allocations: 1.052 GiB, 22.30% gc time) +# 0.198736 seconds (2.26 k allocations: 636.432 MiB, 23.03% gc time) @time reg(df, @formula(ln_y ~ X1 + fe(id1) + fe(id2)), Vcov.cluster(:id1)) -# 0.808992 seconds (3.52 k allocations: 1.272 GiB, 8.68% gc time) +# 1.093537 seconds (1.77 M allocations: 1.125 GiB, 6.85% gc time, 43.38% compilation time) +@time reg(df, @formula(ln_y ~ X1 + fe(id1) + fe(id2)), Vcov.cluster(:id1)) +# 0.661163 seconds (6.27 k allocations: 1.006 GiB, 12.54% gc time) +@time reg(df, @formula(ln_y ~ X1 + fe(id1) + fe(id2) + fe(id3)), Vcov.cluster(:id1)) +# 0.855344 seconds (327.35 k allocations: 1.280 GiB, 10.46% gc time, 17.11% compilation time) @time reg(df, @formula(ln_y ~ X1 + fe(id1) + fe(id2) + fe(id3)), Vcov.cluster(:id1)) -# 0.950808 seconds (4.75 k allocations: 1.496 GiB, 7.48% gc time) +# 0.782286 seconds (8.00 k allocations: 1.265 GiB, 10.83% gc time) diff --git a/benchmark/benchmark.md b/benchmark/benchmark.md index 3b9fb495..7f496a60 100755 --- a/benchmark/benchmark.md +++ b/benchmark/benchmark.md @@ -1,9 +1,10 @@ ### Simple benchmark -![benchmark](https://cdn.rawgit.com/matthieugomez/FixedEffectModels.jl/4c7d1db39377f1ee649624c909c9017f92484114/benchmark/result.svg) +![benchmark](fixedeffectmodels_benchmark.png) +All timings below are from the same machine (Apple M4, 16GB RAM, 2026-08-28). Code to reproduce this graph: - FixedEffectModels.jl v1.9.0 (Julia 1.9) + FixedEffectModels.jl v2.1.0 (Julia 1.12.7, started with `julia -t auto`; timings are warm calls) ```julia using DataFrames, CategoricalArrays, FixedEffectModels N = 10_000_000 @@ -15,19 +16,19 @@ Code to reproduce this graph: y= 3 .* x1 .+ 2 .* x2 .+ sin.(id1) .+ cos.(id2).^2 .+ randn(N) df = DataFrame(id1 = categorical(id1), id2 = categorical(id2), x1 = x1, x2 = x2, y = y) @time reg(df, @formula(y ~ x1 + x2)) - # 0.338749 seconds (450 allocations: 691.441 MiB, 2.30% gc time) + # 0.133884 seconds (313 allocations: 386.310 MiB, 1.42% gc time) @time reg(df, @formula(y ~ x1 + x2 + fe(id1))) - # 0.463058 seconds (1.00 k allocations: 929.129 MiB, 13.31% gc time) + # 0.203769 seconds (2.75 k allocations: 664.707 MiB, 28.83% gc time) @time reg(df, @formula(y ~ x1 + x2 + fe(id1) + fe(id2))) - # 1.006031 seconds (3.22 k allocations: 1.057 GiB, 1.68% gc time) + # 0.489542 seconds (7.22 k allocations: 893.853 MiB, 17.39% gc time) @time reg(df, @formula(y ~ x1 + x2), Vcov.cluster(:id1)) - # 0.380562 seconds (580 allocations: 771.606 MiB, 3.07% gc time) + # 0.117933 seconds (468 allocations: 535.131 MiB, 0.43% gc time) @time reg(df, @formula(y ~ x1 + x2), Vcov.cluster(:id1, :id2)) - #0.765847 seconds (719 allocations: 1.128 GiB, 2.01% gc time) + # 0.384734 seconds (641 allocations: 849.876 MiB, 13.68% gc time) ```` - fixest v0.8.4 (R 4.2.2) + fixest v0.13.2 (R 4.4.2) ```R library(fixest) N = 10000000 @@ -41,23 +42,23 @@ Code to reproduce this graph: df[, "y"] = 3 * df[, "x1"] + 2 * df[, "x2"] + sin(as.numeric(df[, "id1"])) + cos(as.numeric(df[, "id2"])) + runif(N) system.time(feols(y ~ x1 + x2, df)) #> user system elapsed - #> 0.280 0.036 0.317 + #> 0.324 0.028 0.354 system.time(feols(y ~ x1 + x2|id1, df)) #> user system elapsed - #> 0.616 0.089 0.704 + #> 0.396 0.048 0.444 system.time(feols(y ~ x1 + x2|id1 + id2, df)) #> user system elapsed - #> 1.181 0.120 1.297 + #> 0.714 0.056 0.770 system.time(feols(y ~ x1 + x2, cluster = "id1", df)) #> user system elapsed - #> 0.630 0.071 0.700 + #> 0.438 0.046 0.498 system.time(feols(y ~ x1 + x2, cluster = c("id1", "id2"), df)) #> user system elapsed - #> 1.570 0.197 1.803 + #> 1.345 0.102 1.449 ``` - lfe v2.8-8 (R 4.2.2) + lfe v3.1.1 (R 4.4.2) ```R library(lfe) N = 10000000 @@ -72,22 +73,23 @@ Code to reproduce this graph: system.time(felm(y ~ x1 + x2, df)) #> user system elapsed - #> 1.137 0.232 1.596 + #> 0.807 0.107 0.914 system.time(felm(y ~ x1 + x2|id1, df)) #> user system elapsed - #> 7.08 0.41 7.46 + #> 5.589 0.202 5.770 system.time(felm(y ~ x1 + x2|id1 + id2, df)) #> user system elapsed - #> 4.832 0.370 4.615 + #> 3.580 0.300 3.504 system.time(felm(y ~ x1 + x2|0|0|id1, df)) #> user system elapsed - #> 3.712 0.287 3.996 + #> 3.105 0.166 3.271 system.time(felm(y ~ x1 + x2|0|0|id1 + id2, df)) #> user system elapsed - #> 59.119 0.889 59.946 + #> 48.771 0.552 49.466 + ``` - reghdfe version 5.6.8 03mar2019 (Stata 16.1) + reghdfe 5.7.3 / ivreg2 4.1.11 (Stata 19.0 MP, 2 cores) ``` clear all local N = 10000000 @@ -102,13 +104,13 @@ Code to reproduce this graph: set rmsg on reg y x1 x2 - #> r; t=0.61 + #> r; t=0.46 reghdfe y x1 x2, a(id1) - #>r; t=4.64 + #> r; t=2.67 reghdfe y x1 x2, a(id1 id2) - #> r; t==22.99 + #> r; t=15.14 reg y x1 x2, cl(id1) - #> r; t=8.28 + #> r; t=6.44 ivreg2 y x1 x2, cluster(id1 id2) - #> r; t=70.44 + #> r; t=35.71 ```` diff --git a/benchmark/fixedeffectmodels_benchmark.png b/benchmark/fixedeffectmodels_benchmark.png index b79ebed1..90237c76 100644 Binary files a/benchmark/fixedeffectmodels_benchmark.png and b/benchmark/fixedeffectmodels_benchmark.png differ diff --git a/benchmark/result.jl b/benchmark/result.jl index bebb8a44..62a61804 100644 --- a/benchmark/result.jl +++ b/benchmark/result.jl @@ -1 +1,21 @@ -using DataFrames, CSV, StatsPlots df = CSV.read("/Users/matthieugomez/Dropbox/Github/FixedEffectModels.jl/benchmark/benchmark2.csv", DataFrame) df."fixest (R)" = df."fixest (R)" ./ df."FixedEffectModels.jl (Julia)" df."lfe (R)" = df."lfe (R)" ./ df."FixedEffectModels.jl (Julia)" df."reg / reghdfe (Stata)" = df."reg / reghdfe (Stata)" ./ df."FixedEffectModels.jl (Julia)" df."FixedEffectModels.jl (Julia)" = df."FixedEffectModels.jl (Julia)" ./ df."FixedEffectModels.jl (Julia)" mdf = stack(df, Not([:Command, :Order])) mdf = rename(mdf, :variable => :Language) p = @df mdf plot( :Command, :value, group = :Language, yaxis = :log10, xlabel = "Command", ylabel = "Time (Ratio to Julia)", legend = :top, seriestype = :scatter, # or :line / :path depending on what you want palette = :tol_light, dpi = 200, size=(8 * 100 * 2 /3, 5 * 100 * 2 /3) ) savefig("/Users/matthieugomez/Dropbox/Github/FixedEffectModels.jl/benchmark/fixedeffectmodels_benchmark.png") \ No newline at end of file +using DataFrames, CSV, StatsPlots +df = CSV.read(joinpath(@__DIR__, "benchmark.csv"), DataFrame) +df."fixest (R)" = df."fixest (R)" ./ df."FixedEffectModels.jl (Julia)" +df."lfe (R)" = df."lfe (R)" ./ df."FixedEffectModels.jl (Julia)" +df."reg / reghdfe (Stata)" = df."reg / reghdfe (Stata)" ./ df."FixedEffectModels.jl (Julia)" +df."FixedEffectModels.jl (Julia)" = df."FixedEffectModels.jl (Julia)" ./ df."FixedEffectModels.jl (Julia)" +mdf = stack(df, Not([:Command, :Order])) +mdf = rename(mdf, :variable => :Language) +p = @df mdf plot( + :Command, :value, + group = :Language, + yaxis = :log10, + xlabel = "Command", + ylabel = "Time (Ratio to Julia)", + legend = :top, + seriestype = :scatter, + palette = :tol_light, + dpi = 200, + right_margin = 8Plots.mm, + size = (8 * 100 * 2 / 3, 5 * 100 * 2 / 3)) +savefig(joinpath(@__DIR__, "fixedeffectmodels_benchmark.png")) diff --git a/src/FixedEffectModel.jl b/src/FixedEffectModel.jl index ee125f13..e48d4df3 100644 --- a/src/FixedEffectModel.jl +++ b/src/FixedEffectModel.jl @@ -51,6 +51,18 @@ struct FixedEffectModel <: RegressionModel # for IV F_kp::Float64 # First Stage F statistics KP p_kp::Float64 # First Stage p value KP + + # @nospecialize on the formula arguments: their concrete types change with + # every formula shape, and the fields are abstract anyway, so specializing + # the constructor would recompile it on each new formula + function FixedEffectModel(coef, vcov, vcov_type, nclusters, esample, residuals, fe, fekeys, + coefnames, responsename, @nospecialize(formula), @nospecialize(formula_schema), + contrasts, nobs, dof, dof_fes, dof_residual, rss, tss, F, p, + iterations, converged, r2_within, F_kp, p_kp) + return new(coef, vcov, vcov_type, nclusters, esample, residuals, fe, fekeys, + coefnames, responsename, formula, formula_schema, contrasts, nobs, dof, + dof_fes, dof_residual, rss, tss, F, p, iterations, converged, r2_within, F_kp, p_kp) + end end """ diff --git a/src/FixedEffectModels.jl b/src/FixedEffectModels.jl index 6e387a59..9a7aa89b 100644 --- a/src/FixedEffectModels.jl +++ b/src/FixedEffectModels.jl @@ -20,6 +20,7 @@ include("utils/fixedeffects.jl") include("utils/basecol.jl") include("utils/tss.jl") include("utils/formula.jl") +include("utils/vectorterms.jl") include("FixedEffectModel.jl") include("fit.jl") include("partial_out.jl") @@ -39,10 +40,13 @@ Vcov @compile_workload begin - df = DataFrame(x1 = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0], x2 = [1.0, 2.0, 4.0, 4.0, 3.0, 5.0], y = [3.0, 4.0, 4.0, 5.0, 1.0, 2.0], id = [1, 1, 2, 2, 3, 3]) + df = DataFrame(x1 = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0], x2 = [1.0, 2.0, 4.0, 4.0, 3.0, 5.0], y = [3.0, 4.0, 4.0, 5.0, 1.0, 2.0], id = [1, 1, 2, 2, 3, 3], s = ["a", "a", "b", "b", "a", "b"]) reg(df, @formula(y ~ x1 + x2)) reg(df, @formula(y ~ x1 + fe(id))) reg(df, @formula(y ~ x1), Vcov.cluster(:id)) + # cover the categorical and interaction term types (the vector-based + # machinery compiles per term type, not per formula shape) + reg(df, @formula(y ~ x1 + s + x1&x2 + x1&s)) end diff --git a/src/fit.jl b/src/fit.jl index 1a7e0cd5..2d097e1b 100644 --- a/src/fit.jl +++ b/src/fit.jl @@ -49,21 +49,21 @@ df = dataset("plm", "Cigar") fit(FixedEffectModel, @formula(Sales ~ NDI + fe(State) + fe(State)&Year), df) ``` """ -function reg(df, - formula::FormulaTerm, - vcov::CovarianceEstimator = Vcov.simple(); - contrasts::Dict = Dict{Symbol, Any}(), - weights::Union{Symbol, Nothing} = nothing, - save::Union{Bool, Symbol} = :none, - method::Symbol = :cpu, - nthreads::Union{Integer, Nothing} = nothing, - double_precision::Bool = method == :cpu, - tol::Real = 1e-6, - maxiter::Integer = 10000, - drop_singletons::Bool = true, - progress_bar::Bool = true, - subset::Union{Nothing, AbstractVector} = nothing, - first_stage::Bool = true) +function reg(@nospecialize(df), + @nospecialize(formula::FormulaTerm), + @nospecialize(vcov::CovarianceEstimator = Vcov.simple()); + @nospecialize(contrasts::Dict = Dict{Symbol, Any}()), + @nospecialize(weights::Union{Symbol, Nothing} = nothing), + @nospecialize(save::Union{Bool, Symbol} = :none), + @nospecialize(method::Symbol = :cpu), + @nospecialize(nthreads::Union{Integer, Nothing} = nothing), + @nospecialize(double_precision::Bool = method == :cpu), + @nospecialize(tol::Real = 1e-6), + @nospecialize(maxiter::Integer = 10000), + @nospecialize(drop_singletons::Bool = true), + @nospecialize(progress_bar::Bool = true), + @nospecialize(subset::Union{Nothing, AbstractVector} = nothing), + @nospecialize(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 @@ -113,13 +113,19 @@ function StatsAPI.fit(::Type{FixedEffectModel}, ========================================================# formula_origin = formula - if !omitsintercept(formula) && !hasintercept(formula) - formula = FormulaTerm(formula.lhs, InterceptTerm{true}() + formula.rhs) + formula.lhs isa FormulaTerm && throw(ArgumentError("Malformed formula: the left-hand side cannot contain `~`. The instrumental-variable syntax is `y ~ exogenous + (endogenous ~ instruments)`.")) + # Work on Vector{AbstractTerm} rather than the formula's term tuple: the + # vector-based machinery in utils/vectorterms.jl compiles per term type + # rather than per formula shape, which avoids recompilation on every new + # formula (see that file for details). + exo_ts = termvector(formula.rhs) + if !_omitsintercept(exo_ts) && !_hasintercept(exo_ts) + pushfirst!(exo_ts, InterceptTerm{true}()) end - formula, formula_endo, formula_iv = parse_iv(formula) - has_iv = formula_iv != FormulaTerm(ConstantTerm(0), ConstantTerm(0)) - formula, formula_fes = parse_fe(formula) - has_fes = formula_fes != FormulaTerm(ConstantTerm(0), ConstantTerm(0)) + endo_ts, iv_ts = parse_iv!(exo_ts) + has_iv = !isempty(endo_ts) + fe_ts = parse_fe!(exo_ts) + has_fes = !isempty(fe_ts) # HC2/HC3 compute leverage from the partialled-out (demeaned) regressors only, which # omits the absorbed fixed-effect contribution to leverage. The resulting standard # errors are silently anti-conservative, so guard against the combination (mirroring @@ -132,25 +138,28 @@ function StatsAPI.fit(::Type{FixedEffectModel}, has_weights = weights !== nothing # 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[]) + fes, feids, fekeys = has_fes ? parse_fixedeffect(df, fe_ts) : (FixedEffect[], Symbol[], Symbol[]) has_fe_intercept = any(fe.interaction isa UnitWeights for fe in fes) # remove intercept if absorbed by fixed effects if has_fe_intercept - formula = FormulaTerm(formula.lhs, tuple(InterceptTerm{false}(), (term for term in eachterm(formula.rhs) if !isa(term, Union{ConstantTerm,InterceptTerm}))...)) + filter!(t -> !isa(t, Union{ConstantTerm, InterceptTerm}), exo_ts) + pushfirst!(exo_ts, InterceptTerm{false}()) end - has_intercept = hasintercept(formula) + has_intercept = _hasintercept(exo_ts) #======================================================== Create boolean vector esample that is true for observations used in estimation ========================================================# # Collect all variable names needed to detect missing values and build model matrices - exo_vars = unique(StatsModels.termvars(formula)) - iv_vars = unique(StatsModels.termvars(formula_iv)) - endo_vars = unique(StatsModels.termvars(formula_endo)) - fe_vars = unique(StatsModels.termvars(formula_fes)) - all_vars = unique(vcat(exo_vars, endo_vars, iv_vars, fe_vars)) + response_vars = StatsModels.termvars(formula.lhs) + exo_vars = termvars_vector(exo_ts) + endo_vars = termvars_vector(endo_ts) + iv_vars = termvars_vector(iv_ts) + fe_vars = termvars_vector(fe_ts) + model_vars = unique(vcat(response_vars, exo_vars, endo_vars, iv_vars)) + all_vars = unique(vcat(model_vars, fe_vars)) # Create esample that returns obs used in estimation esample = completecases(df, all_vars) @@ -181,7 +190,9 @@ function StatsAPI.fit(::Type{FixedEffectModel}, else weights = uweights(nobs) end - subdf = DataFrame((; (x => disallowmissing(view(df[!, x], esample)) for x in all_vars)...)) + # constructed from a vector of pairs, not a NamedTuple splat: column names + # stay out of the type domain, so nothing recompiles on new variable names + subdf = DataFrame(Pair{Symbol, AbstractVector}[x => disallowmissing(view(df[!, x], esample)) for x in all_vars], copycols = false) subfes = FixedEffect[fe[esample] for fe in fes] vcov_method = Vcov.materialize(view(df, esample, :), vcov) @@ -189,33 +200,35 @@ function StatsAPI.fit(::Type{FixedEffectModel}, Dataframe --> Matrix ========================================================# - s = schema(formula, subdf, contrasts) - - formula_schema = apply_schema(formula, s, FixedEffectModel, has_fe_intercept) + sch = build_schema(model_vars, subdf, contrasts) - # for a Vector{Float64}, convert(Vector{Float64}, y) aliases y - y = convert(Vector{Float64}, response(formula_schema, subdf)) - Xexo = convert(Matrix{Float64}, modelmatrix(formula_schema, subdf)) - response_name, coefnames_exo = coefnames(formula_schema) + exo_ts = apply_schema_vector(exo_ts, sch, StatisticalModel, has_fe_intercept) + lhs_schema = apply_schema(formula.lhs, sch, StatisticalModel) + # the constructor always copies, so y never aliases a column of df + y = Vector{Float64}(modelcols_term(lhs_schema, subdf, nobs)) + Xexo = modelmatrix_vector(exo_ts, subdf, nobs) + response_name = coefnames(lhs_schema) + coefnames_exo = coefnames_vector(exo_ts) Xendo = Array{Float64}(undef, nobs, 0) Z = Array{Float64}(undef, nobs, 0) - coefnames_endo = typeof(coefnames_exo)[] - coefnames_iv = typeof(coefnames_exo)[] + coefnames_endo = String[] + coefnames_iv = String[] if has_iv - formula_endo_schema = apply_schema(formula_endo, schema(formula_endo, subdf, contrasts), StatisticalModel) - Xendo = convert(Matrix{Float64}, modelmatrix(formula_endo_schema, subdf)) - _, coefnames_endo = coefnames(formula_endo_schema) - - formula_iv_schema = apply_schema(formula_iv, schema(formula_iv, subdf, contrasts), StatisticalModel) - _, coefnames_iv = coefnames(formula_iv_schema) - - Z = convert(Matrix{Float64}, modelmatrix(formula_iv_schema, subdf)) + endo_ts = apply_schema_vector(endo_ts, sch, StatisticalModel) + Xendo = modelmatrix_vector(endo_ts, subdf, nobs) + coefnames_endo = coefnames_vector(endo_ts) - # modify formula to use in predict - formula_schema = FormulaTerm(formula_schema.lhs, MatrixTerm(tuple(eachterm(formula_schema.rhs)..., (term for term in eachterm(formula_endo_schema.rhs) if term != ConstantTerm(0))...))) + iv_ts = apply_schema_vector(iv_ts, sch, StatisticalModel) + Z = modelmatrix_vector(iv_ts, subdf, nobs) + coefnames_iv = coefnames_vector(iv_ts) end coef_names = vcat(coefnames_exo, coefnames_endo) + + # formula with concrete terms, stored in the model and used by predict; + # with IV, the endogenous terms are appended since their coefficients are reported + predict_ts = has_iv ? vcat(exo_ts, AbstractTerm[t for t in endo_ts if StatsModels.width(t) > 0]) : exo_ts + formula_schema = FormulaTerm(lhs_schema, MatrixTerm((predict_ts...,))) # compute tss now before potentially demeaning y tss_total = tss(y, has_intercept || has_fe_intercept, weights) diff --git a/src/partial_out.jl b/src/partial_out.jl index 6095d9da..1ab67d6c 100644 --- a/src/partial_out.jl +++ b/src/partial_out.jl @@ -106,8 +106,11 @@ function partial_out( # Compute residualized Y vars = unique(StatsModels.termvars(formula)) subdf = Tables.columntable(disallowmissing!(df[esample, vars])) + # one pass over the variables, skipping the unused per-column statistics + # (see build_schema); the same schema serves both Y and X below + sch = build_schema(vars, subdf, contrasts) formula_y = FormulaTerm(ConstantTerm(0), (ConstantTerm(0), eachterm(formula.lhs)...)) - formula_y_schema = apply_schema(formula_y, schema(formula_y, subdf, contrasts), StatisticalModel) + formula_y_schema = apply_schema(formula_y, sch, StatisticalModel) Y = convert(Matrix{Float64}, modelmatrix(formula_y_schema, subdf)) ynames = coefnames(formula_y_schema)[2] @@ -125,7 +128,7 @@ function partial_out( # Compute residualized X formula_x = FormulaTerm(ConstantTerm(0), formula.rhs) - formula_x_schema = apply_schema(formula_x, schema(formula_x, subdf, contrasts), StatisticalModel) + formula_x_schema = apply_schema(formula_x, sch, StatisticalModel) X = convert(Matrix{Float64}, modelmatrix(formula_x_schema, subdf)) if has_fes _, b, c = solve_residuals!(eachcol(X), feM; maxiter = maxiter, tol = tol, progress_bar = false) diff --git a/src/utils/vectorterms.jl b/src/utils/vectorterms.jl new file mode 100644 index 00000000..f1b4aecc --- /dev/null +++ b/src/utils/vectorterms.jl @@ -0,0 +1,178 @@ +#======================================================== +Vector-based replicas of the StatsModels formula machinery, for latency. + +StatsModels represents a formula right-hand side as a *tuple* of terms and the +data as a NamedTuple keyed by *column names*. Both encode information in the +type domain, so every new formula shape (number of terms) and every new set of +column names compiles fresh specializations of schema/apply_schema/modelcols — +roughly 0.2-0.3s of compilation per regression with a formula never seen +before, even when every individual term type has been precompiled. + +The functions here operate on Vector{AbstractTerm} and look up columns by +runtime Symbol. Dispatch happens once per term *type* (ContinuousTerm, +CategoricalTerm, InteractionTerm, ...), so after the precompile workload a +regression with a brand-new formula runs entirely warm. The semantics +replicate the StatsModels tuple path exactly; in particular the stateful +FullRank contrast promotion is sequential over terms, so applying it in a +vector loop is equivalent to applying it to a tuple. +========================================================# + +# FormulaTerm rhs (tuple of terms or a single term) -> Vector{AbstractTerm} +function termvector(@nospecialize(x)) + ts = AbstractTerm[] + if x isa Tuple + for t in x + push!(ts, t) + end + else + push!(ts, x) + end + return ts +end + +# same definitions as StatsModels.hasintercept/omitsintercept, on a term vector +_hasintercept(ts::Vector{AbstractTerm}) = + any(t -> t == InterceptTerm{true}() || t == ConstantTerm(1), ts) +_omitsintercept(ts::Vector{AbstractTerm}) = + any(t -> t == InterceptTerm{false}() || t == ConstantTerm(0) || t == ConstantTerm(-1), ts) + +function termvars_vector(ts::Vector{AbstractTerm}) + out = Symbol[] + for t in ts + append!(out, StatsModels.termvars(t)) + end + return unique!(out) +end + +# Split the first (endogenous ~ instruments) term out of rhs (mutating it into +# the exogenous part). Returns endo and iv term vectors, empty when there is no +# IV term. Variables appearing on both sides are treated as exogenous controls. +function parse_iv!(rhs::Vector{AbstractTerm}) + i = findfirst(t -> t isa FormulaTerm, rhs) + i === nothing && return AbstractTerm[], AbstractTerm[] + ivpart = rhs[i]::FormulaTerm + endo_all = termvector(ivpart.lhs) + iv_all = termvector(ivpart.rhs) + both = intersect(endo_all, iv_all) + endos = setdiff(endo_all, both) + exos = setdiff(iv_all, both) + isempty(endos) && throw(ArgumentError("There are no endogeneous variables")) + length(exos) < length(endos) && throw(ArgumentError("Model not identified. There must be at least as many instrumental variables as endogeneneous variables")) + filter!(t -> !isa(t, FormulaTerm), rhs) + append!(rhs, both) + # the leading ConstantTerm(0) suppresses the intercept when the schema is + # applied, so categorical endogenous variables get full dummy coding + return pushfirst!(endos, ConstantTerm(0)), pushfirst!(exos, ConstantTerm(0)) +end + +# Split fixed-effect terms out of rhs (mutating it). A main-effect term exactly +# spanned by a continuous-slope fixed effect is removed as well (see parse_fe). +function parse_fe!(rhs::Vector{AbstractTerm}) + fe_ts = AbstractTerm[t for t in rhs if has_fe(t)] + isempty(fe_ts) && return fe_ts + filter!(t -> !has_fe(t) && !any(fe_t -> _is_absorbed_fe_slope(t, fe_t), fe_ts), rhs) + return fe_ts +end + +function parse_fixedeffect(data, ts::Vector{AbstractTerm}) + fes = FixedEffect[] + feids = Symbol[] + fekeys = Symbol[] + for t in ts + result = _parse_fixedeffect(data, t) + if result !== nothing + push!(fes, result[1]) + push!(feids, result[2]) + append!(fekeys, result[3]) + end + end + return fes, feids, unique(fekeys) +end + +getcol(data, s::Symbol) = Tables.getcolumn(data, s) + +# name-free equivalent of schema(f, data, contrasts): one concrete term per +# variable, computed from the column looked up by runtime symbol +function build_schema(vars::Vector{Symbol}, data, contrasts::Dict) + sch = StatsModels.Schema() + for s in vars + t = Term(s) + col = getcol(data, s) + hint = get(contrasts, s, nothing) + if hint === nothing && eltype(col) <: Number + # concrete_term fills ContinuousTerm's mean/var/min/max with three + # O(n) passes over the column (mean_and_var + extrema); nothing here + # reads those fields, so store NaN placeholders instead + sch.schema[t] = ContinuousTerm(s, NaN, NaN, NaN, NaN) + else + sch.schema[t] = StatsModels.concrete_term(t, col, hint) + end + end + return sch +end + +# equivalent of apply_schema(ts, FullRank(sch), Mod) on a tuple: FullRank +# promotion mutates its `already` set term by term, left to right +function apply_schema_vector(ts::Vector{AbstractTerm}, sch::StatsModels.Schema, Mod::Type, + already_intercept::Bool = false) + fullrank = StatsModels.FullRank(sch) + already_intercept && push!(fullrank.already, InterceptTerm{true}()) + out = AbstractTerm[] + for t in ts + applied = apply_schema(t, fullrank, Mod) + if applied isa Tuple + for a in applied + push!(out, a) + end + else + push!(out, applied) + end + end + return out +end + +# per-term model columns; block eltypes are converted when copied into the +# Float64 model matrix +modelcols_term(t::InterceptTerm{true}, data, n::Integer) = ones(n) +modelcols_term(t::InterceptTerm{false}, data, n::Integer) = Matrix{Float64}(undef, n, 0) +modelcols_term(t::ContinuousTerm, data, n::Integer) = getcol(data, t.sym) +modelcols_term(t::CategoricalTerm, data, n::Integer) = t.contrasts[getcol(data, t.sym), :] +modelcols_term(t::InteractionTerm, data, n::Integer) = + StatsModels.row_kron_insideout(*, (modelcols_term(x, data, n) for x in t.terms)...) +modelcols_term(t::FunctionTerm, data, n::Integer) = + t.f.((modelcols_term(x, data, n) for x in t.args)...) +# bare Term and ConstantTerm only appear inside (protected) FunctionTerm arguments +modelcols_term(t::Term, data, n::Integer) = getcol(data, t.sym) +modelcols_term(t::ConstantTerm, data, n::Integer) = t.n +# fallback for exotic term types: use the generic StatsModels path on the +# columns the term needs +function modelcols_term(@nospecialize(t::AbstractTerm), data, n::Integer) + vars = StatsModels.termvars(t) + nt = NamedTuple{(vars...,)}(((getcol(data, s) for s in vars)...,)) + return modelcols(t, nt) +end + +function modelmatrix_vector(ts::Vector{AbstractTerm}, data, n::Integer) + blocks = [modelcols_term(t, data, n) for t in ts if StatsModels.width(t) > 0] + ncols = 0 + for b in blocks + ncols += size(b, 2) + end + X = Matrix{Float64}(undef, n, ncols) + j = 1 + for b in blocks + k = size(b, 2) + copyto!(view(X, :, j:(j + k - 1)), b) + j += k + end + return X +end + +function coefnames_vector(ts::Vector{AbstractTerm}) + out = String[] + for t in ts + StatsModels.width(t) > 0 || continue + append!(out, StatsModels.vectorize(coefnames(t))) + end + return out +end diff --git a/test/fit.jl b/test/fit.jl index 203e2255..e4ac7b75 100644 --- a/test/fit.jl +++ b/test/fit.jl @@ -1,4 +1,4 @@ -using FixedEffectModels, CategoricalArrays, CSV, DataFrames, Test, LinearAlgebra +using FixedEffectModels, CategoricalArrays, CSV, DataFrames, Test, LinearAlgebra, Random using FixedEffectModels: nullloglikelihood_within using CUDA, Metal @@ -873,3 +873,30 @@ end m = reg(df, @formula(y ~ z + fe(id)&x)) @test FixedEffectModels.dof_fes(m) == 2 end + +@testset "reg does not modify the input dataframe" begin + # The model matrices are built from views into the input columns (no + # defensive copy), so the in-place demeaning and weighting steps must + # only ever touch freshly allocated arrays. + Random.seed!(0) + n = 200 + df1 = DataFrame(y = rand(n), x1 = rand(n), x2 = rand(n), + xm = [i % 7 == 0 ? missing : rand() for i in 1:n], + w = rand(n) .+ 0.5, firm = rand(1:10, n), grp = rand(["a", "b", "c"], n), + instr = rand(n)) + snapshot = deepcopy(df1) + reg(df1, @formula(y ~ x1 + x2)) + reg(df1, @formula(y ~ x1 + xm)) + reg(df1, @formula(y ~ x1 + x2 + fe(firm)), save = :all) + reg(df1, @formula(y ~ x1 + fe(firm) + fe(firm)&x2)) + reg(df1, @formula(y ~ x1 + grp + x1&grp)) + reg(df1, @formula(y ~ x1 + (x2 ~ instr))) + reg(df1, @formula(y ~ x1 + log(w))) + reg(df1, @formula(y ~ x1 + x2), weights = :w) + reg(df1, @formula(y ~ x1 + x2 + fe(firm)), weights = :w, save = :residuals) + reg(df1, @formula(y ~ x1 + x2), subset = df1.x1 .> 0.2) + reg(df1, @formula(y ~ x1), Vcov.cluster(:firm)) + for c in names(df1) + @test isequal(df1[!, c], snapshot[!, c]) + end +end