Improving the performance of PyFixest Demeaning Algorithm on both CPU and GPU

[s3alfisc]

Background:

For fast fitting of fixed effects regression models of the form Y = X'beta + D'alpha + u, PyFixest relies on the Frisch-Waugh-Lovell Theorem.

At its core, the FWL approach relies on "demeaning" both dependent variable Y and fixed effects X by the fixed effects fe. In PyFixest, this logic is implemented in the demean function.

Via multiple dispatch, demean allows users to choose from multiple algorithms:

• numba and rust implement the standard MAP algorithm (without any accelerations) in the respective language. See here for details
• scipy and cupy run the demeaning algorithm via the LSMR algorithm on CPU or GPU
• jax runs the MAP algorithm on the GPU in Jax

To do

We would like to explore the following:

• Implement all acceleration tricks and tweaks implemented in fixes. See here for details - both in Rust and optionally in JAX for GPU support
• Implement support for varying slopes
• Experiment with Graph Theoretical solvers - see here for details
• Provide unified convergence criteria for all demeaner backends (LSMR currently deviates from MAP)
• Consider writing a standalone demeaner package with bindings to other languages (Python, R, Rust)
• Write excellent documentation on all approaches