Medical Charges: Take It Further¶
This notebook continues the Chapter 106 case study. We recap the regularized cost model, then work the five extensions from the chapter's Take It Further section: ridge vs lasso on the collinear body-size trio, how unstable plain OLS becomes, tuning the elastic-net mix, a BMI-by-smoker interaction, and pricing a single patient. Each is visual and explained, closing with what the model says in plain terms and the Step 12 plain-English write-up.
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="#059669"; DEEP="#047857"; LIGHT="#6ee7b7"; EM="#059669"; EMDEEP="#047857"; 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/"
import warnings; warnings.filterwarnings('ignore')
import statsmodels.api as sm, statsmodels.formula.api as smf
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LinearRegression, Ridge, Lasso, RidgeCV, LassoCV, ElasticNetCV
pd.set_option('display.max_columns', 40)
def load_clean():
try: raw = pd.read_csv('../../data/medical_charges.csv')
except FileNotFoundError: raw = pd.read_csv(BASE + 'medical_charges.csv')
df = raw.dropna(subset=['charges']).drop_duplicates(subset='patient_id').copy()
df['charges'] = df['charges'].astype(float)
df['sex'] = df['sex'].str.strip().str.upper().map({'MALE':'Male','M':'Male','FEMALE':'Female','F':'Female'})
df['smoker'] = df['smoker'].astype(str).str.strip().str.upper().map({'YES':1,'Y':1,'1':1,'NO':0,'N':0,'0':0})
df['region'] = df['region'].astype(str).str.strip().str.lower()
df['bmi'] = pd.to_numeric(df['bmi'], errors='coerce')
df['bmi'] = df['bmi'].fillna(df['weight_kg'] / (df['height_cm']/100)**2)
df['log_charges'] = np.log(df['charges'])
return df
df = load_clean()
feats = ['age','bmi','height_cm','weight_kg','children','smoker','exercise_hours_wk','diet_score','blood_pressure','cholesterol','income_k','noise_a','noise_b']
X = pd.get_dummies(df[feats].join(df['region']), columns=['region'], drop_first=True).astype(float)
scaler = StandardScaler().fit(X); Xs = scaler.transform(X); y = df['log_charges'].values
lasso = LassoCV(cv=5, random_state=0).fit(Xs, y)
print(f'n={len(df)} features={X.shape[1]} lasso R2={lasso.score(Xs,y):.3f} alpha={lasso.alpha_:.4f}')
n=994 features=16 lasso R2=0.865 alpha=0.0087
trio = ['weight_kg','height_cm','bmi']
idx = [list(X.columns).index(c) for c in trio]
r = Ridge(alpha=10).fit(Xs, y); l = Lasso(alpha=lasso.alpha_).fit(Xs, y)
cmp = pd.DataFrame({'ridge': r.coef_[idx], 'lasso': l.coef_[idx]}, index=trio).round(3)
print(cmp.to_string())
fig, ax = plt.subplots(figsize=(7.2,4)); w=0.36; pos=np.arange(3)
ax.bar(pos-w/2, cmp.ridge, w, color=DEEP, label='ridge (shrinks, keeps all)')
ax.bar(pos+w/2, cmp.lasso, w, color=EM, label='lasso (picks one)')
ax.axhline(0,color=INK,lw=1); ax.set_xticks(pos); ax.set_xticklabels(trio)
ax.set(title='Ridge vs lasso on the collinear body-size trio', ylabel='standardized coefficient'); ax.legend()
plt.tight_layout(); plt.show()
ridge lasso weight_kg 0.024 0.000 height_cm -0.013 0.000 bmi 0.096 0.109
What we learn. The two penalties treat redundancy differently. Ridge keeps all three features but shrinks them, splitting the shared body-size signal across weight, height, and BMI. Lasso picks one representative (BMI) and zeros the other two. Neither is wrong, but the lasso's choice gives a shorter, more interpretable model, one clean body-size term instead of three tangled ones.
Xtrio = sm.add_constant(pd.DataFrame(Xs, columns=X.columns)[trio])
ols_trio = sm.OLS(y, Xtrio).fit()
se = ols_trio.bse[trio]
print('Unpenalized OLS on the collinear trio:')
print(pd.DataFrame({'coef':ols_trio.params[trio].round(2),'std_err':se.round(2),'p_value':ols_trio.pvalues[trio].round(3)}).to_string())
print(f'\nStandard errors are huge and coefficients flip sign -> collinearity has made them meaningless.')
Unpenalized OLS on the collinear trio:
coef std_err p_value
weight_kg 0.13 0.27 0.625
height_cm -0.05 0.13 0.725
bmi 0.01 0.23 0.978
Standard errors are huge and coefficients flip sign -> collinearity has made them meaningless.
What we learn. Feeding all three collinear features to unpenalized OLS produces enormous standard errors and coefficients that swing wildly in sign and size, none of them individually significant, even though together they clearly matter. This is exactly the instability VIF warned about in the chapter, and precisely what regularization cures by shrinking or dropping the redundant terms.
rows = []
for r in [0.1, 0.3, 0.5, 0.7, 0.9, 1.0]:
m = ElasticNetCV(l1_ratio=[r], cv=5, random_state=0).fit(Xs, y)
rows.append({'l1_ratio': r, 'features_kept': int((np.abs(m.coef_)>1e-6).sum()), 'R2': round(m.score(Xs,y),3)})
res = pd.DataFrame(rows); print(res.to_string(index=False))
fig, ax = plt.subplots(figsize=(7.2,4))
ax.plot(res.l1_ratio, res.features_kept, 'o-', color=EM, lw=2.4)
ax.set(title='More lasso (higher l1_ratio) -> fewer features kept', xlabel='l1_ratio (0=ridge, 1=lasso)', ylabel='features kept')
for _,row in res.iterrows(): ax.annotate(f"R2={row.R2}", (row.l1_ratio, row.features_kept), textcoords='offset points', xytext=(0,8), fontsize=8, ha='center')
plt.tight_layout(); plt.show()
l1_ratio features_kept R2
0.1 15 0.865
0.3 12 0.865
0.5 11 0.865
0.7 10 0.865
0.9 10 0.865
1.0 10 0.865
What we learn. As l1_ratio moves from ridge (0) toward lasso (1), the model keeps fewer features while the R-squared barely changes. Elastic net lets you dial the trade-off: pure ridge keeps everything (shrunk), pure lasso is most aggressive at selection, and a middle value is handy when several correlated features should be kept as a group rather than one picked arbitrarily.
d = df.copy(); d['bmi_z'] = (d.bmi - d.bmi.mean())/d.bmi.std()
mi = smf.ols('log_charges ~ bmi_z * smoker + age + blood_pressure + cholesterol', d).fit()
print(mi.summary2().tables[1].loc[['bmi_z','smoker','bmi_z:smoker']].round(3).to_string())
fig, ax = plt.subplots(figsize=(7.2,4)); grid=np.linspace(-2,3,50)
for s,c,lab in [(0,EM,'non-smoker'),(1,RED,'smoker')]:
pr = mi.params['Intercept']+mi.params['bmi_z']*grid+mi.params['smoker']*s+mi.params['bmi_z:smoker']*grid*s+mi.params['age']*d.age.mean()+mi.params['blood_pressure']*d.blood_pressure.mean()+mi.params['cholesterol']*d.cholesterol.mean()
ax.plot(grid, np.exp(pr), color=c, lw=2.4, label=lab)
ax.set(title='Cost vs BMI, by smoking status', xlabel='BMI (standardized)', ylabel='predicted charges'); ax.legend()
ax.yaxis.set_major_formatter(plt.FuncFormatter(lambda v,_: f'${v/1000:.0f}k')); plt.tight_layout(); plt.show()
Coef. Std.Err. t P>|t| [0.025 0.975] bmi_z 0.119 0.012 10.012 0.000 0.096 0.143 smoker 1.435 0.027 53.584 0.000 1.382 1.487 bmi_z:smoker 0.000 0.027 0.002 0.998 -0.053 0.053
What we learn. The interaction term tests whether the BMI slope is steeper for smokers. The interaction is essentially zero (coefficient 0.000, p = 0.998), so BMI has the same proportional effect for smokers and non-smokers. The two fitted lines still diverge in dollars, because the model is additive on the log scale and a constant multiplier widens the dollar gap as cost rises. We tested the idea rather than assuming it, and the data give no reason to carry the interaction into the pricing model.
profile = {'age':45,'bmi':31,'height_cm':175,'weight_kg':95,'children':2,'smoker':1,
'exercise_hours_wk':2,'diet_score':5,'blood_pressure':135,'cholesterol':210,
'income_k':60,'noise_a':50,'noise_b':50}
row = pd.DataFrame([profile]).join(pd.DataFrame([{'region':'southeast'}]))
row = pd.get_dummies(row, columns=['region']).reindex(columns=X.columns, fill_value=0).astype(float)
pt = scaler.transform(row)
resid_sd = (y - lasso.predict(Xs)).std()
logpred = lasso.predict(pt)[0]
lo, hi = np.exp(logpred-1.96*resid_sd), np.exp(logpred+1.96*resid_sd)
print(f'Predicted annual charge (median): ${np.exp(logpred):,.0f}')
print(f'95% range: ${lo:,.0f} to ${hi:,.0f}')
# same patient as a non-smoker
row2=row.copy(); row2['smoker']=0; ns=np.exp(lasso.predict(scaler.transform(row2))[0])
print(f'\nSame patient as a NON-smoker: ${ns:,.0f} (the model predicts ~${np.exp(logpred)-ns:,.0f} more for the smoker)')
Predicted annual charge (median): $63,538 95% range: $33,995 to $118,754 Same patient as a NON-smoker: $15,332 (the model predicts ~$48,206 more for the smoker)
What we learn. The model becomes a pricing tool: feed one patient's profile and it returns a typical (median) annual charge with an approximate range (never a false-precision single number); exponentiating a log prediction gives the median, and the range uses the in-sample residual spread only. Flipping just the smoker flag shows how much higher the model predicts for a smoker with this profile, an association in observational data, the actionable, plain-English payoff of the whole analysis.
What the model says, in plain terms¶
- Smoking is the number-one cost driver by a wide margin, and toggling it for one patient shows a large dollar jump.
- The body-size signal is one thing, not three: ridge spread it across height, weight, and BMI; lasso cleanly kept BMI and dropped the rest.
- Plain OLS was untrustworthy on the collinear features (huge, sign-flipping coefficients); regularization fixed that at no cost in accuracy.
- The elastic-net dial trades feature count against a tiny bit of fit, and the BMI-by-smoking interaction is essentially zero (p = 0.998).
Step 12 - the plain-English write-up¶
What we did. We took the cleaned insurance file and stress-tested the cost model five ways: compared two ways of handling redundant body-size measures, showed why the naive approach fails, tuned the selection dial, checked whether weight matters more for smokers, and produced an actual price quote for a sample patient.
How good is it? The model still explains about 87% of cost differences, and the extensions confirmed it is stable and trustworthy, not an accident of the particular features we picked.
What drives cost: smoking first, by a lot, then age, blood pressure, cholesterol, and body-mass index, with the same proportional BMI effect for smokers and non-smokers.
Bottom line: a short, regularized model gives an honest, explainable price, one body-size measure, a handful of health factors, a clear range, and a smoking flag that is the single biggest lever on the number.