generate_plots.py

"""Regenerate the 4 portfolio plots from the notebook's modeling pipeline."""
import os
import re
import sys
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import r2_score, mean_squared_error, mean_absolute_error

REPO = os.path.dirname(os.path.abspath(__file__))
CSV = os.path.join(REPO, "analytic_data2025_v2.csv")
ASSETS = os.path.join(REPO, "assets")
os.makedirs(ASSETS, exist_ok=True)

DPI = 150
WIDTH_PX = 1400
FIG_W = WIDTH_PX / DPI  # ~9.33 inches at 150 DPI

# ---------- Load + filter (replicates notebook cells 1-5) ----------
df = pd.read_csv(CSV, low_memory=False)
cols = set(df.columns)
state_abbr_col  = "state" if "state" in cols else "State Abbreviation"
county_fips_col = "countycode" if "countycode" in cols else "County FIPS Code"

df[county_fips_col] = pd.to_numeric(df[county_fips_col], errors="coerce").fillna(0).astype(int)
not_national = df[state_abbr_col] != "US"
df_counties = df[(df[county_fips_col] != 0) & (not_national)].copy()
print(f"counties: {df_counties.shape}", flush=True)

# ---------- Build df_model (replicates cell 3) ----------
def norm(s):
    s = s.lower()
    return re.sub(r"[\s\-_/:%()]+", "", s)

cols_norm = {c: norm(c) for c in df_counties.columns}

def find_col(*keywords):
    k = [norm(kw) for kw in keywords]
    for c, cn in cols_norm.items():
        if all(kw in cn for kw in k):
            return c
    return None

target_name = find_col("Premature Death", "raw value")
wanted = [
    ("Access to Exercise Opportunities", "raw value"),
    ("Flu Vaccinations", "raw value"),
    ("Food Environment Index", "raw value"),
    ("Preventable Hospital Stays", "raw value"),
    ("Uninsured", "raw value"),
    ("Primary Care Physicians", "raw value"),
    ("Mammography Screening", "raw value"),
    ("Mental Health Providers", "raw value"),
    ("Dentists", "raw value"),
    ("Air Pollution", "Particulate Matter", "raw value"),
    ("Drinking Water Violations", "raw value"),
    ("Broadband Access", "raw value"),
    ("Library Access", "raw value"),
    ("Severe Housing Problems", "raw value"),
    ("Driving Alone to Work", "raw value"),
    ("Long Commute", "Driving Alone", "raw value"),
    ("High School Completion", "raw value"),
    ("Some College", "raw value"),
    ("Unemployment", "raw value"),
    ("Children in Poverty", "raw value"),
    ("Income Inequality", "raw value"),
    ("Child Care Cost Burden", "raw value"),
    ("Injury Deaths", "raw value"),
    ("Social Associations", "raw value"),
]
predictor_cols = [find_col(*kw) for kw in wanted]
predictor_cols = [c for c in predictor_cols if c]
print(f"predictors found: {len(predictor_cols)}", flush=True)

df_model = df_counties[predictor_cols + [target_name]].copy()
# numeric coerce (cell 11)
df_model = df_model.apply(lambda s: pd.to_numeric(s.astype(str).str.replace(",", "", regex=False), errors="coerce"))
df_model = df_model.fillna(df_model.median(numeric_only=True))
print(f"df_model: {df_model.shape}", flush=True)

# ---------- Train RF (cells 29, 32) ----------
X = df_model.drop(columns=[target_name])
y = df_model[target_name]
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
rf = RandomForestRegressor(n_estimators=300, max_depth=None, random_state=42, n_jobs=-1)
rf.fit(X_train, y_train)
y_pred = rf.predict(X_test)
r2 = r2_score(y_test, y_pred)
rmse = float(np.sqrt(mean_squared_error(y_test, y_pred)))
mae = mean_absolute_error(y_test, y_pred)
print(f"RF R2={r2:.3f} RMSE={rmse:.0f} MAE={mae:.0f}", flush=True)

# ---------- Plot 1: correlation heatmap (replace the smaller notebook one) ----------
plt.figure(figsize=(FIG_W, FIG_W * 0.86))
corr = df_model.corr()
short_labels = [c.replace(" raw value", "") for c in corr.columns]
sns.heatmap(corr, cmap="RdBu_r", center=0, annot=False, fmt=".1f",
            linewidths=0.5, xticklabels=short_labels, yticklabels=short_labels,
            cbar_kws={"shrink": 0.7})
