vif: Allow computing VIF only for variables in the formula

This commit is contained in:
RunasSudo 2022-10-14 14:08:51 +11:00
parent 461e00df78
commit bb534cb285
Signed by: RunasSudo
GPG Key ID: 7234E476BF21C61A
1 changed files with 10 additions and 2 deletions

View File

@ -28,8 +28,16 @@ import itertools
from .utils import Estimate, check_nan, fmt_p_html, fmt_p_text
def vif(df, nan_policy='warn'):
"""Calculate the variance inflation factor for each variable in df"""
def vif(df, formula=None, nan_policy='warn'):
"""
Calculate the variance inflation factor for each variable in df
formula: If specified, calculate the VIF only for the variables in the formula
"""
if formula:
# Only consider columns in the formula
df = df[cols_for_formula(formula)]
# Check for/clean NaNs
df = check_nan(df, nan_policy)