Predicting Contact Lens Prescriptions with Decision Trees

Implementing the Decision Tree Algorithm

In this section, we will implement the core logic for constructing a decision tree using the ID3 algorithm. This involves calculating the Shannon Entropy to measure dataset impurity, splitting the dataset based on features, and recursively building the tree structure.

Calculating Shannon Entropy

The first step is to define a function that calculates the Shannon Entropy of a dataset. This metric quantifies the amount of disorder or uncertainty; higher values indicate a more mixed set of class labels.

import math
from collections import Counter

def calculate_entropy(dataset):
    num_entries = len(dataset)
    label_counts = Counter(row[-1] for row in dataset)
    
    entropy = 0.0
    for count in label_counts.values():
        probability = count / num_entries
        entropy -= probability * math.log(probability, 2)
        
    return entropy

This function iteraets through the dataset, counts the occurrences of each class label, and applies the entropy formula. It serves as the foundation for determining how much information a feature provides.

Splitting the Dataset

To evaluate a feature, we need a helper function that divides the dataset into subsets based on a specific feature value. This function isolates the data where a given feature matches a value and removes that feature column from the resulting subset.

def partition_dataset(dataset, axis, value):
    subset = []
    for row in dataset:
        if row[axis] == value:
            reduced_row = row[:axis] + row[axis+1:]
            subset.append(reduced_row)
    return subset

Selecting the Best Feature

The following function determines which feature offers the highest information gain. It calculates the entropy of the current dataset, then iterates through each unique value of every feature to compute the weighted entropy of the resulting partitions. The feature that results in the largest reduction in entropy (highest information gain) is selected.

def determine_best_feature(dataset):
    num_features = len(dataset[0]) - 1
    base_entropy = calculate_entropy(dataset)
    best_info_gain = -1
    best_feature_index = -1
    
    for i in range(num_features):
        feature_values = set(row[i] for row in dataset)
        new_entropy = 0.0
        
        for value in feature_values:
            sub_dataset = partition_dataset(dataset, i, value)
            probability = len(sub_dataset) / float(len(dataset))
            new_entropy += probability * calculate_entropy(sub_dataset)
            
        info_gain = base_entropy - new_entropy
        
        if info_gain > best_info_gain:
            best_info_gain = info_gain
            best_feature_index = i
            
    return best_feature_index

Building the Decision Tree

The tree construction is a recursive process. If all labels in a subset are identical, that subset becomes a leaf node. If no features remain to split on, the majority class is returned. Otherwice, the function selects the best feature, creates a node, and recursively builds branches for each unique value of that feature.

def get_majority_class(labels):
    return Counter(labels).most_common(1)[0][0]

def construct_tree(dataset, feature_labels):
    class_list = [row[-1] for row in dataset]
    
    # Stop if all classes are the same
    if class_list.count(class_list[0]) == len(class_list):
        return class_list[0]
    
    # Stop if no more features to split
    if len(dataset[0]) == 1:
        return get_majority_class(class_list)
    
    best_feat_idx = determine_best_feature(dataset)
    best_feat_label = feature_labels[best_feat_idx]
    
    tree = {best_feat_label: {}}
    
    # Create a copy of labels to avoid modifying the original list in recursion
    updated_labels = feature_labels[:]
    del(updated_labels[best_feat_idx])
    
    unique_values = set(row[best_feat_idx] for row in dataset)
    
    for value in unique_values:
        sub_dataset = partition_dataset(dataset, best_feat_idx, value)
        tree[best_feat_label][value] = construct_tree(sub_dataset, updated_labels)
        
    return tree

Data Preparation and Loading

We will now load the dataset containing patient attributes and corresponding contact lens recommendations. The data is read from a tab-separated file, parsed into a list of lists, and defined with appropriate feature headers.

# Assuming data is loaded from a file named 'lenses.txt'
# Format: age, prescript, astigmatic, tearRate, class
with open('lenses.txt', 'r') as file_handler:
    raw_data = [line.strip().split('\t') for line in file_handler.readlines()]

feature_headers = ['age', 'prescript', 'astigmatic', 'tearRate']
decision_tree = construct_tree(raw_data, feature_headers)

print(decision_tree)

Visualizing the Decision Tree

To better understend the logic, we can visualize the tree structure using Matplotlib. The following code defines functions to draw nodes, annotate edges, and recursively plot the tree structure.

