pygam_linear_gam_fit_with_partial_dependence_visualization.py

1from pygam import LinearGAM, s, f
2from pygam.datasets import wage
3import matplotlib.pyplot as plt
4
5# Load the dataset
6X, y = wage(return_X_y=True)
7
8# Fit a model with splines for the first two features and a factor for the third
9# s() stands for spline, f() stands for factor
10gam = LinearGAM(s(0) + s(1) + f(2)).gridsearch(X, y)
11
12# Print the model summary
13gam.summary()
14
15# Plot the partial dependence functions
16fig, axs = plt.subplots(1, 3, figsize=(15, 5))
17titles = ['year', 'age', 'education']
18
19for i, ax in enumerate(axs):
20    XX = gam.generate_X_grid(term=i)
21    pdep, confidence = gam.partial_dependence(term=i, width=.95)
22
23    ax.plot(XX[:, i], pdep)
24    ax.plot(XX[:, i], confidence, c='r', ls='--')
25    ax.set_title(titles[i])
26
27plt.show()