{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Decision Trees: Iris Classification\n",
        "\n",
        "This Notebook Studio example accompanies the article [AI Made Easy: Decision Trees for Data Science Success](/posts/ai-made-easy-decision-trees-for-data-science-success/).\n",
        "\n",
        "The goal is to make decision trees concrete: load a small dataset, train an interpretable classifier, inspect its rules, and compare a controlled tree with an overgrown one."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 1. Load the Iris dataset\n",
        "\n",
        "The Iris dataset is a classic beginner-friendly classification dataset. Each row is a flower. The model uses four measurements to predict the Iris species.\n",
        "\n",
        "The features are:\n",
        "\n",
        "- `sepal length (cm)`\n",
        "- `sepal width (cm)`\n",
        "- `petal length (cm)`\n",
        "- `petal width (cm)`\n",
        "\n",
        "The target has three classes: `setosa`, `versicolor`, and `virginica`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import pandas as pd\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "from sklearn.datasets import load_iris\n",
        "from sklearn.metrics import accuracy_score, confusion_matrix\n",
        "from sklearn.model_selection import train_test_split\n",
        "from sklearn.tree import DecisionTreeClassifier, plot_tree, export_text\n",
        "\n",
        "iris = load_iris(as_frame=True)\n",
        "X = iris.data\n",
        "y = iris.target\n",
        "\n",
        "data = X.copy()\n",
        "data[\"species\"] = y.map(dict(enumerate(iris.target_names)))\n",
        "\n",
        "data.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 2. Split the data\n",
        "\n",
        "We keep a test set aside so the tree has to make predictions on flowers it did not see during training. The split is stratified so each species is represented in both train and test sets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "X_train, X_test, y_train, y_test = train_test_split(\n",
        "    X,\n",
        "    y,\n",
        "    test_size=0.30,\n",
        "    random_state=42,\n",
        "    stratify=y,\n",
        ")\n",
        "\n",
        "print(f\"Training rows: {len(X_train)}\")\n",
        "print(f\"Test rows: {len(X_test)}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 3. Train a small decision tree\n",
        "\n",
        "A shallow tree is easier to inspect. Here we limit the tree to depth 3. That means it can ask at most three questions before reaching a prediction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "tree = DecisionTreeClassifier(\n",
        "    criterion=\"gini\",\n",
        "    max_depth=3,\n",
        "    random_state=42,\n",
        ")\n",
        "\n",
        "tree.fit(X_train, y_train)\n",
        "\n",
        "train_pred = tree.predict(X_train)\n",
        "test_pred = tree.predict(X_test)\n",
        "\n",
        "print(f\"Train accuracy: {accuracy_score(y_train, train_pred):.3f}\")\n",
        "print(f\"Test accuracy:  {accuracy_score(y_test, test_pred):.3f}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 4. Inspect the learned rules\n",
        "\n",
        "One reason decision trees are useful is that you can inspect the actual questions the model learned."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "rules = export_text(\n",
        "    tree,\n",
        "    feature_names=list(X.columns),\n",
        ")\n",
        "\n",
        "print(rules)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 5. Visualise the tree\n",
        "\n",
        "The root node is the first question. Each branch narrows the data until the model reaches a leaf with a predicted class."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "fig, ax = plt.subplots(figsize=(14, 8))\n",
        "plot_tree(\n",
        "    tree,\n",
        "    feature_names=list(X.columns),\n",
        "    class_names=list(iris.target_names),\n",
        "    filled=True,\n",
        "    rounded=True,\n",
        "    ax=ax,\n",
        ")\n",
        "ax.set_title(\"Decision tree trained on the Iris dataset\")\n",
        "fig"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 6. Check mistakes with a confusion matrix\n",
        "\n",
        "A confusion matrix shows which classes the model confuses. Rows are true species; columns are predicted species."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "cm = confusion_matrix(y_test, test_pred)\n",
        "cm_df = pd.DataFrame(\n",
        "    cm,\n",
        "    index=[f\"true {name}\" for name in iris.target_names],\n",
        "    columns=[f\"pred {name}\" for name in iris.target_names],\n",
        ")\n",
        "\n",
        "cm_df"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 7. Feature importance\n",
        "\n",
        "Decision trees can also tell us which features were most useful for making splits. On Iris, petal measurements usually dominate."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "importance = pd.Series(\n",
        "    tree.feature_importances_,\n",
        "    index=X.columns,\n",
        ").sort_values(ascending=True)\n",
        "\n",
        "ax = importance.plot(kind=\"barh\", figsize=(8, 4), color=\"#3b82f6\")\n",
        "ax.set_title(\"Decision tree feature importance\")\n",
        "ax.set_xlabel(\"Importance\")\n",
        "plt.tight_layout()\n",
        "ax.figure"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 8. What overfitting looks like\n",
        "\n",
        "Now compare the controlled tree with an unrestricted tree. If a tree grows too freely, it can memorise training details instead of learning a simpler pattern."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "models = {\n",
        "    \"controlled tree\": DecisionTreeClassifier(max_depth=3, random_state=42),\n",
        "    \"unrestricted tree\": DecisionTreeClassifier(random_state=42),\n",
        "}\n",
        "\n",
        "rows = []\n",
        "for name, model in models.items():\n",
        "    model.fit(X_train, y_train)\n",
        "    rows.append({\n",
        "        \"model\": name,\n",
        "        \"depth\": model.get_depth(),\n",
        "        \"leaves\": model.get_n_leaves(),\n",
        "        \"train_accuracy\": accuracy_score(y_train, model.predict(X_train)),\n",
        "        \"test_accuracy\": accuracy_score(y_test, model.predict(X_test)),\n",
        "    })\n",
        "\n",
        "pd.DataFrame(rows)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Takeaway\n",
        "\n",
        "Decision trees are useful because they are both predictive and inspectable. The trade-off is complexity: the deeper the tree, the easier it is to memorise the training set. Parameters such as `max_depth`, `min_samples_leaf`, and pruning help keep the tree understandable and more likely to generalise."
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "pygments_lexer": "ipython3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}