import matplotlib.pyplot as plt

# Define node and arrow styles
decision_node_style = dict(boxstyle="sawtooth", fc="0.8")
leaf_node_style = dict(boxstyle="round4", fc="0.8")
arrow_args = dict(arrowstyle="<-")

def draw_node(node_text, center_point, parent_point, node_style):
    plot.ax1.annotate(node_text, xy=parent_point, xycoords='axes fraction',
                      xytext=center_point, textcoords='axes fraction',
                      va="center", ha="center", bbox=node_style, arrowprops=arrow_args)

def get_leaf_count(tree):
    num_leaves = 0
    root_key = list(tree.keys())[0]
    child_dict = tree[root_key]
    
    for key in child_dict.keys():
        if isinstance(child_dict[key], dict):
            num_leaves += get_leaf_count(child_dict[key])
        else:
            num_leaves += 1
    return num_leaves

def get_tree_depth(tree):
    max_depth = 0
    root_key = list(tree.keys())[0]
    child_dict = tree[root_key]
    
    for key in child_dict.keys():
        if isinstance(child_dict[key], dict):
            current_depth = 1 + get_tree_depth(child_dict[key])
        else:
            current_depth = 1
        if current_depth > max_depth:
            max_depth = current_depth
    return max_depth

def plot_mid_text(center_point, parent_point, text_string):
    x_mid = (parent_point[0] - center_point[0]) / 2.0 + center_point[0]
    y_mid = (parent_point[1] - center_point[1]) / 2.0 + center_point[1]
    plot.ax1.text(x_mid, y_mid, text_string, va="center", ha="center", rotation=30)

def draw_tree(tree, parent_point, node_text):
    num_leaves = get_leaf_count(tree)
    depth = get_tree_depth(tree)
    root_key = list(tree.keys())[0]
    
    # Calculate coordinates for the current node
    center_point = (plot.x_off + (1.0 + float(num_leaves)) / 2.0 / plot.total_width, plot.y_off)
    
    plot_mid_text(center_point, parent_point, node_text)
    draw_node(root_key, center_point, parent_point, decision_node_style)
    
    child_dict = tree[root_key]
    plot.y_off = plot.y_off - 1.0 / plot.total_depth
    
    for key in child_dict.keys():
        if isinstance(child_dict[key], dict):
            draw_tree(child_dict[key], center_point, str(key))
        else:
            plot.x_off = plot.x_off + 1.0 / plot.total_width
            draw_node(child_dict[key], (plot.x_off, plot.y_off), center_point, leaf_node_style)
            plot_mid_text((plot.x_off, plot.y_off), center_point, str(key))
            
    plot.y_off = plot.y_off + 1.0 / plot.total_depth

def create_plot(tree):
    fig = plt.figure(1, facecolor='white')
    fig.clf()
    ax_props = dict(xticks=[], yticks=[])
    plot.ax1 = plt.subplot(111, frameon=False, **ax_props)
    
    plot.total_width = float(get_leaf_count(tree))
    plot.total_depth = float(get_tree_depth(tree))
    plot.x_off = -0.5 / plot.total_width
    plot.y_off = 1.0
    
    draw_tree(tree, (0.5, 1.0), '')
    plt.show()

# Create the visualization
create_plot(decision_tree)

Classifying New Data

Finally, we implement a classification function that traverses the constructed tree to predict the contact lens type for a new patient. The function recursively matches the feature values of the input vector against the tree structure until it reaches a leaf node.

def predict(tree, feature_labels, test_vector):
    root_key = list(tree.keys())[0]
    sub_tree = tree[root_key]
    feature_index = feature_labels.index(root_key)
    
    for key in sub_tree.keys():
        if test_vector[feature_index] == key:
            if isinstance(sub_tree[key], dict):
                class_label = predict(sub_tree[key], feature_labels, test_vector)
            else:
                class_label = sub_tree[key]
    
    return class_label

# Example prediction
# Features: ['age', 'prescript', 'astigmatic', 'tearRate']
# Input: ['pre', 'myope', 'yes', 'normal']
result = predict(decision_tree, feature_headers, ['pre', 'myope', 'yes', 'normal'])
print(result)

Tags: python machine-learning decision-tree id3-algorithm data-visualization

Posted on Fri, 31 Jul 2026 16:23:35 +0000 by PHPiSean