Commit 9b92f340 authored by fouad5's avatar fouad5
Browse files

merge

parents 4301adb9 8b52e2aa
# -*- 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()
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment