CODEWITHSUNDEEP
pythonpython-3.xmachine-learningstatsmodelspatsy
user9102437
user9102437

Reputation: 742

How to plot the confidence interval for statsmodels fit?

I wanted to show the confidence interval on the plot which I have made for the cubic spline of the data, but I have no idea how it should be done. From theory, I know that the CI should diverge from the fitted line when we get closer to the edges, but the only solution I came up with is this janky addition, which does not show the correct CI.

Here is the code:

import pandas as pd
from patsy import dmatrix
import statsmodels.api as sm
import matplotlib.pyplot as plt
import numpy as np

plt.figure(figsize=(7,5))
df = pd.read_csv('http://web.stanford.edu/~oleg2/hse/wage/wage.csv').sort_values(by=['age'])
ind_df = df[['wage', 'age']].copy()

def get_bse(bse, k, m, ma, x):
  prev, ans = m, []
  k.append(ma+1)
  for i, k_ in enumerate(k):
    ans += [bse[i]]*np.sum( ((x >= prev) & (x < k_)) ); prev = k_
  return np.array(ans)


plt.scatter(df.age, df.wage, color='none', edgecolor='silver', s=10)
plt.xlabel('Age', fontsize=15)
plt.ylabel('Wage', fontsize=15)
plt.ylim((0,333))

d = 4
knots = [df.age.quantile(0.25), df.age.quantile(0.5), df.age.quantile(0.75)]

my_spline_transformation = f"bs(train, knots={knots}, degree={d}, include_intercept=True)"

transformed = dmatrix( my_spline_transformation, {"train": df.age}, return_type='dataframe' )

ft = sm.GLS(df.wage, transformed).fit()

lft = sm.Logit( (df.age > 250), transformed )
y_grid1 = lft.predict(transformed.transpose())
y_grid = ft.predict(transformed)
plt.plot(df.age, y_grid, color='crimson', linewidth=2)
plt.plot(df.age, y_grid + get_bse(ft.bse, knots, df.age.min(), df.age.max(), df.age), color='crimson', linewidth=2, linestyle='--')
plt.plot(df.age, y_grid - get_bse(ft.bse, knots, df.age.min(), df.age.max(), df.age), color='crimson', linewidth=2, linestyle='--')

plt.show()

Output

Note that the graph is supposed to be a natural cubic spline function with four degrees of freedom, but I am not sure that my solution is correct. What is the proper way to achieve it?

Upvotes: 5

Views: 2664

Answers (1)

Jzbach
Jzbach

Reputation: 368

If you use:

predictions = ft.get_prediction()
df_predictions = predictions.summary_frame()
df_predictions.index = df.age.values

you will have a DataFrame with the results from the fitting with CI. Then, changing a bit the plotting to:

y_grid1 = lft.predict(transformed.transpose())
y_grid = ft.predict(transformed)
plt.plot(df.age, y_grid, color='crimson', linewidth=2)
#plt.plot(df.age, y_grid + get_bse(ft.bse, knots, df.age.min(), df.age.max(), df.age), color='crimson', linewidth=2, linestyle='--')
#plt.plot(df.age, y_grid - get_bse(ft.bse, knots, df.age.min(), df.age.max(), df.age), color='crimson', linewidth=2, linestyle='--')
predictions = ft.get_prediction()
df_predictions = predictions.summary_frame()
df_predictions.index = df.age.values
plt.plot(df_predictions['mean'], color='crimson')
plt.fill_between(df_predictions.index, df_predictions.mean_ci_lower, df_predictions.mean_ci_upper, alpha=.1, color='crimson')
plt.fill_between(df_predictions.index, df_predictions.obs_ci_lower, df_predictions.obs_ci_upper, alpha=.1, color='crimson')
plt.show()

yields:

answer

where the dark red area is the mean CI and the lighter red area is the CI of the observations in your dataset.

Upvotes: 3

Related Questions