Reduce latency with vector-based formula

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"

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)

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.612,1 hd fe,0.463 ,0.704 ,14.831, 4.643,2 hd fe,1.00,1.297 ,10.626, 22.994, 1 cluster se,0.38058,0.700 ,9.255, 8.285, 2 clusters se,0.765,1.803,96.958, 70.44
+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

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,29 +39,36 @@ 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
 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)

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
   ````

benchmark/result.jl

@@ -1 +1,21 @@
-using DataFrames, CSV, StatsPlotsdf = 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")
+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"))

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
 
 """

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