diff --git a/pyfixest/estimation/feols_.py b/pyfixest/estimation/feols_.py index 17cbc4a9..54605880 100644 --- a/pyfixest/estimation/feols_.py +++ b/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):