An epidemiological study: Predicting 5-year mortality

The NHANES study has been runs dietary interviews and blood tests for several decade to measure the nutritional status of U.S. adults and children and relate it to health.

Here we use models to predict 5-year mortality from the NHANES covariates. We use both a simple model (a linear, that is a classic tool in epidemiology) and a more flexible machine learning model.

We show how to understand these models within the context of our specific question: mortality prediction. For this, we inspect features importance, and look at partial dependence plots to understand how the model uses covariates to predict.

Learning objectives and take home messages

This notebook introduces a dataset and can be seen as a baseline for the later notebook on survival analysis.

It can be skipped unless you have a particular interest in this type of datasets.

Load NHANES 5-year-outcome dataset

This dataset was built from the NHANES study:

import pandas as pd

df = pd.read_csv("nhanes_1999_2018_mortality_5yr_horizon.csv")

X = df.drop(columns=["death_within_5y"])
y = df["death_within_5y"]

print("Covariates used:", X.columns.tolist())
print("\n5-year mortality rate:", y.mean())

Covariates used: ['age', 'sex', 'race_eth', 'education', 'bmi', 'waist_cm', 'systolic_bp_mmhg', 'diastolic_bp_mmhg', 'EVER_SMOKED']

5-year mortality rate: 0.07133276161576181

Model fitting and prediction

Train / test split

Before fitting models, we split train and test data, in order to have untouched hold-out data (“test”) to evaluate the model. We stratify the split, the outcome is rare (~7%).

from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, stratify=y, random_state=0
)

Fit a linear model: logistic regression

We use a LogisticRegression from scikit-learn, but wrap it in a tabular_pipeline from skrub that does some data preparation.

from sklearn.linear_model import LogisticRegression
from skrub import tabular_pipeline

model_linear = tabular_pipeline(LogisticRegression())
model_linear.fit(X_train, y_train)

Pipeline(steps=[('tablevectorizer',
                 TableVectorizer(datetime=DatetimeEncoder(periodic_encoding='spline'))),
                ('simpleimputer', SimpleImputer(add_indicator=True)),
                ('squashingscaler', SquashingScaler(max_absolute_value=5)),
                ('logisticregression', LogisticRegression())])

In a Jupyter environment, please rerun this cell to show the HTML representation or trust the notebook.
On GitHub, the HTML representation is unable to render, please try loading this page with nbviewer.org.

Evaluate the model on held-out data

from sklearn.metrics import roc_auc_score

y_pred_proba_linear = model_linear.predict_proba(X_test)[:, 1]
auc_linear = roc_auc_score(y_test, y_pred_proba_linear)
print(f"Held-out AUC, linear model (logistic regression): {auc_linear:.3f}")

Held-out AUC, linear model (logistic regression): 0.869

Fit a non-linear model: gradient boosting

This is skrub’s default classification pipeline: a TableVectorizer followed by a HistGradientBoostingClassifier.

model_nonlinear = tabular_pipeline("classifier")
model_nonlinear.fit(X_train, y_train)

Pipeline(steps=[('tablevectorizer',
                 TableVectorizer(low_cardinality=ToCategorical())),
                ('histgradientboostingclassifier',
                 HistGradientBoostingClassifier())])

In a Jupyter environment, please rerun this cell to show the HTML representation or trust the notebook.
On GitHub, the HTML representation is unable to render, please try loading this page with nbviewer.org.

Once again, evaluate it on left-out data

y_pred_proba_nonlinear = model_nonlinear.predict_proba(X_test)[:, 1]
auc_nonlinear = roc_auc_score(y_test, y_pred_proba_nonlinear)
print(f"Held-out AUC, non-linear model (gradient boosting): {auc_nonlinear:.3f}")

Held-out AUC, non-linear model (gradient boosting): 0.871

The non-linear model predicts slightly better than the linear one. Let’s now see what drives this prediction.

Model inspection: how do the models predict

Permutation importance: finding the important variables

We use permutation importance to see which variables drive the prediction

from sklearn.inspection import permutation_importance

perm_linear = permutation_importance(
    model_linear, X_test, y_test, scoring="roc_auc", n_repeats=10, random_state=0
)

importances_linear = pd.DataFrame({
    "feature": X_test.columns,
    "importance_mean": perm_linear.importances_mean,
    "importance_std": perm_linear.importances_std,
}).sort_values("importance_mean", ascending=False)

print("Permutation importance (drop in prediction performance when a feature is shuffled):")
print(importances_linear.to_string(index=False))

