ManimML/manim_ml/decision_tree/decision_tree.py

"""
    Module for visualizing decision trees in Manim. 
    It parses a decision tree classifier from sklearn. 

    TODO return a map from nodes to split animation for BFS tree expansion
    TODO reimplement the decision 2D decision tree surface drawing. 
"""
from manim import *
from manim_ml.decision_tree.classification_areas import (
    compute_decision_areas,
    merge_overlapping_polygons,
)
import manim_ml.decision_tree.helpers as helpers
from manim_ml.one_to_one_sync import OneToOneSync

import numpy as np
from PIL import Image

class LeafNode(Group):
    """Leaf node in tree"""

    def __init__(
        self, class_index, display_type="image", class_image_paths=[], class_colors=[]
    ):
        super().__init__()
        self.display_type = display_type
        self.class_image_paths = class_image_paths
        self.class_colors = class_colors
        assert self.display_type in ["image", "text"]
        if self.display_type == "image":
            self._construct_image_node(class_index)
        else:
            raise NotImplementedError()

    def _construct_image_node(self, class_index):
        """Make an image node"""
        # Get image
        image_path = self.class_image_paths[class_index]
        pil_image = Image.open(image_path)
        node = ImageMobject(pil_image)
        node.scale(1.5)
        rectangle = Rectangle(
            width=node.width + 0.05,
            height=node.height + 0.05,
            color=self.class_colors[class_index],
            stroke_width=6,
        )
        rectangle.move_to(node.get_center())
        rectangle.shift([-0.02, 0.02, 0])
        self.add(rectangle)
        self.add(node)

class SplitNode(VGroup):
    """Node for splitting decision in tree"""

    def __init__(self, feature, threshold):
        super().__init__()
        node_text = f"{feature}\n<=  {threshold:.2f} cm"
        # Draw decision text
        decision_text = Text(node_text, color=WHITE)
        # Draw the surrounding box
        bounding_box = SurroundingRectangle(decision_text, buff=0.3, color=WHITE)
        self.add(bounding_box)
        self.add(decision_text)

