I'm using ID3 algorithm to build a classifier and was wondering if there is any way to visualize the decision tree that the algorithm builds. This is my code for a decision tree in Python:

import numpy as np
import matplotlib.pyplot as plt # data visualization
import seaborn as sns # statistical data visualization
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier
from sklearn.preprocessing import LabelEncoder, OneHotEncoder
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
from collections import defaultdict
import pandas as pd
import math

def id3(data, target, features, default=0):
    Implements the ID3 algorithm for decision tree classification.

        data: A list of dictionaries, where each dictionary represents a data point
              with keys for features and the target variable.
        target: The name of the target variable in the data.
        features: A list of feature names.
        default: The default class value to return if no majority class is found
                 at a node (optional, defaults to 0).

        A decision tree represented as a dictionary.

    # Check if all data points have the same target value (base case)
    if all(data[0][target] == row[target] for row in data):
        return {target: data[0][target]}

    # Check if there are no remaining features (base case)
    if not features:
        # Majority vote for class label
        class_counts = defaultdict(int)
        for row in data:
            class_counts[row[target]] += 1
        return {target: max(class_counts, key=class_counts.get)}

    # Calculate entropy of the target variable
    entropy = _calculate_entropy(data, target)

    # Find the feature with the highest information gain
    best_feature, best_gain = None, -float('inf')
    for feature in features:
        gain = _calculate_information_gain(data, target, feature, entropy)
        if gain > best_gain:
            best_feature = feature
            best_gain = gain

    # Print information gain for this feature
        print(f"Feature '{feature}': Information Gain = {gain:.4f}")

    # Create a new decision tree node
    tree = {best_feature: {}}

    # Remove the best feature from the remaining features list
    remaining_features = features.copy()

    # Group data based on the best feature's values
    feature_values = set(row[best_feature] for row in data)
    for value in feature_values:
        # Filter data points with the current feature value
        subset_data = [row for row in data if row[best_feature] == value]

        # Print detailed calculations:
        print(f"Data for feature '{best_feature}' = '{value}':")
        print(f"- Entropy: {entropy:.4f}")  # Format entropy with 4 decimal places
        print(f"- Feature value: '{value}'")
        print(f"- Subset data:", subset_data)

        # Recursively build subtrees for each value of the best feature
        subtree = id3(subset_data, target, remaining_features, default)
        tree[best_feature][value] = subtree

    return tree

def _calculate_entropy(data, target):
    """Calculates the entropy of the target variable."""
    class_counts = defaultdict(int)
    for row in data:
        class_counts[row[target]] += 1

    entropy = 0
    for count in class_counts.values():
        if count > 0:
            p = count / len(data)
            entropy -= p * math.log2(p)

    return entropy

def _calculate_information_gain(data, target, feature, entropy):
    """Calculates the information gain for a given feature."""
    feature_values = set(row[feature] for row in data)
    information_gain = entropy

    for value in feature_values:
        subset_data = [row for row in data if row[feature] == value]
        subset_entropy = _calculate_entropy(subset_data, target)
        information_gain -= len(subset_data) / len(data) * subset_entropy

    return information_gain

#Import DataSet
dataC1=pd.read_csv("DataClass_Table.csv", sep=",", engine='python')
#Declare feature vector and target variable 
features = list(dataC1.columns)
target = features.pop()

# Converted data
data_list = dataC1.to_dict('records')

# Split data (replace 0.2 with your desired test size)
X_train, X_test, y_train, y_test = train_test_split(data_list, target, test_size=0.2)

#Call ID3 Function
decision_tree = id3(X_train, target, features)

How can I make the tree drawing with nodes like this? Decision Tree Visualization



Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.