Permutation importance (drop in prediction performance when a feature is shuffled):
          feature  importance_mean  importance_std
              age         0.280406        0.013239
         waist_cm         0.013071        0.001735
      EVER_SMOKED         0.006360        0.001324
              bmi         0.005419        0.001247
              sex         0.004085        0.000723
         race_eth         0.003396        0.000683
        education         0.001913        0.000354
diastolic_bp_mmhg         0.001632        0.000299
 systolic_bp_mmhg         0.000371        0.000509

For the non-linear model

perm_nonlinear = permutation_importance(
    model_nonlinear, X_test, y_test, scoring="roc_auc", n_repeats=10, random_state=0
)

importances_nonlinear = pd.DataFrame({
    "feature": X_test.columns,
    "importance_mean": perm_nonlinear.importances_mean,
}).sort_values("importance_mean", ascending=False)

print("Permutation importance (drop in prediction performance when a feature is shuffled):")
print(importances_nonlinear.to_string(index=False))

Permutation importance (drop in prediction performance when a feature is shuffled):
          feature  importance_mean
              age         0.266011
         waist_cm         0.013757
              bmi         0.012261
      EVER_SMOKED         0.007410
              sex         0.005385
 systolic_bp_mmhg         0.005004
         race_eth         0.003856
        education         0.002729
diastolic_bp_mmhg         0.001337

The ordering of which variables are important for prediction are quite similar Across the two models. Let’s now understand the difference in the predictions

Partial dependence on age

Here we plot “partial dependencies”, that show how the prediction of the model for a given feature changes, on average, across the population.

For the linear model

We plot both the prediction of the model (averaged across the population), and the average observed mortality for each age bin.

import matplotlib.pyplot as plt
from sklearn.inspection import partial_dependence
sex_colors = {"M": "tab:blue", "F": "tab:orange"}
sex_markers = {"M": "o", "F": "^"}

# Plot the partial dependence of age for the whole population
pd_population_linear = partial_dependence(model_linear, X_test, features=["age"], grid_resolution=30)
plt.plot(pd_population_linear["grid_values"][0], pd_population_linear["average"][0],
        color="black", linewidth=2, label="Model prediction, averaged over whole population")

# Now plot for each sex, and overlay the average observed mortality for the corresponding sex
for sex_value in ["M", "F"]:
    is_sex = X_test["sex"] == sex_value
    pd_sex_linear = partial_dependence(model_linear, X_test[is_sex], features=["age"], grid_resolution=30)
    plt.plot(pd_sex_linear["grid_values"][0], pd_sex_linear["average"][0],
            color=sex_colors[sex_value], linewidth=2,
            label=f"Model prediction, averaged for sex = {sex_value}")
    age_sex_mean = (
        pd.DataFrame({"age": X_test.loc[is_sex, "age"], "y_test": y_test[is_sex]})
        .groupby("age", as_index=False)["y_test"].mean()
    )
    plt.plot(age_sex_mean["age"], age_sex_mean["y_test"],
            color=sex_colors[sex_value], linewidth=1, linestyle="--", marker=sex_markers[sex_value], markersize=4,
            label=f"Average 5 year mortality, sex = {sex_value}")


plt.xlabel("age")
plt.ylabel("predicted probability of death within 5 years")
plt.title("Partial dependence of age, by sex - linear model")
plt.legend()
plt.tight_layout()
plt.show()

What we see is that the model appears a bit as a “smoother” compared to the bin-wise average observed mortality. This is a good picture to have in mind for machine learning.

This “smoothing” is a tradeoff to have in mind. More smoothing removes noise, but can also remove useful trends.

The non-linear model is more flexible, and can capture more complex trends, but it can also capture more noise (overfitting).

For the non-linear model

pd_population_nonlinear = partial_dependence(model_nonlinear, X_test, features=["age"], grid_resolution=30)

plt.plot(pd_population_nonlinear["grid_values"][0], pd_population_nonlinear["average"][0],
        color="black", linewidth=2, label="Model prediction, averaged over whole population")

for sex_value in ["M", "F"]:
    is_sex = X_test["sex"] == sex_value
    pd_sex_nonlinear = partial_dependence(model_nonlinear, X_test[is_sex], features=["age"], grid_resolution=30)
    plt.plot(pd_sex_nonlinear["grid_values"][0], pd_sex_nonlinear["average"][0],
            color=sex_colors[sex_value], linewidth=2,
            label=f"Model prediction, averaged for sex = {sex_value}")
    age_sex_mean = (
        pd.DataFrame({"age": X_test.loc[is_sex, "age"], "y_test": y_test[is_sex]})
        .groupby("age", as_index=False)["y_test"].mean()
    )
    plt.plot(age_sex_mean["age"], age_sex_mean["y_test"],
            color=sex_colors[sex_value], linewidth=1, linestyle="--", marker=sex_markers[sex_value], markersize=4,
            label=f"Average 5 year mortality, sex = {sex_value}")

