Skip to content

Commit

Permalink
jit compile multicollinearity detection #466 (#579)
Browse files Browse the repository at this point in the history
  • Loading branch information
s3alfisc authored Aug 11, 2024
1 parent 72e1967 commit 1e221ef
Showing 1 changed file with 6 additions and 4 deletions.
10 changes: 6 additions & 4 deletions pyfixest/estimation/feols_.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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]:
Expand All @@ -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]

Expand All @@ -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

Expand All @@ -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):
Expand Down

0 comments on commit 1e221ef

Please sign in to comment.