Linear Regression plots with Python

Part 1: Linear Regression with One Explanatory Variable

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

# Sample dataset (replace with your actual data)
# mydata = pd.read_csv("your_data.csv")
X = mydata[['Explanatory1']]
X = sm.add_constant(X) # Adds intercept
y = mydata['Response']

mod1 = sm.OLS(y, X).fit()
print(mod1.summary())

Based on this model, we can now graph the regression


# Create prediction data
graphdata1 = pd.DataFrame({
    'Explanatory1': np.linspace(mydata['Explanatory1'].min(), mydata['Explanatory1'].max(), 100)
})
X_new = sm.add_constant(graphdata1)

predictions = mod1.get_prediction(X_new)
pred_summary = predictions.summary_frame(alpha=0.05)

graphdata1['fit'] = pred_summary['mean']
graphdata1['LL'] = pred_summary['mean_ci_lower']
graphdata1['UL'] = pred_summary['mean_ci_upper']

# Plot
plt.figure(figsize=(10, 6))
plt.plot(graphdata1['Explanatory1'], graphdata1['fit'], color='red', linewidth=2)
plt.fill_between(graphdata1['Explanatory1'], graphdata1['LL'], graphdata1['UL'], color='gray', alpha=0.3)
plt.scatter(graphdata1['Explanatory1'], graphdata1['fit'], color='black', s=50)
plt.xlabel('Explanatory1')
plt.ylabel('Predicted Response')
plt.title('Linear Regression Prediction with Confidence Interval')
plt.show()

Part 2: Regression with Two Explanatory Variables

# Assuming mydata is already loaded
X = mydata[['Explanatory1', 'Explanatory2']]
X = sm.add_constant(X)
y = mydata['Response']

mod2 = sm.OLS(y, X).fit()
print(mod2.summary())

# Create grid of interesting values
ex1_vals = np.linspace(mydata['Explanatory1'].min(), mydata['Explanatory1'].max(), 50)
ex2_vals = mydata['Explanatory2'].unique() # or custom values

graphdata2 = pd.DataFrame([(x1, x2) for x2 in ex2_vals for x1 in ex1_vals],
columns=['Explanatory1', 'Explanatory2'])
X_new2 = sm.add_constant(graphdata2)

predictions2 = mod2.get_prediction(X_new2)
pred_summary2 = predictions2.summary_frame(alpha=0.05)

graphdata2['fit'] = pred_summary2['mean']
graphdata2['LL'] = pred_summary2['mean_ci_lower']
graphdata2['UL'] = pred_summary2['mean_ci_upper']

# Plot using seaborn for color grouping
import seaborn as sns

plt.figure(figsize=(10, 6))
sns.lineplot(data=graphdata2, x='Explanatory1', y='fit', hue='Explanatory2', palette='tab10', linewidth=2)
for _, group in graphdata2.groupby('Explanatory2'):
plt.fill_between(group['Explanatory1'], group['LL'], group['UL'], alpha=0.2)
plt.scatter(group['Explanatory1'], group['fit'], color='black', s=20)

plt.xlabel('Explanatory1')
plt.ylabel('Predicted Response')
plt.title('Regression Predictions with Two Explanatory Variables')
plt.legend(title='Explanatory2')
plt.show()