データフレームを得て VIF を計算する関数

ChatGPT 3.5 に「VIF を計算する関数を書いて」とお願いすると，すぐに書いてくれるけど，for を使ったり，statsmodels 使ったりして，冗長だなあ。結果だって，別にデータフレームで返してくれなくてよい。

import pandas as pd
import numpy as np
from statsmodels.stats.outliers_influence import variance_inflation_factor
from statsmodels.tools.tools import add_constant

def calculate_vif(data):
    data_with_const = add_constant(data)
    vif_values = np.zeros(data_with_const.shape[1])
    for i in range(data_with_const.shape[1]):
        vif_values[i] = variance_inflation_factor(data_with_const.values, i)
    variables = data_with_const.columns
    vif_data = pd.DataFrame({'Variable': variables, 'VIF': vif_values})
    vif_data = vif_data[vif_data["Variable"] != "const"]
    return vif_data

data = {
    'x1': np.random.rand(100),
    'x2': np.random.rand(100),
    'x3': np.random.rand(100),
    'y': np.random.rand(100)
}
df = pd.DataFrame(data)

vif_result = calculate_vif(df)
print(vif_result)

  Variable       VIF
1       x1  1.038079
2       x2  1.002299
3       x3  1.081315
4        y  1.049966

関数自体は，一行で書ける。
データフレームを渡すと VIF をベクトルで返す。

import pandas as pd
import numpy as np
from scipy import linalg

def VIF(df): return np.diag(linalg.inv(df.corr()))
print(VIF(df))

[1.03807931 1.00229873 1.08131533 1.04996588]