jit compile multicollinearity detection #466 (#579)

Author: s3alfisc

pyfixest/estimation/feols_.py

@@ -4,6 +4,7 @@
 from importlib import import_module
 from typing import Optional, Union
 
+import numba as nb
 import numpy as np
 import pandas as pd
 import polars as pl
@@ -2190,6 +2191,7 @@ def _drop_multicollinear_variables(
     return X, names_array.tolist(), collin_vars, collin_index
 
 
+@nb.njit(parallel=False)
 def _find_collinear_variables(
     X: np.ndarray, tol: float = 1e-10
 ) -> tuple[np.ndarray, int, bool]:
@@ -2215,7 +2217,7 @@ def _find_collinear_variables(
     """
     K = X.shape[1]
     R = np.zeros((K, K))
-    id_excl = np.zeros(K, dtype=bool)
+    id_excl = np.zeros(K, dtype=np.int32)
     n_excl = 0
     min_norm = X[0, 0]
 
@@ -2228,11 +2230,11 @@ def _find_collinear_variables(
 
         if R_jj < tol:
             n_excl += 1
-            id_excl[j] = True
+            id_excl[j] = 1
 
             if n_excl == K:
                 all_removed = True
-                return id_excl, n_excl, all_removed
+                return id_excl.astype(np.bool_), n_excl, all_removed
 
             continue
 
@@ -2250,7 +2252,7 @@ def _find_collinear_variables(
                 value -= R[k, i] * R[k, j]
             R[j, i] = value / R_jj
 
-    return id_excl, n_excl, False
+    return id_excl.astype(np.bool_), n_excl, False
 
 
 def _check_vcov_input(vcov: Union[str, dict[str, str]], data: pd.DataFrame):