Decision Tree for classification tasks in Python

Question

I've decided to implement the ID3 Decision Tree algorithm in Python based on what I've learned from George F. Luger's textbook on AI (and other secondary readings). As far as I know, the code is working correctly:

import math

from random import choice
from collections import defaultdict, namedtuple

class TreeNode:
    """Used to implement the tree structure."""
    def __init__(self, type = None, value = None):
        #type may be "property" or "leaf". 
        self.type = type
        
        #The edges of the tree are represented as dict keys
        #and the nodes as the dict values.
        self.branches = dict()
        
        #This variable contains the name of the property represented by the node
        #or the label, if it's a leaf node.
        self.value = value

class DecisionTree:
    """This class implements a simple Decision Tree for classification tasks.
    After providing a dataset, the train method should be used to create the tree. 
    After that, the classify method may be used to classify a new example 
    by traversing the tree."""
    
   
    def __init__(self, dataset):
        """The dataset must be a sequence of named tuples representing each
        sample of the training set, which must contain the values for each attribute
        and also their label. Attributes might have any name but the samples must 
        have a 'label' attribute."""
        
        self.dataset = dataset
        self.tree = None
        self.labels = None 
    
    def train(self):
        """Create the decision tree based on the dataset and stores it in self.tree"""

        if len(self.dataset) == 0:
            print("The dataset was not provided.")
            return False

        #Find out the list of properties (assuming all examples have the same structure)
        properties = list(dataset[0]._fields)
        properties.remove("label")
        
        #Find out the possible labels for each sample. 
        self.labels = set([getattr(x, "label") for x in self.dataset])
        self.tree = self._induce_tree(self.dataset, properties)
    
    def classify(self, example):
        """Return a string (or a list of strings) containing the label
        (or possible labels) of the given example."""
        tree = self.tree
        while tree.type != "leaf":
            tree = tree.branches[getattr(example, tree.value)]
        return tree.value

    def _calc_entropy(self, probabilities):
        """Return the entropy of a given random variable given the
        probabilities (in the form of a list/tuple) of each of its possible outcomes"""
        
        entropy = 0
        for p in probabilities:
            entropy -= 0 if p == 0 else p * math.log(p, 2)
        return entropy

    def _calc_expected_entropy(self, partition):
        """Return the expected information needed to complete the tree"""
        
        #Get the total number of elements (population)
        total_pop = 0
        for p in partition.values():
            total_pop += len(p)

        expectation = 0

        for p in partition.values():
            p_pop = len(p)
            p_labels = [getattr(s, "label") for s in p]
            
            #Get the probability of each label ocurrency in the given
            #partition: No. of elements with the label divided by
            #the total no. of elements of the partitions.
            probs = []
            for label in self.labels:
                probs.append(p_labels.count(label)/p_pop)
            #Expected information needed =
            #Sigma |Partition_i| / Total population * Entropy of Partition_i
            #for i from 1 to n (no. of partitions).
            expectation += (p_pop/total_pop) * self._calc_entropy(probs)

        return expectation

    def _choose_property(self, example_set, properties):
        rank = list()
        for p in properties:
            #Partition the example set in groups of samples that 
            #share the same value for the property p. 
            partition = defaultdict(lambda: list())
            for s in example_set:
                partition[getattr(s, p)].append(s)

            rank.append((p, 
                self._calc_expected_entropy(partition), 
                partition)) 
        
        #Sort by expected entropy in increasing order.
        rank.sort(key = lambda x: x[1])
        #Since rank is a list of tuples containing the property,  
        #the expected entropy and partition for a subtree with that property
        #rank[0][0] in the sorted rank will be the property
        #with the least expected entropy (or with the most significant 
        #information gain) and rank[0][2] will be the partition of
        #the example set for that property.
        return (rank[0][0], rank[0][2])
    
    def _induce_tree(self, example_set, properties):
        """Recursive algorithm for inducing a decision tree based in ID3
        algorithm as explained by George F. Luger's textbook (2009).
         
        example_set: a list named tuples containing the samples
        properties: a list of strings containing the remaining properties
        to be chosen in the tree induction"""
        
        #If all samples belong to the same class, then
        #produce a leaf node for that class (label). 
        labels = set([getattr(s, "label") for s in example_set])
        if len(labels) == 1:
            return TreeNode(type = "leaf", value = list(labels)[0])
        
        #Else if there are no more properties to evaluate
        #but there are still examples with different labels
        #pick one at random and produce a leaf node with it.
        if not properties:
            return TreeNode(type = "leaf", value = choice(example_set).getatrr("label"))
        
        #Else, chose a property and make it the root of a new subtree.
        p, partition = self._choose_property(example_set, properties)
        new_root = TreeNode(type = "property", value = p)

        new_properties = list(properties)
        new_properties.remove(p)
        
        #Create a new branch in the subtree for each partition. 
        for value, block in partition.items():
            new_root.branches[value] = self._induce_tree(block, new_properties)
        return new_root

