import pyfixest as pf
import numpy as np
import pandas as pd
import time
rng = np.random.default_rng(737)
N = 10_000_000
benchmark_data = pd.DataFrame({
"Y": rng.normal(0, 1, N),
"X1": rng.normal(0, 1, N),
"X2": rng.normal(0, 1, N),
"X3": rng.normal(0, 1, N),
"f1": rng.integers(0, 10_000, N),
"f2": rng.integers(0, 1_000, N),
"f3": rng.integers(0, 10, N),
})
# burn-in for numba
fit_nb_warmup = pf.feols(
fml = "Y ~ X1 + X2 + X3 | f1 + f2 + f3", data = benchmark_data[:100_000]
)
# benchmark for numba backend
tic = time.time()
fit_nb = pf.feols(
fml = "Y ~ X1 + X2 + X3 | f1 + f2 + f3", data = benchmark_data
)
toc = time.time()
print(f"Numba backend took {toc-tic}.")
# benchmark for rust backend
tic = time.time()
fit_rust = pf.feols(
fml = "Y ~ X1 + X2 + X3 | f1 + f2 + f3", data = benchmark_data,
demeaner_backend = "rust"
)
toc = time.time()
print(f"Rust backend took {toc-tic}.")