import pandas as pd
from sklearn.datasets import load_diabetes
from sklearn.tree import DecisionTreeRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error, r2_score
import matplotlib.pyplot as plt
from sklearn.tree import plot_tree
# 1. LOAD AND SPLIT DATA
# The Diabetes dataset is used to predict a quantitative measure of disease progression
diabetes = load_diabetes()
X_train, X_test, y_train, y_test = train_test_split(
diabetes.data, diabetes.target, test_size=0.3, random_state=42
)
print(f"Total samples: {len(diabetes.data)}")
print(f"Training samples: {len(X_train)}")
print("--- Data Loaded and Split Successfully ---")
# 2. TRAIN THE DECISION TREE REGRESSOR
# We set max_depth=3 for visualization and to prevent initial overfitting.
dt_reg = DecisionTreeRegressor(
max_depth=3, # Controls complexity (pruning)
criterion='squared_error', # The default splitting criteria for regression (MSE)
random_state=42
)
# Train the model on the training data
dt_reg.fit(X_train, y_train)
# 3. MAKE PREDICTIONS AND EVALUATE
y_pred = dt_reg.predict(X_test)
# Evaluation Metrics:
# MSE (Mean Squared Error): Measures the average squared difference between actual and predicted values. Lower is better.
mse = mean_squared_error(y_test, y_pred)
# R2 Score (Coefficient of Determination): Represents the proportion of the variance
# in the dependent variable that is predictable from the independent variables. Closer to 1 is better.
r2 = r2_score(y_test, y_pred)
print("\n--- Model Performance Evaluation ---")
print(f"Decision Tree Depth: {dt_reg.get_depth()}")
print(f"Mean Squared Error (MSE): {mse:.2f}")
print(f"R2 Score: {r2:.4f}")
# 4. VISUALIZE THE TREE (Optional but highly recommended for instruction)
plt.figure(figsize=(18, 10))
plot_tree(
dt_reg,
feature_names=diabetes.feature_names,
filled=True,
rounded=True,
fontsize=10,
# The 'value' in the leaf node is the mean target value (the prediction)
# The 'mse' is the error (variance) in that node.
)
plt.title("Decision Tree Regressor (max_depth=3)")
plt.show()