class DecisionTreeDiagram(Group):
    """Decision Tree Diagram Class for Manim"""

    def __init__(
        self,
        sklearn_tree,
        feature_names=None,
        class_names=None,
        class_images_paths=None,
        class_colors=[RED, GREEN, BLUE],
    ):
        super().__init__()
        self.tree = sklearn_tree
        self.feature_names = feature_names
        self.class_names = class_names
        self.class_image_paths = class_images_paths
        self.class_colors = class_colors
        # Make graph container for the tree
        self.tree_group, self.nodes_map, self.edge_map = self._make_tree()
        self.add(self.tree_group)

    def _make_node(
        self,
        node_index,
    ):
        """Make node"""
        is_split_node = (
            self.tree.children_left[node_index] != self.tree.children_right[node_index]
        )
        if is_split_node:
            node_feature = self.tree.feature[node_index]
            node_threshold = self.tree.threshold[node_index]
            node = SplitNode(self.feature_names[node_feature], node_threshold)
        else:
            # Get the most abundant class for the given leaf node
            # Make the leaf node object
            tree_class_index = np.argmax(self.tree.value[node_index])
            node = LeafNode(
                class_index=tree_class_index,
                class_colors=self.class_colors,
                class_image_paths=self.class_image_paths,
            )
        return node

    def _make_connection(self, top, bottom, is_leaf=False):
        """Make a connection from top to bottom"""
        top_node_bottom_location = top.get_center()
        top_node_bottom_location[1] -= top.height / 2
        bottom_node_top_location = bottom.get_center()
        bottom_node_top_location[1] += bottom.height / 2

        line = Line(top_node_bottom_location, bottom_node_top_location, color=WHITE)

        return line

    def _make_tree(self):
        """Construct the tree diagram"""
        tree_group = Group()
        max_depth = self.tree.max_depth
        # Make the root node
        nodes_map = {}
        root_node = self._make_node(
            node_index=0,
        )
        nodes_map[0] = root_node
        tree_group.add(root_node)
        # Save some information
        node_height = root_node.height
        node_width = root_node.width
        scale_factor = 1.0
        edge_map = {}
        # tree height
        tree_height = scale_factor * node_height * max_depth
        tree_width = scale_factor * 2**max_depth * node_width
        # traverse tree
        def recurse(node_index, depth, direction, parent_object, parent_node):
            # make the node object
            is_leaf = (
                self.tree.children_left[node_index]
                == self.tree.children_right[node_index]
            )
            node_object = self._make_node(node_index=node_index)
            nodes_map[node_index] = node_object
            node_height = node_object.height
            # set the node position
            direction_factor = -1 if direction == "left" else 1
            shift_right_amount = (
                0.9 * direction_factor * scale_factor * tree_width / (2**depth) / 2
            )
            if is_leaf:
                shift_down_amount = -1.0 * scale_factor * node_height
            else:
                shift_down_amount = -1.8 * scale_factor * node_height
            node_object.match_x(parent_object).match_y(parent_object).shift(
                [shift_right_amount, shift_down_amount, 0]
            )
            tree_group.add(node_object)
            # make a connection
            connection = self._make_connection(
                parent_object, node_object, is_leaf=is_leaf
            )
            edge_name = str(parent_node) + "," + str(node_index)
            edge_map[edge_name] = connection
            tree_group.add(connection)
            # recurse
            if not is_leaf:
                recurse(
                    self.tree.children_left[node_index],
                    depth + 1,
                    "left",
                    node_object,
                    node_index,
                )
                recurse(
                    self.tree.children_right[node_index],
                    depth + 1,
                    "right",
                    node_object,
                    node_index,
                )

        recurse(self.tree.children_left[0], 1, "left", root_node, 0)
        recurse(self.tree.children_right[0], 1, "right", root_node, 0)

        tree_group.scale(0.35)
        return tree_group, nodes_map, edge_map

    def create_level_order_expansion_decision_tree(self, tree):
        """Expands the decision tree in level order"""
        raise NotImplementedError()
    
    def create_bfs_expansion_decision_tree(self, tree):
        """Expands the tree using BFS"""
        animations = []
        split_node_animations = {} # Dictionary mapping split node to animation
        # Compute parent mapping
        parent_mapping = helpers.compute_node_to_parent_mapping(self.tree)
        # Create the root node as most common class
        placeholder_class_nodes = {}
        root_node_class_index = np.argmax(
            self.tree.value[0]
        )
        root_placeholder_node = LeafNode(
            class_index=root_node_class_index,
            class_colors=self.class_colors,
            class_image_paths=self.class_image_paths,
        )
        root_placeholder_node.move_to(self.nodes_map[0])
        placeholder_class_nodes[0] = root_placeholder_node
        root_create_animation = AnimationGroup(
            FadeIn(root_placeholder_node),
            lag_ratio=0.0
        )
        animations.append(root_create_animation)
        # Iterate through the nodes
        queue = [0]
        while len(queue) > 0:
            node_index = queue.pop(0)
            # Check if a node is a split node or not
            left_child_index = self.tree.children_left[node_index]
            right_child_index = self.tree.children_right[node_index]
            is_leaf_node = left_child_index == right_child_index
            if not is_leaf_node:
                # Remove the currently placeholder class node
                fade_out_animation = FadeOut(
                    placeholder_class_nodes[node_index]
                )
                animations.append(fade_out_animation)
                # Fade in the split node
                fade_in_animation = FadeIn(
                    self.nodes_map[node_index]
                )
                animations.append(fade_in_animation)
                # Split the node by creating the children and connecting them
                # to the parent
                # Handle left child
                assert left_child_index in self.nodes_map.keys()
                left_node = self.nodes_map[left_child_index]
                left_parent_edge = self.edge_map[f"{node_index},{left_child_index}"]
                # Get the children of the left node
                left_node_left_index = self.tree.children_left[left_child_index]
                left_node_right_index = self.tree.children_right[left_child_index]
                left_is_leaf = left_node_left_index == left_node_right_index
                if left_is_leaf:
                    # If a child is a leaf then just create it
                    left_animation = FadeIn(left_node)
                else:
                    # If the child is a split node find the dominant class and make a temp
                    left_node_class_index = np.argmax(
                        self.tree.value[left_child_index]
                    )
                    new_leaf_node = LeafNode(
                        class_index=left_node_class_index,
                        class_colors=self.class_colors,
                        class_image_paths=self.class_image_paths,
                    )
                    new_leaf_node.move_to(self.nodes_map[leaf_child_index])
                    placeholder_class_nodes[left_child_index] = new_leaf_node
                    left_animation = AnimationGroup(
                        FadeIn(new_leaf_node),
                        Create(left_parent_edge),
                        lag_ratio=0.0
                    )
                # Handle right child
                assert right_child_index in self.nodes_map.keys()
                right_node = self.nodes_map[right_child_index]
                right_parent_edge = self.edge_map[f"{node_index},{right_child_index}"]
                # Get the children of the left node
                right_node_left_index = self.tree.children_left[right_child_index]
                right_node_right_index = self.tree.children_right[right_child_index]
                right_is_leaf = right_node_left_index == right_node_right_index
                if right_is_leaf:
                    # If a child is a leaf then just create it
                    right_animation = FadeIn(right_node)
                else:
                    # If the child is a split node find the dominant class and make a temp
                    right_node_class_index = np.argmax(
                        self.tree.value[right_child_index]
                    )
                    new_leaf_node = LeafNode(
                        class_index=right_node_class_index,
                        class_colors=self.class_colors,
                        class_image_paths=self.class_image_paths,
                    )
                    placeholder_class_nodes[right_child_index] = new_leaf_node
                    right_animation = AnimationGroup(
                        FadeIn(new_leaf_node),
                        Create(right_parent_edge),
                        lag_ratio=0.0
                    )
                # Combine the animations
                split_animation = AnimationGroup(
                    left_animation,
                    right_animation,
                    lag_ratio=0.0,
                )
                animations.append(split_animation)
                # Add the split animation to the split node dict
                split_node_animations[node_index] = split_animation
            # Add the children to the queue
            if left_child_index != -1:
                queue.append(left_child_index)
            if right_child_index != -1:
                queue.append(right_child_index)

        return Succession(
            *animations, 
            lag_ratio=1.0
        ), split_node_animations

    def make_expand_tree_animation(self, node_expand_order):
        """
            Make an animation for expanding the decision tree

            Shows each split node as a leaf node initially, and
            then when it comes up shows it as a split node. The 
            reason for this is for purposes of animating each of the 
            splits in a decision surface.    
        """
        # Show the root node as a leaf node
        # Iterate through the nodes in the traversal order
        for node_index in node_expand_order[1:]:
            # Figure out if it is a leaf or not
            # If it is not a leaf then remove the placeholder leaf node
            #       then show the split node
            # If it is a leaf then just show the leaf node
        pass

    @override_animation(Create)
    def create_decision_tree(self, traversal_order="bfs"):
        """Makes a create animation for the decision tree"""
        # Comptue the node expand order
        if traversal_order == "level":
            node_expand_order = helpers.compute_level_order_traversal(self.tree)
        elif traversal_order == "bfs":
            node_expand_order = helpers.compute_bfs_traversal(self.tree)
        else:
            raise Exception(f"Uncrecognized traversal: {traversal_order}")
        # Make the animation
        expand_tree_animation = self.make_expand_tree_animation(node_expand_order)
        return expand_tree_animation

class DecisionTreeContainer(OneToOneSync):
    """Connects the DecisionTreeDiagram to the DecisionTreeEmbedding"""

    def __init__(self, sklearn_tree, points, classes):
        self.sklearn_tree = sklearn_tree
        self.points = points
        self.classes = classes

    def make_unfold_tree_animation(self):
        """Unfolds the tree through an in order traversal

        This animations unfolds the tree diagram as well as showing the splitting
        of a shaded region in the Decision Tree embedding.
        """
        # Draw points in the embedding
        # Start the tree splitting animation
        pass