import wooldridge as woo

import pandas as pd

import statsmodels.formula.api as smf

k401ksubs = woo.dataWoo('401ksubs')

# subsetting data:

k401ksubs_sub = k401ksubs[k401ksubs['fsize'] == 1]

# OLS (only for singles, i.e. 'fsize'==1):

reg_ols = smf.ols(formula='nettfa ~ inc + I((age-25)**2) + male + e401k',

data=k401ksubs_sub)

results_ols = reg_ols.fit(cov_type='HC0')

# print regression table:

table_ols = pd.DataFrame({'b': round(results_ols.params, 4),

'se': round(results_ols.bse, 4),

't': round(results_ols.tvalues, 4),

'pval': round(results_ols.pvalues, 4)})

print(f'table_ols: \n{table_ols}\n')

# WLS:

wls_weight = list(1 / k401ksubs_sub['inc'])

reg_wls = smf.wls(formula='nettfa ~ inc + I((age-25)**2) + male + e401k',

weights=wls_weight, data=k401ksubs_sub)

results_wls = reg_wls.fit()

# print regression table:

table_wls = pd.DataFrame({'b': round(results_wls.params, 4),

'se': round(results_wls.bse, 4),

't': round(results_wls.tvalues, 4),

'pval': round(results_wls.pvalues, 4)})

print(f'table_wls: \n{table_wls}\n')