plt.xlabel("age")
plt.ylabel("predicted probability of death within 5 years")
plt.title("Partial dependence of age, by sex - non-linear model")
plt.legend()
plt.tight_layout()
plt.show()

Partial dependence of age, by race/ethnicity

race_markers = {
    "Mexican American": "o",
    "Other Hispanic": "^",
    "Non-Hispanic White": "s",
    "Non-Hispanic Black": "D",
    "Other Race - Including Multi-Racial": "P",
}

plt.figure(figsize=(7, 5))

plt.plot(pd_population_nonlinear["grid_values"][0], pd_population_nonlinear["average"][0],
         color="black", linewidth=2, label="Model prediction, averaged over whole population")

for race_value in sorted(X_test["race_eth"].unique()):
    is_race = X_test["race_eth"] == race_value
    pd_race = partial_dependence(model_nonlinear, X_test[is_race], features=["age"], grid_resolution=30)
    line, = plt.plot(pd_race["grid_values"][0], pd_race["average"][0],
                     linewidth=2, label=f"Model prediction, averaged for race/ethnicity = {race_value}")
    age_race_mean = (
        pd.DataFrame({"age": X_test.loc[is_race, "age"], "y_test": y_test[is_race]})
        .groupby("age", as_index=False)["y_test"].mean()
    )
    plt.plot(
        age_race_mean["age"],
        age_race_mean["y_test"],
        color=line.get_color(),
        linewidth=1,
        linestyle="--",
        marker=race_markers.get(race_value, "o"),
        markersize=4,
        label=f"Average 5 year mortality, race/ethnicity = {race_value}",
    )

plt.xlabel("age")
plt.ylabel("predicted probability of death within 5 years")
plt.title("Partial dependence of age, by race/ethnicity")
plt.legend(fontsize=8)
plt.tight_layout()
plt.show()

Partial dependence of age, by education

Education codes 7 (“refused”) and 9 (“don’t know”) concern too few participants to give a meaningful curve, so they are excluded here.

education_labels = {
    1.0: "less than 9th grade",
    2.0: "9-11th grade",
    3.0: "high school grad / GED",
    4.0: "some college / AA degree",
    5.0: "college graduate or above",
}
education_markers = {
    1.0: "o",
    2.0: "^",
    3.0: "s",
    4.0: "D",
    5.0: "P",
}

plt.figure(figsize=(7, 5))

plt.plot(pd_population_nonlinear["grid_values"][0], pd_population_nonlinear["average"][0],
         color="black", linewidth=2, label="Model prediction, averaged over whole population")

for education_value, education_label in education_labels.items():
    is_education = X_test["education"] == education_value
    pd_education = partial_dependence(model_nonlinear, X_test[is_education], features=["age"], grid_resolution=30)
    line, = plt.plot(
        pd_education["grid_values"][0],
        pd_education["average"][0],
        linewidth=2,
        label=f"Model prediction, averaged for education = {education_label}",
    )
    age_education_mean = (
        pd.DataFrame({"age": X_test.loc[is_education, "age"], "y_test": y_test[is_education]})
        .groupby("age", as_index=False)["y_test"].mean()
    )
    plt.plot(
        age_education_mean["age"],
        age_education_mean["y_test"],
        color=line.get_color(),
        linewidth=1,
        linestyle="--",
        marker=education_markers[education_value],
        markersize=4,
        label=f"Average 5 year mortality, education = {education_label}",
    )

plt.xlabel("age")
plt.ylabel("predicted probability of death within 5 years")
plt.title("Partial dependence of age, by education")
plt.legend(fontsize=8)
plt.tight_layout()
plt.show()

2D partial dependence: age and waist circumference

We drop rows with a missing waist circumference are dropped, to ease plotting, though the model handes missing values

X_test_complete = X_test.dropna(subset=["age", "waist_cm"])
pd_age_waist = partial_dependence(
    model_nonlinear, X_test_complete, features=["age", "waist_cm"], grid_resolution=30
)
age_grid, waist_grid = pd_age_waist["grid_values"]

fig, ax = plt.subplots(figsize=(7, 5))
cs = ax.contourf(age_grid, waist_grid, pd_age_waist["average"][0].T, levels=20, cmap="viridis")
fig.colorbar(cs, ax=ax, label="predicted probability of death within 5 years")
ax.set_xlabel("age")
ax.set_ylabel("waist circumference (cm)")
ax.set_title("Partial dependence of age and waist circumference")
fig.tight_layout()
plt.show()