#2 - Add gini and entropy measures

rename get_dataset to load_dataset
add features and impurity to  __str__ of node

main.py

@@ -1,72 +1,15 @@
 import time
 from sklearn.model_selection import train_test_split
+from sklearn.datasets import load_iris
 from stree import Stree
 
 random_state = 1
 
+X, y = load_iris(return_X_y=True)
 
-def load_creditcard(n_examples=0):
-    import pandas as pd
-    import numpy as np
-    import random
-
-    df = pd.read_csv("data/creditcard.csv")
-    print(
-        "Fraud: {0:.3f}% {1}".format(
-            df.Class[df.Class == 1].count() * 100 / df.shape[0],
-            df.Class[df.Class == 1].count(),
-        )
-    )
-    print(
-        "Valid: {0:.3f}% {1}".format(
-            df.Class[df.Class == 0].count() * 100 / df.shape[0],
-            df.Class[df.Class == 0].count(),
-        )
-    )
-    y = np.expand_dims(df.Class.values, axis=1)
-    X = df.drop(["Class", "Time", "Amount"], axis=1).values
-    if n_examples > 0:
-        # Take first n_examples samples
-        X = X[:n_examples, :]
-        y = y[:n_examples, :]
-    else:
-        # Take all the positive samples with a number of random negatives
-        if n_examples < 0:
-            Xt = X[(y == 1).ravel()]
-            yt = y[(y == 1).ravel()]
-            indices = random.sample(range(X.shape[0]), -1 * n_examples)
-            X = np.append(Xt, X[indices], axis=0)
-            y = np.append(yt, y[indices], axis=0)
-    print("X.shape", X.shape, " y.shape", y.shape)
-    print(
-        "Fraud: {0:.3f}% {1}".format(
-            len(y[y == 1]) * 100 / X.shape[0], len(y[y == 1])
-        )
-    )
-    print(
-        "Valid: {0:.3f}% {1}".format(
-            len(y[y == 0]) * 100 / X.shape[0], len(y[y == 0])
-        )
-    )
-    Xtrain, Xtest, ytrain, ytest = train_test_split(
-        X,
-        y,
-        train_size=0.7,
-        shuffle=True,
-        random_state=random_state,
-        stratify=y,
-    )
-    return Xtrain, Xtest, ytrain, ytest
-
-
-# data = load_creditcard(-5000) # Take all true samples + 5000 of the others
-# data = load_creditcard(5000)  # Take the first 5000 samples
-data = load_creditcard()  # Take all the samples
-
-Xtrain = data[0]
-Xtest = data[1]
-ytrain = data[2]
-ytest = data[3]
+Xtrain, Xtest, ytrain, ytest = train_test_split(
+    X, y, test_size=0.2, random_state=random_state
+)
 
 now = time.time()
 clf = Stree(C=0.01, random_state=random_state)