separate decisionTree from main()

main.py

@@ -17,6 +17,7 @@ from PatchType import PatchType
 from data.enum.CompanySize import CompanySize
 from data.enum.Direction import Direction
 from data.enum.Priority import Priority
+from tree.DecisionTree import DecisionTree
 
 colors = [
     'blue', 'cyan', 'orange', 'yellow', 'magenta', 'purple', '#103d3e', '#9fc86c',
@@ -60,80 +61,24 @@ def agent_portrayal(agent):
 
 
 if __name__ == '__main__':
-    test = ClientParamsFactory()
 
-    header = ['DELAY',
+    base = 512
-              'PAYED',
+    gridWidth = 10
-              'NET-WORTH',
+    gridHeight = 10
-              'INFLUENCE',
+    scale = base / gridWidth
-              'SKARBOWKA',
-              'MEMBER',
-              'HAT',
-              'SIZE']
 
-    with open("data/TEST/generatedData.csv", 'w', newline='') as file:
+    diagram4 = GridWithWeights(gridWidth, gridHeight)
-        writer = csv.writer(file)
+    diagram4.walls = [(6, 5), (6, 6), (6, 7), (6, 8), (2, 3), (2, 4), (3, 4), (4, 4), (6, 4)]
 
-        writer.writerow(header)
+    diagram5 = GridWithWeights(gridWidth, gridHeight)
+    diagram5.puddles = [(2, 2), (2, 5), (2, 6), (5, 4)]
 
-        for i in range(200):
+    grid = CanvasGrid(agent_portrayal, gridWidth, gridHeight, scale * gridWidth, scale * gridHeight)
-            data = test.get_client_params()
 
-            writer.writerow([data.payment_delay,
+    server = ModularServer(GameModel,
-                             data.payed,
+                           [grid],
-                             data.net_worth,
+                           "Automatyczny Wózek Widłowy",
-                             data.infuence_rate,
+                           {"width": gridHeight, "height": gridWidth, "graph": diagram4, "graph2": diagram5},)
-                             data.is_skarbowka,
-                             data.membership,
-                             data.is_hat,
-                             data.company_size])
 
-    file.close()
+    server.port = 8888
-
+    server.launch()
-    data_input = pandas.read_csv('data/TEST/importedData.csv', delimiter=",")
-
-    X_headers = ['DELAY','PAYED','NET-WORTH','INFLUENCE','SKARBOWKA','MEMBER','HAT','SIZE']
-
-    X = data_input[X_headers].values
-    Y = data_input["PRIORITY"]
-
-    label_BP = preprocessing.LabelEncoder()
-    label_BP.fit(['CompanySize.NO', 'CompanySize.SMALL', 'CompanySize.NORMAL', 'CompanySize.BIG', 'CompanySize.HUGE', 'CompanySize.GIGANTISHE'])
-    X[:, 7] = label_BP.transform(X[:, 7])
-
-    X_train, X_test, y_train, y_test = train_test_split(X, Y, test_size=0.2, train_size=0.8)
-
-    drugTree = DecisionTreeClassifier(criterion="entropy", max_depth=4)
-
-    clf = drugTree.fit(X_train, y_train)
-    predicted = drugTree.predict(X_test)
-
-    y_test = y_test.to_list()
-
-    for i in range(len(predicted)):
-        print("{}. {} \n predicted: {}, actual: {}".format(i, X_test[i,:], predicted[i], y_test[i]))
-
-    print("\nDecisionTrees's Accuracy: ", metrics.accuracy_score(y_test, predicted))
-
-    print(sklearn.tree.export_text(clf, feature_names=X_headers))
-
-    # base = 512
-    # gridWidth = 10
-    # gridHeight = 10
-    # scale = base / gridWidth
-    #
-    # diagram4 = GridWithWeights(gridWidth, gridHeight)
-    # diagram4.walls = [(6, 5), (6, 6), (6, 7), (6, 8), (2, 3), (2, 4), (3, 4), (4, 4), (6, 4)]
-    #
-    # diagram5 = GridWithWeights(gridWidth, gridHeight)
-    # diagram5.puddles = [(2, 2), (2, 5), (2, 6), (5, 4)]
-    #
-    # grid = CanvasGrid(agent_portrayal, gridWidth, gridHeight, scale * gridWidth, scale * gridHeight)
-    #
-    # server = ModularServer(GameModel,
-    #                        [grid],
-    #                        "Automatyczny Wózek Widłowy",
-    #                        {"width": gridHeight, "height": gridWidth, "graph": diagram4, "graph2": diagram5},)
-    #
-    # server.port = 8888
-    # server.launch()

tree/DecisionTree.py

@@ -0,0 +1,86 @@
+import csv
+
+import pandas
+import sklearn
+from sklearn import metrics, preprocessing
+from sklearn.model_selection import train_test_split
+from sklearn.tree import DecisionTreeClassifier
+
+from ClientParamsFactory import ClientParamsFactory
+
+
+class DecisionTree:
+    def __init__(self) -> None:
+        super().__init__()
+
+    def generate_data(self, generator: ClientParamsFactory, n:int):
+        header = ['DELAY',
+                  'PAYED',
+                  'NET-WORTH',
+                  'INFLUENCE',
+                  'SKARBOWKA',
+                  'MEMBER',
+                  'HAT',
+                  'SIZE']
+
+        with open("data/TEST/generatedData.csv", 'w', newline='') as file:
+            writer = csv.writer(file)
+
+            writer.writerow(header)
+
+            for i in range(n):
+                data = generator.get_client_params()
+
+                writer.writerow([data.payment_delay,
+                                 data.payed,
+                                 data.net_worth,
+                                 data.infuence_rate,
+                                 data.is_skarbowka,
+                                 data.membership,
+                                 data.is_hat,
+                                 data.company_size])
+
+        file.close()
+
+    def get_normalized_data(self, X):
+        label_BP = preprocessing.LabelEncoder()
+        label_BP.fit(
+            ['CompanySize.NO', 'CompanySize.SMALL', 'CompanySize.NORMAL', 'CompanySize.BIG', 'CompanySize.HUGE',
+             'CompanySize.GIGANTISHE'])
+        X[:, 7] = label_BP.transform(X[:, 7])
+
+        return X
+
+    def print_logs(self, x, y, prediction):
+        for i in range(len(prediction)):
+            print("{}. {} \n predicted: {}, actual: {}".format(i, x[i, :], prediction[i], y[i]))
+
+        print("\nDecisionTrees's Accuracy: ", metrics.accuracy_score(y, prediction))
+
+
+
+
+    def get_decision_tree(self) -> DecisionTreeClassifier:
+        data_input = pandas.read_csv('data/TEST/importedData.csv', delimiter=",")
+
+        X_headers = ['DELAY', 'PAYED', 'NET-WORTH', 'INFLUENCE', 'SKARBOWKA', 'MEMBER', 'HAT', 'SIZE']
+
+        X = data_input[X_headers].values
+        Y = data_input["PRIORITY"]
+
+        X = self.get_normalized_data(X)
+
+        X_train, X_test, y_train, y_test = train_test_split(X, Y, test_size=0.2, train_size=0.8)
+
+        drugTree = DecisionTreeClassifier(criterion="entropy", max_depth=4)
+
+        clf = drugTree.fit(X_train, y_train)
+        predicted = drugTree.predict(X_test)
+
+        y_test = y_test.to_list()
+
+        self.print_logs(X_test, y_test, predicted)
+
+        print(sklearn.tree.export_text(clf, feature_names=X_headers))
+
+        return drugTree