# Weighted Least Squares
# Module 8 - Weighted Least Squares Intuition
# Dataset: MODULE8_INCOME_SAVINGS_SYNTHETIC

# Weighted Least Squares
#
# Module 8 notebook lab. This notebook uses original Ceteris Lab teaching data and does not report real empirical findings.

# Learning goal
# Choose inverse-variance weights and compare OLS with WLS.
#
# Dataset: MODULE8_INCOME_SAVINGS_SYNTHETIC. Variables: household_id, income, savings, age, education, family_size.

# %% Cell 3
import pandas as pd
import statsmodels.api as sm

df = pd.read_csv("/data/module-8/module8_income_savings_heteroskedastic.csv")
X = sm.add_constant(df[["income", "age", "education"]])
weights = 1 / (df["error_scale"] ** 2)
ols = sm.OLS(df["savings"], X).fit(cov_type="HC1")
wls = sm.WLS(df["savings"], X, weights=weights).fit()
print("Robust OLS coefficients")
print(ols.params.round(4))
print("WLS coefficients")
print(wls.params.round(4))

# Reflection
# Write two sentences: one sentence explaining what the diagnostic or robust result says, and one sentence explaining a limitation or next step.