import numpy as np, pandas as pd, matplotlib.pyplot as plt
from scipy import stats
import seaborn as sns   # seaborn = high-level statistical plots (heatmaps, regplots, pairplots)
import statsmodels.api as sm
from statsmodels.formula.api import ols
from statsmodels.stats.outliers_influence import variance_inflation_factor
from statsmodels.nonparametric.smoothers_lowess import lowess
from sklearn.linear_model import LinearRegression, Ridge, Lasso, LogisticRegression
from sklearn.preprocessing import StandardScaler, PolynomialFeatures
from sklearn.pipeline import make_pipeline
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.metrics import r2_score, mean_squared_error
ORG="#ea580c"; DEEP="#c2410c"; LIGHT="#fdba74"; INK="#1a2138"; GRID="#e6e9f2"; GREEN="#059669"; RED="#ef4444"; AMBER="#d97706"; BLUE="#2563eb"; PUR="#9333ea"; GREY="#94a3b8"; SLATE="#475569"
plt.rcParams.update({"figure.facecolor":"white","axes.facecolor":"white","figure.dpi":110,"font.size":11,
   "axes.edgecolor":GRID,"axes.grid":True,"grid.color":GRID,"axes.axisbelow":True,"axes.spines.top":False,
   "axes.spines.right":False,"axes.titlesize":12,"axes.titleweight":"bold","legend.frameon":False})
sns.set_style("whitegrid")
BASE="https://raw.githubusercontent.com/johnfisher-ai/Statistics-Data-Science-AI-Visual-Book/main/data/"

try:
    wages = pd.read_excel('../../data/the-ols-framework--wages.xlsx', sheet_name='Workers')
except FileNotFoundError:
    wages = pd.read_excel(BASE + 'the-ols-framework--wages.xlsx', sheet_name='Workers')
print(wages.shape); wages.head(3)

(300, 4)

# EXPLORE FIRST: how wage moves with each predictor
fig, ax = plt.subplots(1,2, figsize=(11,4))
sns.regplot(data=wages, x='education', y='wage', ax=ax[0], scatter_kws=dict(s=18,color=INK,alpha=0.4), line_kws=dict(color=DEEP))
sns.regplot(data=wages, x='experience', y='wage', ax=ax[1], scatter_kws=dict(s=18,color=INK,alpha=0.4), line_kws=dict(color=DEEP))
ax[0].set_title('wage vs education'); ax[1].set_title('wage vs experience')
plt.tight_layout(); plt.show()

# Build the design matrix X (with a column of 1s) and solve the normal equations
X = np.column_stack([np.ones(len(wages)), wages['education'], wages['experience']])
y = wages['wage'].values
beta_hand = np.linalg.solve(X.T @ X, X.T @ y)          # (X'X)^-1 X'y, done stably
model = ols('wage ~ education + experience', wages).fit()
print('by hand   :', np.round(beta_hand, 4))
print('statsmodels:', np.round(model.params.values, 4))
print('identical? ', np.allclose(beta_hand, model.params.values))
print(f'R-squared  : {model.rsquared:.3f}')

by hand   : [2.415  2.3758 0.3193]
statsmodels: [2.415  2.3758 0.3193]
identical?  True
R-squared  : 0.570

H = X @ np.linalg.inv(X.T @ X) @ X.T            # the hat matrix
yhat = H @ y
e = y - yhat
print('y-hat = Hy matches model.fittedvalues? ', np.allclose(yhat, model.fittedvalues))
print('X.T @ e (should be ~0):', np.round(X.T @ e, 6))
print(f'trace(H) = {np.trace(H):.3f}  (equals the number of parameters, {X.shape[1]})')
SST = ((y-y.mean())**2).sum(); SSE = (e**2).sum(); SSR = ((yhat-y.mean())**2).sum()
print(f'SST = SSR + SSE ?  {SST:.1f} = {SSR:.1f} + {SSE:.1f}  ->  {np.isclose(SST, SSR+SSE)}')

y-hat = Hy matches model.fittedvalues?  True
X.T @ e (should be ~0): [-0. -0. -0.]
trace(H) = 3.000  (equals the number of parameters, 3)
SST = SSR + SSE ?  49392.2 = 28170.8 + 21221.4  ->  True

rng = np.random.default_rng(0)
true_slope = 3.0; n = 40
ols_slopes, twopt_slopes = [], []
for _ in range(3000):
    x = np.linspace(0, 10, n)
    yy = 2 + true_slope*x + rng.normal(0, 4, n)
    b = np.polyfit(x, yy, 1)[0]                          # OLS slope
    tp = (yy[-1] - yy[0]) / (x[-1] - x[0])               # 'two-point' unbiased estimator
    ols_slopes.append(b); twopt_slopes.append(tp)
ols_slopes = np.array(ols_slopes); twopt_slopes = np.array(twopt_slopes)
fig, ax = plt.subplots(figsize=(7.6,4.4))
ax.hist(twopt_slopes, bins=40, color=GREY, alpha=0.5, label=f'two-point  (sd={twopt_slopes.std():.2f})')
ax.hist(ols_slopes,  bins=40, color=ORG, alpha=0.7, label=f'OLS  (sd={ols_slopes.std():.2f})')
ax.axvline(true_slope, color=GREEN, ls='--', lw=2, label='true slope')
ax.set(title='Both unbiased, but OLS has the smaller variance (BLUE)', xlabel='estimated slope', ylabel='count'); ax.legend()
plt.tight_layout(); plt.show()
print(f'means: OLS {ols_slopes.mean():.2f}, two-point {twopt_slopes.mean():.2f}  (both ~ {true_slope})')

means: OLS 3.00, two-point 2.98  (both ~ 3.0)

from statsmodels.stats.diagnostic import het_breuschpagan
bp = het_breuschpagan(model.resid, model.model.exog)[1]
fig, ax = plt.subplots(figsize=(7.4,4.2))
ax.scatter(model.fittedvalues, model.resid, s=20, color=ORG, alpha=0.5); ax.axhline(0, color=INK, lw=1)
ax.set(title=f'Residuals fan out -> heteroscedastic (Breusch-Pagan p={bp:.4f})', xlabel='fitted wage', ylabel='residual')
plt.tight_layout(); plt.show()

robust = ols('wage ~ education + experience', wages).fit(cov_type='HC3')
cmp = pd.DataFrame({'coef': model.params.round(3),
                    'SE_classical': model.bse.round(3),
                    'SE_HC3_robust': robust.bse.round(3)})
print(cmp)
print('\nCoefficients are identical; only the standard errors (and thus p-values, CIs) change.')

             coef  SE_classical  SE_HC3_robust
Intercept   2.415         2.060          1.764
education   2.376         0.127          0.134
experience  0.319         0.050          0.051

Coefficients are identical; only the standard errors (and thus p-values, CIs) change.

	worker_id	education	experience	wage
0	W1000	20	8	47.76
1	W1001	20	16	66.54
2	W1002	16	22	41.16

The OLS Framework

The OLS framework, from formula to failure mode¶