plt.title("Correlation Matrix: 24 Health Predictors + Premature Death", fontsize=13)
plt.xticks(rotation=60, ha="right", fontsize=8)
plt.yticks(fontsize=8)
plt.tight_layout()
out = os.path.join(ASSETS, "correlation-heatmap.png")
plt.savefig(out, dpi=DPI, bbox_inches="tight", facecolor="white")
plt.close()
print(f"wrote {out}", flush=True)

# ---------- Plot 2: predicted vs actual ----------
plt.figure(figsize=(FIG_W, FIG_W * 0.75))
plt.scatter(y_test, y_pred, alpha=0.4, s=18, c="#2563eb", edgecolors="none")
lo, hi = float(y_test.min()), float(y_test.max())
plt.plot([lo, hi], [lo, hi], "r--", lw=2, label="y = x")
plt.xlabel("Actual Premature Death Rate (YPLL per 100K)")
plt.ylabel("Predicted Premature Death Rate (YPLL per 100K)")
plt.title(f"Random Forest: Predicted vs Actual (R² = {r2:.3f})")
plt.legend(loc="upper left")
plt.grid(alpha=0.25)
plt.tight_layout()
out = os.path.join(ASSETS, "predicted-vs-actual.png")
plt.savefig(out, dpi=DPI, bbox_inches="tight", facecolor="white")
plt.close()
print(f"wrote {out}", flush=True)

# ---------- Plot 3: target distribution ----------
plt.figure(figsize=(FIG_W, FIG_W * 0.46))
national_avg = float(y.mean())
plt.hist(y, bins=60, color="#2563eb", edgecolor="white")
plt.axvline(x=national_avg, color="red", linestyle="--", lw=2,
            label=f"County mean: {national_avg:,.0f}")
plt.xlabel("Years of Potential Life Lost per 100K")
plt.ylabel("Number of Counties")
plt.title(f"Distribution of Premature Death Rates Across {len(y):,} U.S. Counties")
plt.legend()
plt.grid(alpha=0.25)
plt.tight_layout()
out = os.path.join(ASSETS, "target-distribution.png")
plt.savefig(out, dpi=DPI, bbox_inches="tight", facecolor="white")
plt.close()
print(f"wrote {out}", flush=True)

# ---------- Plot 4: RF feature importance ----------
importances = pd.Series(rf.feature_importances_, index=X.columns).sort_values(ascending=True)
short_idx = [c.replace(" raw value", "") for c in importances.index]
plt.figure(figsize=(FIG_W, FIG_W * 0.75))
plt.barh(short_idx, importances.values, color="#2563eb")
plt.xlabel("Feature Importance (mean decrease in impurity)")
plt.title("Random Forest Feature Importance — 24 Health Predictors")
plt.grid(axis="x", alpha=0.25)
plt.tight_layout()
out = os.path.join(ASSETS, "rf-feature-importance.png")
plt.savefig(out, dpi=DPI, bbox_inches="tight", facecolor="white")
plt.close()
print(f"wrote {out}", flush=True)

# ---------- Plot 5: residuals ----------
residuals = y_test - y_pred
fig, axes = plt.subplots(1, 2, figsize=(FIG_W, FIG_W * 0.42))
axes[0].scatter(y_pred, residuals, alpha=0.4, s=18, c="#2563eb", edgecolors="none")
axes[0].axhline(0, color="red", linestyle="--", lw=1.5)
axes[0].set_xlabel("Predicted Premature Death Rate")
axes[0].set_ylabel("Residual (Actual - Predicted)")
axes[0].set_title("Residuals vs Predicted")
axes[0].grid(alpha=0.25)
axes[1].hist(residuals, bins=50, color="#2563eb", edgecolor="white")
axes[1].axvline(0, color="red", linestyle="--", lw=1.5)
axes[1].set_xlabel("Residual")
axes[1].set_ylabel("Number of Counties")
axes[1].set_title("Residual Distribution")
axes[1].grid(alpha=0.25)
plt.suptitle(f"Random Forest Residual Analysis (RMSE = {rmse:,.0f})", fontsize=13)
plt.tight_layout()
out = os.path.join(ASSETS, "residuals.png")
plt.savefig(out, dpi=DPI, bbox_inches="tight", facecolor="white")
plt.close()
print(f"wrote {out}", flush=True)

print("done", flush=True)