merge

src/test/test-trees.py

+# -*- coding: utf-8 -*-
+"""Test for decision trees."""
+import pandas as pd
+import subprocess
+from copy import deepcopy
+from sklearn.tree import DecisionTreeClassifier, export_graphviz
+
+# The decision tree of sklearn works with numerical values, hence all values
+# are mapped as integers
+DIRECTION = 'direction'
+RIGHT = 0
+LEFT = 1
+
+COLOR = 'color'
+RED = 0
+GREEN = 1
+BLUE = 2
+
+SIZE = 'size'
+SMALL = 0
+MEDIUM = 1
+BIG = 2
+
+SHAPE = 'shape'
+CUBE = 0
+SPHERE = 1
+CYLINDER = 2
+
+# Assuming input data will be between 5 and 10 samples to show what is going
+# left and what is going right the following objects are defined as tests for
+# the required input
+
+# Green objects intended to go to the Right
+training_data_0 = [{
+    COLOR: GREEN,
+    SIZE: MEDIUM,
+    SHAPE: CYLINDER,
+    DIRECTION: RIGHT
+}, {
+    COLOR: GREEN,
+    SIZE: SMALL,
+    SHAPE: CUBE,
+    DIRECTION: RIGHT
+}, {
+    COLOR: GREEN,
+    SIZE: BIG,
+    SHAPE: CYLINDER,
+    DIRECTION: RIGHT
+}, {
+    COLOR: BLUE,
+    SIZE: MEDIUM,
+    SHAPE: SPHERE,
+    DIRECTION: LEFT
+}, {
+    COLOR: RED,
+    SIZE: SMALL,
+    SHAPE: CUBE,
+    DIRECTION: LEFT
+}, {
+    COLOR: BLUE,
+    SIZE: BIG,
+    SHAPE: CYLINDER,
+    DIRECTION: LEFT
+}, {
+    COLOR: RED,
+    SIZE: MEDIUM,
+    SHAPE: SPHERE,
+    DIRECTION: LEFT
+}]
+
+# Cylinders and Red objects intended to go to the right
+training_data_1 = [{
+    COLOR: RED,
+    SIZE: MEDIUM,
+    SHAPE: CYLINDER,
+    DIRECTION: RIGHT
+}, {
+    COLOR: RED,
+    SIZE: SMALL,
+    SHAPE: CUBE,
+    DIRECTION: RIGHT
+}, {
+    COLOR: BLUE,
+    SIZE: BIG,
+    SHAPE: CYLINDER,
+    DIRECTION: RIGHT
+}, {
+    COLOR: GREEN,
+    SIZE: SMALL,
+    SHAPE: CYLINDER,
+    DIRECTION: RIGHT
+}, {
+    COLOR: RED,
+    SIZE: BIG,
+    SHAPE: CYLINDER,
+    DIRECTION: RIGHT
+}, {
+    COLOR: BLUE,
+    SIZE: MEDIUM,
+    SHAPE: SPHERE,
+    DIRECTION: LEFT
+}, {
+    COLOR: BLUE,
+    SIZE: BIG,
+    SHAPE: CUBE,
+    DIRECTION: LEFT
+}, {
+    COLOR: BLUE,
+    SIZE: SMALL,
+    SHAPE: CUBE,
+    DIRECTION: LEFT
+}, {
+    COLOR: GREEN,
+    SIZE: SMALL,
+    SHAPE: CUBE,
+    DIRECTION: LEFT
+}, {
+    COLOR: GREEN,
+    SIZE: MEDIUM,
+    SHAPE: SPHERE,
+    DIRECTION: LEFT
+}]
+
+# Only red big cubes intended to go Right
+training_data_2 = [{
+    COLOR: RED,
+    SIZE: BIG,
+    SHAPE: CUBE,
+    DIRECTION: RIGHT
+}, {
+    COLOR: RED,
+    SIZE: SMALL,
+    SHAPE: CYLINDER,
+    DIRECTION: LEFT
+}, {
+    COLOR: BLUE,
+    SIZE: BIG,
+    SHAPE: CUBE,
+    DIRECTION: LEFT
+}, {
+    COLOR: BLUE,
+    SIZE: MEDIUM,
+    SHAPE: SPHERE,
+    DIRECTION: LEFT
+}, {
+    COLOR: BLUE,
+    SIZE: SMALL,
+    SHAPE: SPHERE,
+    DIRECTION: LEFT
+}, {
+    COLOR: GREEN,
+    SIZE: BIG,
+    SHAPE: CUBE,
+    DIRECTION: LEFT
+}, {
+    COLOR: GREEN,
+    SIZE: MEDIUM,
+    SHAPE: CYLINDER,
+    DIRECTION: LEFT
+}]
+
+# Blue Cubes and Small Spheres intended to go Right
+training_data_3 = [{
+    COLOR: BLUE,
+    SIZE: BIG,
+    SHAPE: CUBE,
+    DIRECTION: RIGHT
+}, {
+    COLOR: RED,
+    SIZE: SMALL,
+    SHAPE: SPHERE,
+    DIRECTION: RIGHT
+}, {
+    COLOR: BLUE,
+    SIZE: SMALL,
+    SHAPE: CUBE,
+    DIRECTION: RIGHT
+}, {
+    COLOR: GREEN,
+    SIZE: SMALL,
+    SHAPE: SPHERE,
+    DIRECTION: RIGHT
+}, {
+    COLOR: RED,
+    SIZE: MEDIUM,
+    SHAPE: SPHERE,
+    DIRECTION: LEFT
+}, {
+    COLOR: RED,
+    SIZE: SMALL,
+    SHAPE: CUBE,
+    DIRECTION: LEFT
+}, {
+    COLOR: BLUE,
+    SIZE: SMALL,
+    SHAPE: CYLINDER,
+    DIRECTION: LEFT
+}, {
+    COLOR: GREEN,
+    SIZE: MEDIUM,
+    SHAPE: CUBE,
+    DIRECTION: LEFT
+}]
+
+test_data_0 = [{
+    COLOR: RED,
+    SIZE: BIG,
+    SHAPE: CUBE
+}, {
+    COLOR: RED,
+    SIZE: MEDIUM,
+    SHAPE: CYLINDER
+}, {
+    COLOR: RED,
+    SIZE: SMALL,
+    SHAPE: CUBE
+}, {
+    COLOR: BLUE,
+    SIZE: MEDIUM,
+    SHAPE: CYLINDER
+}, {
+    COLOR: BLUE,
+    SIZE: SMALL,
+    SHAPE: SPHERE
+}, {
+    COLOR: GREEN,
+    SIZE: SMALL,
+    SHAPE: SPHERE
+}, {
+    COLOR: GREEN,
+    SIZE: BIG,
+    SHAPE: CUBE
+}]
+
+# Trees definition
+# tree1 = DecisionTreeClassifier(
+#     criterion="entropy", random_state=100, max_depth=3)
+# tree2 = DecisionTreeClassifier(
+#     criterion="gini", random_state=100, max_depth=3)
+tree0 = DecisionTreeClassifier(criterion='entropy')
+tree1 = DecisionTreeClassifier(criterion='gini')
+
+
+# TODO: Make this method work if required
+def visualize_tree(tree, feature_names):
+    """Create tree png using graphviz.
+
+    Args
+    ----
+    tree -- scikit-learn DecsisionTree.
+    feature_names -- list of feature names.
+    """
+    with open("dt.dot", 'w') as f:
+        export_graphviz(tree, out_file=f, feature_names=feature_names)
+
+    command = ["dot", "-Tpng", "dt.dot", "-o", "dt.png"]
+    try:
+        subprocess.check_call(command)
+    except:
+        exit("Could not run dot, ie graphviz, to " "produce visualization")
+
+
+def get_trees():
+    """Get a list of trees to use in the experiment."""
+    return [tree0, tree1]
+
+
+def get_training_data():
+    """Get the list of data frames to train the trees."""
+    training_dfs = []
+    for td in [
+            training_data_0, training_data_1, training_data_2, training_data_3
+    ]:
+
+        df = pd.DataFrame(td)
+
+        # TODO: hardcode these values as they would not change
+        features = list(df.columns[:])
+        features.remove(DIRECTION)
+
+        y = df[DIRECTION]
+        X = df[features]
+
+        training_dfs.append((deepcopy(X), deepcopy(y)))
+    return training_dfs
+
+
+def get_test_data():
+    """Get the list of data for testing the trees."""
+    return [test_data_0]
+
+
+def run():
+    """Main method to run the tests."""
+    for i_t, tree in enumerate(get_trees()):
+        # print("Tree %i" % i_t)
+        for i_td, td in enumerate(get_training_data()):
+            tree.fit(td[0], td[1])
+            for test_d in get_test_data():
+                to_test = pd.DataFrame(test_d)
+                predicted = tree.predict(to_test)
+                predicted = [
+                    'RIGHT' if x == RIGHT else 'LEFT' for x in predicted
+                ]
+                print("Tree %i, training data %i, prediction %s" % (i_t, i_td,
+                                                                    predicted))
+
+
+run()