#This is just auxiliary code...
def generate_pydot_graph(graph, node, edge_name = "none"):
    """A recursive function to convert the generated tree
    into a pydot graph"""
    node_name = edge_name + '-' + node.value
    if node.type == "leaf":
        shape = "box"
    else:
        shape = "ellipse"

    graph.add_node(pydot.Node(node_name, label = node.value, shape = shape))
    if node.type == "property":
        for edge, sibling in node.branches.items():
            sibling_name = edge+'-'+sibling.value
            graph.add_edge(pydot.Edge(node_name, sibling_name, label = edge))
            generate_pydot_graph(graph, sibling, edge)
     

if __name__ == "__main__":
    Sample = namedtuple("Sample", ['label', 'credit_history', 'debit', 'collateral', 'income'])
    dataset = list()
    dataset.append(Sample('high', 'bad', 'high', 'none', '0-15k'))
    dataset.append(Sample('high', 'unknown', 'high', 'none', '15-35k'))
    dataset.append(Sample('moderate', 'unknown', 'low', 'none', '15-35k'))
    dataset.append(Sample('high', 'unknown', 'low', 'none', '0-15k'))
    dataset.append(Sample('low', 'unknown', 'low', 'none', '35k+'))
    dataset.append(Sample('low', 'unknown', 'low', 'adequate', '35k+'))
    dataset.append(Sample('high', 'bad', 'low', 'none', '0-15k'))
    dataset.append(Sample('moderate', 'bad', 'low', 'adequate', '35k+'))
    dataset.append(Sample('low', 'good', 'low', 'none', '35k+'))
    dataset.append(Sample('low', 'good', 'high', 'adequate', '35k+'))
    dataset.append(Sample('high', 'good', 'high', 'none', '0-15k'))
    dataset.append(Sample('moderate', 'good', 'high', 'none', '15-35k'))
    dataset.append(Sample('low', 'good', 'high', 'none', '35k+'))
    dataset.append(Sample('high', 'bad', 'high', 'none', '15-35k'))

    decision_tree = DecisionTree(dataset)
    decision_tree.train()
    
    #Testing the tree... 
    inst_fields = list(Sample._fields)
    inst_fields.remove("label")
    NewInstance = namedtuple("NewInstance", inst_fields)
    
    tests = list()
    tests.append(NewInstance("unknown", "low", "none", "15-35k")) #moderate
    tests.append(NewInstance("good", "node", "none", "0-15k"))  #high
    tests.append(NewInstance("unknown", "high", "adequate", "15-35k")) #high
    for t in tests:
        print(decision_tree.classify(t))
    
    #Generate a PNG image with a graphical representation of the tree 
    #to compare with the one in Luger's book :)
    import pydot
    graph = pydot.Dot("decision_tree", graph_type = "graph")
    generate_pydot_graph(graph, decision_tree.tree)
    graph.write_png("output.png") 

I'm particularly interested in feedback on the choices I've made concerning the data structures used (sequences of namedtuples for the dataset, dicts for partitions and so on). I'd also like to know if I commented and documented the code properly and if there are "more pythonic" ways of doing things.

My main goal was to implement the algorithm in a way that it would be easy to read and understand, so performance was not a primary concern of mine. Nonetheless, I'd appreciate tips on that too. Thank you!

Answer 1

You need to type-hint your code. Some of this will be made difficult by the dynamic nature of your field lookups, but you can get 90% of the way there.

property and type are poor variable name choices because they shadow built-ins.

Don't default type and value as None; you always know what these are going to be so there's no point in leaking nullity into those variables.

dict() can just be the literal {}.

