rust-accelerated for benchmarks

Author: s3alfisc

pyfixest/estimation/backends.py

@@ -1,6 +1,7 @@
 from pyfixest.core.collinear import find_collinear_variables
 from pyfixest.core.crv1 import crv1_meat_loop
 from pyfixest.core.demean import demean
+from pyfixest.core.demean_accelerated import demean_accelerated
 from pyfixest.core.nested_fixed_effects import count_fixef_fully_nested_all
 from pyfixest.estimation.demean_ import demean as demean_nb
 from pyfixest.estimation.numba.find_collinear_variables_nb import (
@@ -58,6 +59,12 @@
         "crv1_meat": crv1_meat_loop,
         "nonnested": count_fixef_fully_nested_all,
     },
+    "rust-accelerated": {
+        "demean": demean_accelerated,
+        "collinear": find_collinear_variables,
+        "crv1_meat": crv1_meat_loop,
+        "nonnested": count_fixef_fully_nested_all,
+    },
     "jax": {
         "demean": demean_jax_fn,
         "collinear": find_collinear_variables_jax,

pyfixest/estimation/demean_.py

@@ -349,6 +349,10 @@ def _set_demeaner_backend(
         from pyfixest.core.demean import demean as demean_rs
 
         return demean_rs
+    elif demeaner_backend == "rust-accelerated":
+        from pyfixest.core.demean_accelerated import demean_accelerated
+
+        return demean_accelerated
     elif demeaner_backend == "numba":
         return demean
     elif demeaner_backend == "jax":

pyfixest/estimation/literals.py

@@ -12,7 +12,7 @@
     "jax",
 ]
 DemeanerBackendOptions = Literal[
-    "numba", "jax", "rust", "cupy", "cupy32", "cupy64", "scipy"
+    "numba", "jax", "rust", "rust-accelerated", "cupy", "cupy32", "cupy64", "scipy"
 ]
 PredictionErrorOptions = Literal["prediction"]
 QuantregMethodOptions = Literal["fn", "pfn"]