Repository: kutaslab/fitgrid

Issue:
add VIF calculator

Discussion:

Patch:
diff --git a/fitgrid/utils/__init__.py b/fitgrid/utils/__init__.py
index 404d977..2328b30 100644
--- a/fitgrid/utils/__init__.py
+++ b/fitgrid/utils/__init__.py
@@ -1,1 +1,1 @@
-from . import lmer
+from . import lm, lmer
diff --git a/fitgrid/utils/lm.py b/fitgrid/utils/lm.py
new file mode 100644
index 0000000..60ecd54
--- /dev/null
+++ b/fitgrid/utils/lm.py
@@ -0,0 +1,20 @@
+import pandas as pd
+import patsy
+from tqdm import tqdm
+from statsmodels.stats.outliers_influence import (
+    variance_inflation_factor as vif,
+)
+
+
+def get_vifs(epochs, RHS):
+    def get_single_vif(group, RHS):
+        dmatrix = patsy.dmatrix(formula_like=RHS, data=group)
+        vifs = {
+            name: vif(dmatrix, index)
+            for name, index in dmatrix.design_info.column_name_indexes.items()
+        }
+        return pd.Series(vifs)
+
+    tqdm.pandas(desc="Time")
+
+    return epochs._snapshots.progress_apply(get_single_vif, RHS=RHS)


Test patch:
diff --git a/tests/test_utils_lm.py b/tests/test_utils_lm.py
new file mode 100644
index 0000000..5df68a6
--- /dev/null
+++ b/tests/test_utils_lm.py
@@ -0,0 +1,8 @@
+from .context import fitgrid
+
+
+def test_smoke_get_vifs():
+
+    epochs = fitgrid.generate()
+    RHS = 'continuous + categorical'
+    fitgrid.utils.lm.get_vifs(epochs, RHS)