DecisionTree is troubled by the separation of train from its constructor. There's basically never utility in separating these two methods, since it's expected that train() will always be called after construction, and in the interim the object is in an invalid state. Combining them reveals that the class doesn't actually need to hold onto a dataset at all.

Don't return False if the dataset is missing; raise an exception.

Your self.labels = set([ is better represented by a set comprehension.

_calc_entropy can convert its loop into a sum() with a generator. Also, don't unconditionally subtract with something that might be 0; instead conditionally subtract.

Don't getattr(s, "label"); just write s.label.

Don't write defaultdict(lambda: list()); list itself is callable and does not need a lambda.

(rank[0][0], rank[0][2]) will be made more clear by tuple-unpacking rank[0].

Avoid declaring your Sample and NewInstance in the local namespace; declare them globally with explicit types.

Don't import pydot locally; do this at the top.

Your __main__ guard is not enough. All of those symbols are still global, and need to be moved into a function. When you do that, a bug is revealed: one of your classes had a reference to the global dataset when it should have used self.dataset.

Don't sort followed by [0]; this is equivalent to min().

Suggested

import math
from collections import defaultdict
from random import choice
from typing import Collection, Iterable, Literal, NamedTuple, Sequence

import pydot


class Sample(NamedTuple):
    label: str
    credit_history: str
    debit: str
    collateral: str
    income: str


class NewInstance(NamedTuple):
    credit_history: str
    debit: str
    collateral: str
    income: str


class TreeNode:
    def __init__(self, type: Literal['property', 'leaf'], value: str) -> None:
        self.type = type

        # The edges of the tree are represented as dict keys
        # and the nodes as the dict values.
        self.branches: dict[str, TreeNode] = {}

        # This variable contains the name of the property represented by the node
        # or the label, if it's a leaf node.
        self.value = value


class DecisionTree:
    """This class implements a simple Decision Tree for classification tasks.
    After providing a dataset, the train method should be used to create the tree.
    After that, the classify method may be used to classify a new example
    by traversing the tree."""

    def __init__(self, dataset: Sequence[Sample]) -> None:
        """The dataset must be a sequence of named tuples representing each
        sample of the training set, which must contain the values for each attribute
        and also their label. Attributes might have any name but the samples must
        have a 'label' attribute."""

        if len(dataset) == 0:
            raise ValueError("The dataset was not provided.")

        # Find out the list of properties (assuming all examples have the same structure)
        properties = [f for f in dataset[0]._fields if f != 'label']

        # Find out the possible labels for each sample.
        self.labels = {sample.label for sample in dataset}

        # Create the decision tree based on the dataset and stores it in self.tree
        self.tree = self._induce_tree(dataset, properties)

    def classify(self, example: NewInstance) -> str:
        """Return a string (or a list of strings) containing the label
        (or possible labels) of the given example."""
        tree = self.tree
        while tree.type != "leaf":
            tree = tree.branches[getattr(example, tree.value)]
        return tree.value

    @staticmethod
    def _calc_entropy(probabilities: Iterable[float]) -> float:
        """Return the entropy of a given random variable given the
        probabilities (in the form of a list/tuple) of each of its possible outcomes"""

        entropy = -sum(
            p * math.log(p, 2)
            for p in probabilities
            if p != 0
        )
        return entropy

    def _calc_expected_entropy(self, partition: Collection[Collection[Sample]]) -> float:
        """Return the expected information needed to complete the tree"""

        # Get the total number of elements (population)
        total_pop = sum(len(p) for p in partition)
        expectation = 0

        for p in partition:
            p_pop = len(p)
            p_labels = [s.label for s in p]

            # Get the probability of each label occurrence in the given
            # partition: No. of elements with the label divided by
            # the total no. of elements of the partitions.
            probs = [
                p_labels.count(label) / p_pop
                for label in self.labels
            ]

            # Expected information needed =
            # Sigma |Partition_i| / Total population * Entropy of Partition_i
            # for i from 1 to n (no. of partitions).
            expectation += p_pop / total_pop * self._calc_entropy(probs)

        return expectation

    def _choose_property(
        self,
        example_set: Collection[Sample],
        properties: Iterable[str],
    ) -> tuple[
        str,  # property
        dict[str, list[Sample]],  # partition
    ]:
        rank = []
        for p in properties:
            # Partition the example set in groups of samples that
            # share the same value for the property p.
            partition = defaultdict(list)
            for s in example_set:
                partition[getattr(s, p)].append(s)

            rank.append((p,
                         self._calc_expected_entropy(partition.values()),
                         partition))

        # Sort by expected entropy in increasing order.

        # Since rank is a list of tuples containing the property,
        # the expected entropy and partition for a subtree with that property
        # rank[0][0] in the sorted rank will be the property
        # with the least expected entropy (or with the most significant
        # information gain) and rank[0][2] will be the partition of
        # the example set for that property.
        property, entropy, partition = min(rank, key=lambda x: x[1])
        return property, partition

    def _induce_tree(
        self,
        example_set: Sequence[Sample],
        properties: Collection[str],
    ) -> TreeNode:
        """Recursive algorithm for inducing a decision tree based in ID3
        algorithm as explained by George F. Luger's textbook (2009).

        example_set: a list named tuples containing the samples
        properties: a list of strings containing the remaining properties
        to be chosen in the tree induction"""

        # If all samples belong to the same class, then
        # produce a leaf node for that class (label).
        labels = {s.label for s in example_set}
        if len(labels) == 1:
            only_label, = labels
            return TreeNode(type="leaf", value=only_label)

        # Else if there are no more properties to evaluate
        # but there are still examples with different labels
        # pick one at random and produce a leaf node with it.
        if not properties:
            return TreeNode(type="leaf", value=choice(example_set).label)

        # Else, chose a property and make it the root of a new subtree.
        p, partition = self._choose_property(example_set, properties)
        new_root = TreeNode(type="property", value=p)

        new_properties = [prop for prop in properties if prop != p]

        # Create a new branch in the subtree for each partition.
        for value, block in partition.items():
            new_root.branches[value] = self._induce_tree(block, new_properties)
        return new_root


def generate_pydot_graph(graph: pydot.Dot, node: TreeNode, edge_name: str = "none") -> None:
    """A recursive function to convert the generated tree
    into a pydot graph"""
    node_name = edge_name + '-' + node.value
    if node.type == "leaf":
        shape = "box"
    else:
        shape = "ellipse"

    graph.add_node(pydot.Node(node_name, label=node.value, shape=shape))
    if node.type == "property":
        for edge, sibling in node.branches.items():
            sibling_name = edge + '-' + sibling.value
            graph.add_edge(pydot.Edge(node_name, sibling_name, label=edge))
            generate_pydot_graph(graph, sibling, edge)


def main() -> None:
    dataset = (
        Sample(    'high',     'bad', 'high',     'none',  '0-15k'),
        Sample(    'high', 'unknown', 'high',     'none', '15-35k'),
        Sample('moderate', 'unknown',  'low',     'none', '15-35k'),
        Sample(    'high', 'unknown',  'low',     'none',  '0-15k'),
        Sample(     'low', 'unknown',  'low',     'none',   '35k+'),
        Sample(     'low', 'unknown',  'low', 'adequate',   '35k+'),
        Sample(    'high',     'bad',  'low',     'none',  '0-15k'),
        Sample('moderate',     'bad',  'low', 'adequate',   '35k+'),
        Sample(     'low',    'good',  'low',     'none',   '35k+'),
        Sample(     'low',    'good', 'high', 'adequate',   '35k+'),
        Sample(    'high',    'good', 'high',     'none',  '0-15k'),
        Sample('moderate',    'good', 'high',     'none', '15-35k'),
        Sample(     'low',    'good', 'high',     'none',   '35k+'),
        Sample(    'high',     'bad', 'high',     'none', '15-35k'),
    )

    decision_tree = DecisionTree(dataset)

    tests = (
        NewInstance("unknown",  "low",     "none", "15-35k"),  # moderate
        NewInstance(   "good", "node",     "none",  "0-15k"),  # high
        NewInstance("unknown", "high", "adequate", "15-35k"),  # high
    )
    for t in tests:
        print(decision_tree.classify(t))

    # Generate a PNG image with a graphical representation of the tree
    # to compare with the one in Luger's book :)

    graph = pydot.Dot("decision_tree", graph_type="graph")
    generate_pydot_graph(graph, decision_tree.tree)
    graph.write_png("output.png")


if __name__ == "__main__":
    main()

Output