#2 Cosmetic and style updates

#2 update benchmark notebook
#2 Refactor Stree & create Splitter
2025-08-17 16:36:01 +00:00 · 2020-06-15 11:09:11 +02:00 · 2020-06-15 10:33:51 +02:00 · 2020-06-15 00:22:57 +02:00 · 2020-06-14 14:00:21 +02:00 · 2020-06-14 03:08:55 +02:00
16 changed files with 832 additions and 1190 deletions
--- a/codecov.yml
+++ b/codecov.yml
@@ -3,9 +3,6 @@ overage:
    project:
      default:
        target: 90%
-    patch:
-      default:
-        target: 90%
 comment:
  layout: "reach, diff, flags, files"
  behavior: default
--- a/main.py
+++ b/main.py
@@ -1,88 +1,29 @@
 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()
+print("Predicting with max_features=sqrt(n_features)")
+clf = Stree(C=0.01, random_state=random_state, max_features="auto")
+clf.fit(Xtrain, ytrain)
+print(f"Took {time.time() - now:.2f} seconds to train")
+print(clf)
+print(f"Classifier's accuracy (train): {clf.score(Xtrain, ytrain):.4f}")
+print(f"Classifier's accuracy (test) : {clf.score(Xtest, ytest):.4f}")
+print("=" * 40)
+print("Predicting with max_features=n_features")
 clf = Stree(C=0.01, random_state=random_state)
 clf.fit(Xtrain, ytrain)
 print(f"Took {time.time() - now:.2f} seconds to train")
 print(clf)
 print(f"Classifier's accuracy (train): {clf.score(Xtrain, ytrain):.4f}")
 print(f"Classifier's accuracy (test) : {clf.score(Xtest, ytest):.4f}")
-proba = clf.predict_proba(Xtest)
-print(
-    "Checking that we have correct probabilities, these are probabilities of "
-    "sample belonging to class 1"
-)
-res0 = proba[proba[:, 0] == 0]
-res1 = proba[proba[:, 0] == 1]
-print("++++++++++res0 > .8++++++++++++")
-print(res0[res0[:, 1] > 0.8])
-print("**********res1 < .4************")
-print(res1[res1[:, 1] < 0.4])
--- a/notebooks/benchmark.ipynb
+++ b/notebooks/benchmark.ipynb
--- a/notebooks/features.ipynb
+++ b/notebooks/features.ipynb
@@ -64,7 +64,7 @@
    {
     "output_type": "stream",
     "name": "stdout",
-     "text": "Fraud: 0.173% 492\nValid: 99.827% 284315\nX.shape (1492, 28)  y.shape (1492,)\nFraud: 33.110% 494\nValid: 66.890% 998\n"
+     "text": "Fraud: 0.173% 492\nValid: 99.827% 284315\nX.shape (1492, 28)  y.shape (1492,)\nFraud: 33.244% 496\nValid: 66.756% 996\n"
    }
   ],
   "source": [
@@ -135,7 +135,7 @@
    {
     "output_type": "stream",
     "name": "stdout",
-     "text": "Accuracy of Train without weights 0.9789272030651341\nAccuracy of Train with    weights 0.9952107279693486\nAccuracy of Tests without weights 0.9598214285714286\nAccuracy of Tests with    weights 0.9508928571428571\n"
+     "text": "Accuracy of Train without weights 0.9808429118773946\nAccuracy of Train with    weights 0.9904214559386973\nAccuracy of Tests without weights 0.9441964285714286\nAccuracy of Tests with    weights 0.9375\n"
    }
   ],
   "source": [
@@ -162,7 +162,7 @@
    {
     "output_type": "stream",
     "name": "stdout",
-     "text": "Time: 0.27s\tKernel: linear\tAccuracy_train: 0.9683908045977011\tAccuracy_test: 0.953125\nTime: 0.09s\tKernel: rbf\tAccuracy_train: 0.9875478927203065\tAccuracy_test: 0.9598214285714286\nTime: 0.06s\tKernel: poly\tAccuracy_train: 0.9885057471264368\tAccuracy_test: 0.9464285714285714\n"
+     "text": "Time: 0.13s\tKernel: linear\tAccuracy_train: 0.9693486590038314\tAccuracy_test: 0.9598214285714286\nTime: 0.09s\tKernel: rbf\tAccuracy_train: 0.9923371647509579\tAccuracy_test: 0.953125\nTime: 0.09s\tKernel: poly\tAccuracy_train: 0.9913793103448276\tAccuracy_test: 0.9375\n"
    }
   ],
   "source": [
@@ -195,7 +195,7 @@
    {
     "output_type": "stream",
     "name": "stdout",
-     "text": "************** C=0.001 ****************************\nClassifier's accuracy (train): 0.9531\nClassifier's accuracy (test) : 0.9621\nroot\nroot - Down, <cgaf> - Leaf class=1 belief= 0.983713 counts=(array([0, 1]), array([  5, 302]))\nroot - Up, <cgaf> - Leaf class=0 belief= 0.940299 counts=(array([0, 1]), array([693,  44]))\n\n**************************************************\n************** C=0.01 ****************************\nClassifier's accuracy (train): 0.9569\nClassifier's accuracy (test) : 0.9621\nroot\nroot - Down, <cgaf> - Leaf class=1 belief= 0.990228 counts=(array([0, 1]), array([  3, 304]))\nroot - Up, <cgaf> - Leaf class=0 belief= 0.943012 counts=(array([0, 1]), array([695,  42]))\n\n**************************************************\n************** C=1 ****************************\nClassifier's accuracy (train): 0.9655\nClassifier's accuracy (test) : 0.9643\nroot\nroot - Down\nroot - Down - Down, <pure> - Leaf class=1 belief= 1.000000 counts=(array([1]), array([310]))\nroot - Down - Up, <pure> - Leaf class=0 belief= 1.000000 counts=(array([0]), array([5]))\nroot - Up, <cgaf> - Leaf class=0 belief= 0.950617 counts=(array([0, 1]), array([693,  36]))\n\n**************************************************\n************** C=5 ****************************\nClassifier's accuracy (train): 0.9684\nClassifier's accuracy (test) : 0.9598\nroot\nroot - Down\nroot - Down - Down, <pure> - Leaf class=1 belief= 1.000000 counts=(array([1]), array([311]))\nroot - Down - Up, <pure> - Leaf class=0 belief= 1.000000 counts=(array([0]), array([8]))\nroot - Up\nroot - Up - Down\nroot - Up - Down - Down, <pure> - Leaf class=1 belief= 1.000000 counts=(array([1]), array([1]))\nroot - Up - Down - Up, <pure> - Leaf class=0 belief= 1.000000 counts=(array([0]), array([2]))\nroot - Up - Up\nroot - Up - Up - Down, <pure> - Leaf class=0 belief= 1.000000 counts=(array([0]), array([2]))\nroot - Up - Up - Up\nroot - Up - Up - Up - Down\nroot - Up - Up - Up - Down - Down, <pure> - Leaf class=1 belief= 1.000000 counts=(array([1]), array([1]))\nroot - Up - Up - Up - Down - Up, <pure> - Leaf class=0 belief= 1.000000 counts=(array([0]), array([1]))\nroot - Up - Up - Up - Up, <cgaf> - Leaf class=0 belief= 0.954039 counts=(array([0, 1]), array([685,  33]))\n\n**************************************************\n************** C=17 ****************************\nClassifier's accuracy (train): 0.9751\nClassifier's accuracy (test) : 0.9464\nroot\nroot - Down\nroot - Down - Down, <pure> - Leaf class=1 belief= 1.000000 counts=(array([1]), array([304]))\nroot - Down - Up, <pure> - Leaf class=0 belief= 1.000000 counts=(array([0]), array([8]))\nroot - Up\nroot - Up - Down\nroot - Up - Down - Down, <pure> - Leaf class=1 belief= 1.000000 counts=(array([1]), array([4]))\nroot - Up - Down - Up, <pure> - Leaf class=0 belief= 1.000000 counts=(array([0]), array([3]))\nroot - Up - Up\nroot - Up - Up - Down\nroot - Up - Up - Down - Down, <pure> - Leaf class=1 belief= 1.000000 counts=(array([1]), array([4]))\nroot - Up - Up - Down - Up, <pure> - Leaf class=0 belief= 1.000000 counts=(array([0]), array([2]))\nroot - Up - Up - Up\nroot - Up - Up - Up - Down\nroot - Up - Up - Up - Down - Down, <pure> - Leaf class=1 belief= 1.000000 counts=(array([1]), array([3]))\nroot - Up - Up - Up - Down - Up, <pure> - Leaf class=0 belief= 1.000000 counts=(array([0]), array([1]))\nroot - Up - Up - Up - Up\nroot - Up - Up - Up - Up - Down\nroot - Up - Up - Up - Up - Down - Down, <pure> - Leaf class=1 belief= 1.000000 counts=(array([1]), array([3]))\nroot - Up - Up - Up - Up - Down - Up, <pure> - Leaf class=0 belief= 1.000000 counts=(array([0]), array([3]))\nroot - Up - Up - Up - Up - Up\nroot - Up - Up - Up - Up - Up - Down, <pure> - Leaf class=1 belief= 1.000000 counts=(array([1]), array([2]))\nroot - Up - Up - Up - Up - Up - Up, <cgaf> - Leaf class=0 belief= 0.963225 counts=(array([0, 1]), array([681,  26]))\n\n**************************************************\n0.6869 secs\n"
+     "text": "************** C=0.001 ****************************\nClassifier's accuracy (train): 0.9588\nClassifier's accuracy (test) : 0.9487\nroot feaures=(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27) impurity=0.4438\nroot - Down feaures=(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27) impurity=0.0374\nroot - Down - Down, <cgaf> - Leaf class=1 belief= 0.984076 impurity=0.0313 counts=(array([0, 1]), array([  5, 309]))\nroot - Down - Up, <pure> - Leaf class=0 belief= 1.000000 impurity=0.0000 counts=(array([0]), array([1]))\nroot - Up, <cgaf> - Leaf class=0 belief= 0.947874 impurity=0.0988 counts=(array([0, 1]), array([691,  38]))\n\n**************************************************\n************** C=0.01 ****************************\nClassifier's accuracy (train): 0.9588\nClassifier's accuracy (test) : 0.9531\nroot feaures=(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27) impurity=0.4438\nroot - Down feaures=(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27) impurity=0.0192\nroot - Down - Down, <cgaf> - Leaf class=1 belief= 0.993506 impurity=0.0129 counts=(array([0, 1]), array([  2, 306]))\nroot - Down - Up, <pure> - Leaf class=0 belief= 1.000000 impurity=0.0000 counts=(array([0]), array([1]))\nroot - Up, <cgaf> - Leaf class=0 belief= 0.944218 impurity=0.1053 counts=(array([0, 1]), array([694,  41]))\n\n**************************************************\n************** C=1 ****************************\nClassifier's accuracy (train): 0.9665\nClassifier's accuracy (test) : 0.9643\nroot feaures=(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27) impurity=0.4438\nroot - Down feaures=(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27) impurity=0.0189\nroot - Down - Down, <pure> - Leaf class=1 belief= 1.000000 impurity=0.0000 counts=(array([1]), array([312]))\nroot - Down - Up, <pure> - Leaf class=0 belief= 1.000000 impurity=0.0000 counts=(array([0]), array([3]))\nroot - Up, <cgaf> - Leaf class=0 belief= 0.951989 impurity=0.0914 counts=(array([0, 1]), array([694,  35]))\n\n**************************************************\n************** C=5 ****************************\nClassifier's accuracy (train): 0.9665\nClassifier's accuracy (test) : 0.9621\nroot feaures=(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27) impurity=0.4438\nroot - Down feaures=(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27) impurity=0.0250\nroot - Down - Down, <pure> - Leaf class=1 belief= 1.000000 impurity=0.0000 counts=(array([1]), array([312]))\nroot - Down - Up, <pure> - Leaf class=0 belief= 1.000000 impurity=0.0000 counts=(array([0]), array([4]))\nroot - Up, <cgaf> - Leaf class=0 belief= 0.951923 impurity=0.0915 counts=(array([0, 1]), array([693,  35]))\n\n**************************************************\n************** C=17 ****************************\nClassifier's accuracy (train): 0.9703\nClassifier's accuracy (test) : 0.9665\nroot feaures=(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27) impurity=0.4438\nroot - Down feaures=(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27) impurity=0.0367\nroot - Down - Down, <pure> - Leaf class=1 belief= 1.000000 impurity=0.0000 counts=(array([1]), array([315]))\nroot - Down - Up, <pure> - Leaf class=0 belief= 1.000000 impurity=0.0000 counts=(array([0]), array([6]))\nroot - Up feaures=(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27) impurity=0.0846\nroot - Up - Down, <pure> - Leaf class=1 belief= 1.000000 impurity=0.0000 counts=(array([1]), array([1]))\nroot - Up - Up, <cgaf> - Leaf class=0 belief= 0.957064 impurity=0.0822 counts=(array([0, 1]), array([691,  31]))\n\n**************************************************\n0.4375 secs\n"
    }
   ],
   "source": [
@@ -227,7 +227,7 @@
    {
     "output_type": "stream",
     "name": "stdout",
-     "text": "root\nroot - Down\nroot - Down - Down, <pure> - Leaf class=1 belief= 1.000000 counts=(array([1]), array([304]))\nroot - Down - Up, <pure> - Leaf class=0 belief= 1.000000 counts=(array([0]), array([8]))\nroot - Up\nroot - Up - Down\nroot - Up - Down - Down, <pure> - Leaf class=1 belief= 1.000000 counts=(array([1]), array([4]))\nroot - Up - Down - Up, <pure> - Leaf class=0 belief= 1.000000 counts=(array([0]), array([3]))\nroot - Up - Up\nroot - Up - Up - Down\nroot - Up - Up - Down - Down, <pure> - Leaf class=1 belief= 1.000000 counts=(array([1]), array([4]))\nroot - Up - Up - Down - Up, <pure> - Leaf class=0 belief= 1.000000 counts=(array([0]), array([2]))\nroot - Up - Up - Up\nroot - Up - Up - Up - Down\nroot - Up - Up - Up - Down - Down, <pure> - Leaf class=1 belief= 1.000000 counts=(array([1]), array([3]))\nroot - Up - Up - Up - Down - Up, <pure> - Leaf class=0 belief= 1.000000 counts=(array([0]), array([1]))\nroot - Up - Up - Up - Up\nroot - Up - Up - Up - Up - Down\nroot - Up - Up - Up - Up - Down - Down, <pure> - Leaf class=1 belief= 1.000000 counts=(array([1]), array([3]))\nroot - Up - Up - Up - Up - Down - Up, <pure> - Leaf class=0 belief= 1.000000 counts=(array([0]), array([3]))\nroot - Up - Up - Up - Up - Up\nroot - Up - Up - Up - Up - Up - Down, <pure> - Leaf class=1 belief= 1.000000 counts=(array([1]), array([2]))\nroot - Up - Up - Up - Up - Up - Up, <cgaf> - Leaf class=0 belief= 0.963225 counts=(array([0, 1]), array([681,  26]))\n"
+     "text": "root feaures=(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27) impurity=0.4438\nroot - Down feaures=(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27) impurity=0.0367\nroot - Down - Down, <pure> - Leaf class=1 belief= 1.000000 impurity=0.0000 counts=(array([1]), array([315]))\nroot - Down - Up, <pure> - Leaf class=0 belief= 1.000000 impurity=0.0000 counts=(array([0]), array([6]))\nroot - Up feaures=(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27) impurity=0.0846\nroot - Up - Down, <pure> - Leaf class=1 belief= 1.000000 impurity=0.0000 counts=(array([1]), array([1]))\nroot - Up - Up, <cgaf> - Leaf class=0 belief= 0.957064 impurity=0.0822 counts=(array([0, 1]), array([691,  31]))\n"
    }
   ],
   "source": [
@@ -244,7 +244,7 @@
    {
     "output_type": "stream",
     "name": "stdout",
-     "text": "root\nroot - Down\nroot - Down - Down, <pure> - Leaf class=1 belief= 1.000000 counts=(array([1]), array([304]))\nroot - Down - Up, <pure> - Leaf class=0 belief= 1.000000 counts=(array([0]), array([8]))\nroot - Up\nroot - Up - Down\nroot - Up - Down - Down, <pure> - Leaf class=1 belief= 1.000000 counts=(array([1]), array([4]))\nroot - Up - Down - Up, <pure> - Leaf class=0 belief= 1.000000 counts=(array([0]), array([3]))\nroot - Up - Up\nroot - Up - Up - Down\nroot - Up - Up - Down - Down, <pure> - Leaf class=1 belief= 1.000000 counts=(array([1]), array([4]))\nroot - Up - Up - Down - Up, <pure> - Leaf class=0 belief= 1.000000 counts=(array([0]), array([2]))\nroot - Up - Up - Up\nroot - Up - Up - Up - Down\nroot - Up - Up - Up - Down - Down, <pure> - Leaf class=1 belief= 1.000000 counts=(array([1]), array([3]))\nroot - Up - Up - Up - Down - Up, <pure> - Leaf class=0 belief= 1.000000 counts=(array([0]), array([1]))\nroot - Up - Up - Up - Up\nroot - Up - Up - Up - Up - Down\nroot - Up - Up - Up - Up - Down - Down, <pure> - Leaf class=1 belief= 1.000000 counts=(array([1]), array([3]))\nroot - Up - Up - Up - Up - Down - Up, <pure> - Leaf class=0 belief= 1.000000 counts=(array([0]), array([3]))\nroot - Up - Up - Up - Up - Up\nroot - Up - Up - Up - Up - Up - Down, <pure> - Leaf class=1 belief= 1.000000 counts=(array([1]), array([2]))\nroot - Up - Up - Up - Up - Up - Up, <cgaf> - Leaf class=0 belief= 0.963225 counts=(array([0, 1]), array([681,  26]))\n"
+     "text": "root feaures=(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27) impurity=0.4438\nroot - Down feaures=(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27) impurity=0.0367\nroot - Down - Down, <pure> - Leaf class=1 belief= 1.000000 impurity=0.0000 counts=(array([1]), array([315]))\nroot - Down - Up, <pure> - Leaf class=0 belief= 1.000000 impurity=0.0000 counts=(array([0]), array([6]))\nroot - Up feaures=(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27) impurity=0.0846\nroot - Up - Down, <pure> - Leaf class=1 belief= 1.000000 impurity=0.0000 counts=(array([1]), array([1]))\nroot - Up - Up, <cgaf> - Leaf class=0 belief= 0.957064 impurity=0.0822 counts=(array([0, 1]), array([691,  31]))\n"
    }
   ],
   "source": [
@@ -268,7 +268,7 @@
    {
     "output_type": "stream",
     "name": "stdout",
-     "text": "1 functools.partial(<function check_no_attributes_set_in_init at 0x1254f13b0>, 'Stree')\n2 functools.partial(<function check_estimators_dtypes at 0x1254e84d0>, 'Stree')\n3 functools.partial(<function check_fit_score_takes_y at 0x1254e83b0>, 'Stree')\n4 functools.partial(<function check_sample_weights_pandas_series at 0x1254e0cb0>, 'Stree')\n5 functools.partial(<function check_sample_weights_not_an_array at 0x1254e0dd0>, 'Stree')\n6 functools.partial(<function check_sample_weights_list at 0x1254e0ef0>, 'Stree')\n7 functools.partial(<function check_sample_weights_shape at 0x1254e2050>, 'Stree')\n8 functools.partial(<function check_sample_weights_invariance at 0x1254e2170>, 'Stree')\n9 functools.partial(<function check_estimators_fit_returns_self at 0x1254eb4d0>, 'Stree')\n10 functools.partial(<function check_estimators_fit_returns_self at 0x1254eb4d0>, 'Stree', readonly_memmap=True)\n11 functools.partial(<function check_complex_data at 0x1254e2320>, 'Stree')\n12 functools.partial(<function check_dtype_object at 0x1254e2290>, 'Stree')\n13 functools.partial(<function check_estimators_empty_data_messages at 0x1254e85f0>, 'Stree')\n14 functools.partial(<function check_pipeline_consistency at 0x1254e8290>, 'Stree')\n15 functools.partial(<function check_estimators_nan_inf at 0x1254e8710>, 'Stree')\n16 functools.partial(<function check_estimators_overwrite_params at 0x1254f1290>, 'Stree')\n17 functools.partial(<function check_estimator_sparse_data at 0x1254e0b90>, 'Stree')\n18 functools.partial(<function check_estimators_pickle at 0x1254e8950>, 'Stree')\n19 functools.partial(<function check_classifier_data_not_an_array at 0x1254f15f0>, 'Stree')\n20 functools.partial(<function check_classifiers_one_label at 0x1254eb050>, 'Stree')\n21 functools.partial(<function check_classifiers_classes at 0x1254eba70>, 'Stree')\n22 functools.partial(<function check_estimators_partial_fit_n_features at 0x1254e8a70>, 'Stree')\n23 functools.partial(<function check_classifiers_train at 0x1254eb170>, 'Stree')\n24 functools.partial(<function check_classifiers_train at 0x1254eb170>, 'Stree', readonly_memmap=True)\n25 functools.partial(<function check_classifiers_train at 0x1254eb170>, 'Stree', readonly_memmap=True, X_dtype='float32')\n26 functools.partial(<function check_classifiers_regression_target at 0x1254f40e0>, 'Stree')\n27 functools.partial(<function check_supervised_y_no_nan at 0x1254da9e0>, 'Stree')\n28 functools.partial(<function check_supervised_y_2d at 0x1254eb710>, 'Stree')\n29 functools.partial(<function check_estimators_unfitted at 0x1254eb5f0>, 'Stree')\n30 functools.partial(<function check_non_transformer_estimators_n_iter at 0x1254f1c20>, 'Stree')\n31 functools.partial(<function check_decision_proba_consistency at 0x1254f4200>, 'Stree')\n32 functools.partial(<function check_fit2d_predict1d at 0x1254e2830>, 'Stree')\n33 functools.partial(<function check_methods_subset_invariance at 0x1254e29e0>, 'Stree')\n34 functools.partial(<function check_fit2d_1sample at 0x1254e2b00>, 'Stree')\n35 functools.partial(<function check_fit2d_1feature at 0x1254e2c20>, 'Stree')\n36 functools.partial(<function check_fit1d at 0x1254e2d40>, 'Stree')\n37 functools.partial(<function check_get_params_invariance at 0x1254f1e60>, 'Stree')\n38 functools.partial(<function check_set_params at 0x1254f1f80>, 'Stree')\n39 functools.partial(<function check_dict_unchanged at 0x1254e2440>, 'Stree')\n40 functools.partial(<function check_dont_overwrite_parameters at 0x1254e2710>, 'Stree')\n41 functools.partial(<function check_fit_idempotent at 0x1254f43b0>, 'Stree')\n42 functools.partial(<function check_n_features_in at 0x1254f4440>, 'Stree')\n43 functools.partial(<function check_requires_y_none at 0x1254f44d0>, 'Stree')\n"
+     "text": "1 functools.partial(<function check_no_attributes_set_in_init at 0x12735b3b0>, 'Stree')\n2 functools.partial(<function check_estimators_dtypes at 0x1273514d0>, 'Stree')\n3 functools.partial(<function check_fit_score_takes_y at 0x1273513b0>, 'Stree')\n4 functools.partial(<function check_sample_weights_pandas_series at 0x12734acb0>, 'Stree')\n5 functools.partial(<function check_sample_weights_not_an_array at 0x12734add0>, 'Stree')\n6 functools.partial(<function check_sample_weights_list at 0x12734aef0>, 'Stree')\n7 functools.partial(<function check_sample_weights_shape at 0x12734d050>, 'Stree')\n8 functools.partial(<function check_sample_weights_invariance at 0x12734d170>, 'Stree')\n9 functools.partial(<function check_estimators_fit_returns_self at 0x1273564d0>, 'Stree')\n10 functools.partial(<function check_estimators_fit_returns_self at 0x1273564d0>, 'Stree', readonly_memmap=True)\n11 functools.partial(<function check_complex_data at 0x12734d320>, 'Stree')\n12 functools.partial(<function check_dtype_object at 0x12734d290>, 'Stree')\n13 functools.partial(<function check_estimators_empty_data_messages at 0x1273515f0>, 'Stree')\n14 functools.partial(<function check_pipeline_consistency at 0x127351290>, 'Stree')\n15 functools.partial(<function check_estimators_nan_inf at 0x127351710>, 'Stree')\n16 functools.partial(<function check_estimators_overwrite_params at 0x12735b290>, 'Stree')\n17 functools.partial(<function check_estimator_sparse_data at 0x12734ab90>, 'Stree')\n18 functools.partial(<function check_estimators_pickle at 0x127351950>, 'Stree')\n19 functools.partial(<function check_classifier_data_not_an_array at 0x12735b5f0>, 'Stree')\n20 functools.partial(<function check_classifiers_one_label at 0x127356050>, 'Stree')\n21 functools.partial(<function check_classifiers_classes at 0x127356a70>, 'Stree')\n22 functools.partial(<function check_estimators_partial_fit_n_features at 0x127351a70>, 'Stree')\n23 functools.partial(<function check_classifiers_train at 0x127356170>, 'Stree')\n24 functools.partial(<function check_classifiers_train at 0x127356170>, 'Stree', readonly_memmap=True)\n25 functools.partial(<function check_classifiers_train at 0x127356170>, 'Stree', readonly_memmap=True, X_dtype='float32')\n26 functools.partial(<function check_classifiers_regression_target at 0x12735f0e0>, 'Stree')\n27 functools.partial(<function check_supervised_y_no_nan at 0x1273449e0>, 'Stree')\n28 functools.partial(<function check_supervised_y_2d at 0x127356710>, 'Stree')\n29 functools.partial(<function check_estimators_unfitted at 0x1273565f0>, 'Stree')\n30 functools.partial(<function check_non_transformer_estimators_n_iter at 0x12735bc20>, 'Stree')\n31 functools.partial(<function check_decision_proba_consistency at 0x12735f200>, 'Stree')\n32 functools.partial(<function check_fit2d_predict1d at 0x12734d830>, 'Stree')\n33 functools.partial(<function check_methods_subset_invariance at 0x12734d9e0>, 'Stree')\n34 functools.partial(<function check_fit2d_1sample at 0x12734db00>, 'Stree')\n35 functools.partial(<function check_fit2d_1feature at 0x12734dc20>, 'Stree')\n36 functools.partial(<function check_fit1d at 0x12734dd40>, 'Stree')\n37 functools.partial(<function check_get_params_invariance at 0x12735be60>, 'Stree')\n38 functools.partial(<function check_set_params at 0x12735bf80>, 'Stree')\n39 functools.partial(<function check_dict_unchanged at 0x12734d440>, 'Stree')\n40 functools.partial(<function check_dont_overwrite_parameters at 0x12734d710>, 'Stree')\n41 functools.partial(<function check_fit_idempotent at 0x12735f3b0>, 'Stree')\n42 functools.partial(<function check_n_features_in at 0x12735f440>, 'Stree')\n43 functools.partial(<function check_requires_y_none at 0x12735f4d0>, 'Stree')\n"
    }
   ],
   "source": [
@@ -306,7 +306,7 @@
    {
     "output_type": "stream",
     "name": "stdout",
-     "text": "== Not Weighted ===\nSVC train score ..: 0.9521072796934866\nSTree train score : 0.9578544061302682\nSVC test score ...: 0.9553571428571429\nSTree test score .: 0.9575892857142857\n==== Weighted =====\nSVC train score ..: 0.9616858237547893\nSTree train score : 0.9616858237547893\nSVC test score ...: 0.9642857142857143\nSTree test score .: 0.9598214285714286\n*SVC test score ..: 0.951413553411694\n*STree test score : 0.9480517444389333\n"
+     "text": "== Not Weighted ===\nSVC train score ..: 0.9578544061302682\nSTree train score : 0.960727969348659\nSVC test score ...: 0.9508928571428571\nSTree test score .: 0.9553571428571429\n==== Weighted =====\nSVC train score ..: 0.9636015325670498\nSTree train score : 0.9626436781609196\nSVC test score ...: 0.9553571428571429\nSTree test score .: 0.9553571428571429\n*SVC test score ..: 0.9447820728419238\n*STree test score : 0.9447820728419238\n"
    }
   ],
   "source": [
@@ -338,12 +338,49 @@
    {
     "output_type": "stream",
     "name": "stdout",
-     "text": "root\nroot - Down\nroot - Down - Down, <cgaf> - Leaf class=1 belief= 0.969325 counts=(array([0, 1]), array([ 10, 316]))\nroot - Down - Up, <pure> - Leaf class=0 belief= 1.000000 counts=(array([0]), array([1]))\nroot - Up, <cgaf> - Leaf class=0 belief= 0.958159 counts=(array([0, 1]), array([687,  30]))\n\n"
+     "text": "root feaures=(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27) impurity=0.4438\nroot - Down, <cgaf> - Leaf class=1 belief= 0.978261 impurity=0.0425 counts=(array([0, 1]), array([  7, 315]))\nroot - Up, <cgaf> - Leaf class=0 belief= 0.955679 impurity=0.0847 counts=(array([0, 1]), array([690,  32]))\n\n"
    }
   ],
   "source": [
    "print(clf)"
   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Test max_features"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "metadata": {},
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": "****************************************\nmax_features None = 28\nTrain score : 0.9664750957854407\nTest score .: 0.9642857142857143\nTook 0.09 seconds\n****************************************\nmax_features auto = 5\nTrain score : 0.9511494252873564\nTest score .: 0.9441964285714286\nTook 0.37 seconds\n****************************************\nmax_features log2 = 4\nTrain score : 0.935823754789272\nTest score .: 0.9330357142857143\nTook 0.10 seconds\n****************************************\nmax_features 7 = 7\nTrain score : 0.9568965517241379\nTest score .: 0.9397321428571429\nTook 3.36 seconds\n****************************************\nmax_features 0.5 = 14\nTrain score : 0.960727969348659\nTest score .: 0.9486607142857143\nTook 112.42 seconds\n****************************************\nmax_features 0.1 = 2\nTrain score : 0.8793103448275862\nTest score .: 0.8839285714285714\nTook 0.06 seconds\n****************************************\nmax_features 0.7 = 19\nTrain score : 0.9655172413793104\nTest score .: 0.9553571428571429\nTook 10.59 seconds\n"
+    }
+   ],
+   "source": [
+    "for max_features in [None, \"auto\", \"log2\", 7, .5, .1, .7]:\n",
+    "    now = time.time()\n",
+    "    print(\"*\"*40)\n",
+    "    clf = Stree(random_state=random_state, max_features=max_features)\n",
+    "    clf.fit(Xtrain, ytrain)\n",
+    "    print(f\"max_features {max_features} = {clf.max_features_}\")\n",
+    "    print(\"Train score :\", clf.score(Xtrain, ytrain))\n",
+    "    print(\"Test score .:\", clf.score(Xtest, ytest))\n",
+    "    print(f\"Took {time.time() - now:.2f} seconds\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
  }
 ],
 "metadata": {
--- a/notebooks/test_graphs.ipynb
+++ b/notebooks/test_graphs.ipynb
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,4 @@
 numpy
 scikit-learn
 pandas
-matplotlib
 ipympl
--- a/setup.py
+++ b/setup.py
@@ -30,7 +30,7 @@ setuptools.setup(
        "Topic :: Scientific/Engineering :: Artificial Intelligence",
        "Intended Audience :: Science/Research",
    ],
-    install_requires=["scikit-learn>=0.23.0", "numpy", "matplotlib", "ipympl"],
+    install_requires=["scikit-learn>=0.23.0", "numpy", "ipympl"],
    test_suite="stree.tests",
    zip_safe=False,
 )
--- a/stree/Strees.py
+++ b/stree/Strees.py
@@ -7,19 +7,22 @@ Build an oblique tree classifier based on SVM Trees
 """

 import os
-
+import numbers
+import random
+import warnings
+from itertools import combinations
 import numpy as np
 from sklearn.base import BaseEstimator, ClassifierMixin
 from sklearn.svm import SVC, LinearSVC
 from sklearn.utils import check_consistent_length
 from sklearn.utils.multiclass import check_classification_targets
+from sklearn.exceptions import ConvergenceWarning
 from sklearn.utils.validation import (
    check_X_y,
    check_array,
    check_is_fitted,
    _check_sample_weight,
 )
-from sklearn.utils.sparsefuncs import count_nonzero
 from sklearn.metrics._classification import _weighted_sum, _check_targets


@@ -28,7 +31,15 @@ class Snode:
    dataset assigned to it
    """

-    def __init__(self, clf: SVC, X: np.ndarray, y: np.ndarray, title: str):
+    def __init__(
+        self,
+        clf: SVC,
+        X: np.ndarray,
+        y: np.ndarray,
+        features: np.array,
+        impurity: float,
+        title: str,
+    ):
        self._clf = clf
        self._title = title
        self._belief = 0.0
@@ -38,10 +49,21 @@ class Snode:
        self._down = None
        self._up = None
        self._class = None
+        self._feature = None
+        self._sample_weight = None
+        self._features = features
+        self._impurity = impurity

    @classmethod
    def copy(cls, node: "Snode") -> "Snode":
-        return cls(node._clf, node._X, node._y, node._title)
+        return cls(
+            node._clf,
+            node._X,
+            node._y,
+            node._features,
+            node._impurity,
+            node._title,
+        )

    def set_down(self, son):
        self._down = son
@@ -82,11 +104,15 @@ class Snode:
            count_values = np.unique(self._y, return_counts=True)
            result = (
                f"{self._title} - Leaf class={self._class} belief="
-                f"{self._belief: .6f} counts={count_values}"
+                f"{self._belief: .6f} impurity={self._impurity:.4f} "
+                f"counts={count_values}"
            )
            return result
        else:
-            return f"{self._title}"
+            return (
+                f"{self._title} feaures={self._features} impurity="
+                f"{self._impurity:.4f}"
+            )


 class Siterator:
@@ -97,9 +123,6 @@ class Siterator:
        self._stack = []
        self._push(tree)

-    def __iter__(self):
-        return self
-
    def _push(self, node: Snode):
        if node is not None:
            self._stack.append(node)
@@ -113,6 +136,169 @@ class Siterator:
        return node


+class Splitter:
+    def __init__(
+        self,
+        clf: SVC = None,
+        criterion: str = None,
+        splitter_type: str = None,
+        criteria: str = None,
+        min_samples_split: int = None,
+        random_state=None,
+    ):
+        self._clf = clf
+        self._random_state = random_state
+        if random_state is not None:
+            random.seed(random_state)
+        self._criterion = criterion
+        self._min_samples_split = min_samples_split
+        self._criteria = criteria
+        self._splitter_type = splitter_type
+
+        if clf is None:
+            raise ValueError(f"clf has to be a sklearn estimator, got({clf})")
+
+        if criterion not in ["gini", "entropy"]:
+            raise ValueError(
+                f"criterion must be gini or entropy got({criterion})"
+            )
+
+        if criteria not in ["min_distance", "max_samples"]:
+            raise ValueError(
+                f"split_criteria has to be min_distance or \
+                max_samples got ({criteria})"
+            )
+
+        if splitter_type not in ["random", "best"]:
+            raise ValueError(
+                f"splitter must be either random or best got({splitter_type})"
+            )
+        self.criterion_function = getattr(self, f"_{self._criterion}")
+        self.decision_criteria = getattr(self, f"_{self._criteria}")
+
+    def impurity(self, y: np.array) -> np.array:
+        return self.criterion_function(y)
+
+    @staticmethod
+    def _gini(y: np.array) -> float:
+        _, count = np.unique(y, return_counts=True)
+        return 1 - np.sum(np.square(count / np.sum(count)))
+
+    @staticmethod
+    def _entropy(y: np.array) -> float:
+        _, count = np.unique(y, return_counts=True)
+        proportion = count / np.sum(count)
+        return -np.sum(proportion * np.log2(proportion))
+
+    def information_gain(
+        self, labels_up: np.array, labels_dn: np.array
+    ) -> float:
+        card_up = labels_up.shape[0] if labels_up is not None else 0
+        card_dn = labels_dn.shape[0] if labels_dn is not None else 0
+        samples = card_up + card_dn
+        up = card_up / samples * self.criterion_function(labels_up)
+        dn = card_dn / samples * self.criterion_function(labels_dn)
+        return up + dn
+
+    def _select_best_set(
+        self, dataset: np.array, labels: np.array, features_sets: list
+    ) -> list:
+        min_impurity = 1
+        selected = None
+        warnings.filterwarnings("ignore", category=ConvergenceWarning)
+        for feature_set in features_sets:
+            self._clf.fit(dataset[:, feature_set], labels)
+            node = Snode(
+                self._clf, dataset, labels, feature_set, 0.0, "subset"
+            )
+            self.partition(dataset, node)
+            y1, y2 = self.part(labels)
+            impurity = self.information_gain(y1, y2)
+            if impurity < min_impurity:
+                min_impurity = impurity
+                selected = feature_set
+        return selected
+
+    def _get_subspaces_set(
+        self, dataset: np.array, labels: np.array, max_features: int
+    ) -> np.array:
+        features = range(dataset.shape[1])
+        features_sets = list(combinations(features, max_features))
+        if len(features_sets) > 1:
+            if self._splitter_type == "random":
+                return features_sets[random.randint(0, len(features_sets) - 1)]
+            else:
+                return self._select_best_set(dataset, labels, features_sets)
+        else:
+            return features_sets[0]
+
+    def get_subspace(
+        self, dataset: np.array, labels: np.array, max_features: int
+    ) -> list:
+        """Return the best subspace to make a split
+        """
+        indices = self._get_subspaces_set(dataset, labels, max_features)
+        return dataset[:, indices], indices
+
+    @staticmethod
+    def _min_distance(data: np.array, _) -> np.array:
+        # chooses the lowest distance of every sample
+        indices = np.argmin(np.abs(data), axis=1)
+        return np.array(
+            [data[x, y] for x, y in zip(range(len(data[:, 0])), indices)]
+        )
+
+    @staticmethod
+    def _max_samples(data: np.array, y: np.array) -> np.array:
+        # select the class with max number of samples
+        _, samples = np.unique(y, return_counts=True)
+        selected = np.argmax(samples)
+        return data[:, selected]
+
+    def partition(self, samples: np.array, node: Snode):
+        """Set the criteria to split arrays
+
+        """
+        data = self._distances(node, samples)
+        if data.shape[0] < self._min_samples_split:
+            self._down = np.ones((data.shape[0]), dtype=bool)
+            return
+        if data.ndim > 1:
+            # split criteria for multiclass
+            data = self.decision_criteria(data, node._y)
+        self._down = data > 0
+
+    @staticmethod
+    def _distances(node: Snode, data: np.ndarray) -> np.array:
+        """Compute distances of the samples to the hyperplane of the node
+
+        :param node: node containing the svm classifier
+        :type node: Snode
+        :param data: samples to find out distance to hyperplane
+        :type data: np.ndarray
+        :return: array of shape (m, 1) with the distances of every sample to
+        the hyperplane of the node
+        :rtype: np.array
+        """
+        return node._clf.decision_function(data[:, node._features])
+
+    def part(self, origin: np.array) -> list:
+        """Split an array in two based on indices (down) and its complement
+
+        :param origin: dataset to split
+        :type origin: np.array
+        :param down: indices to use to split array
+        :type down: np.array
+        :return: list with two splits of the array
+        :rtype: list
+        """
+        up = ~self._down
+        return [
+            origin[up] if any(up) else None,
+            origin[self._down] if any(self._down) else None,
+        ]
+
+
 class Stree(BaseEstimator, ClassifierMixin):
    """Estimator that is based on binary trees of svm nodes
    can deal with sample_weights in predict, used in boosting sklearn methods
@@ -131,7 +317,11 @@ class Stree(BaseEstimator, ClassifierMixin):
        tol: float = 1e-4,
        degree: int = 3,
        gamma="scale",
+        split_criteria: str = "max_samples",
+        criterion: str = "gini",
        min_samples_split: int = 0,
+        max_features=None,
+        splitter: str = "random",
    ):
        self.max_iter = max_iter
        self.C = C
@@ -142,63 +332,18 @@ class Stree(BaseEstimator, ClassifierMixin):
        self.gamma = gamma
        self.degree = degree
        self.min_samples_split = min_samples_split
+        self.split_criteria = split_criteria
+        self.max_features = max_features
+        self.criterion = criterion
+        self.splitter = splitter

    def _more_tags(self) -> dict:
-        """Required by sklearn to tell that this estimator is a binary classifier
+        """Required by sklearn to supply features of the classifier

        :return: the tag required
        :rtype: dict
        """
-        return {"binary_only": True, "requires_y": True}
-
-    def _split_array(self, origin: np.array, down: np.array) -> list:
-        """Split an array in two based on indices passed as down and its complement
-
-        :param origin: dataset to split
-        :type origin: np.array
-        :param down: indices to use to split array
-        :type down: np.array
-        :return: list with two splits of the array
-        :rtype: list
-        """
-        up = ~down
-        return (
-            origin[up[:, 0]] if any(up) else None,
-            origin[down[:, 0]] if any(down) else None,
-        )
-
-    def _distances(self, node: Snode, data: np.ndarray) -> np.array:
-        """Compute distances of the samples to the hyperplane of the node
-
-        :param node: node containing the svm classifier
-        :type node: Snode
-        :param data: samples to find out distance to hyperplane
-        :type data: np.ndarray
-        :return: array of shape (m, 1) with the distances of every sample to
-        the hyperplane of the node
-        :rtype: np.array
-        """
-        res = node._clf.decision_function(data)
-        if res.ndim == 1:
-            return np.expand_dims(res, 1)
-        elif res.shape[1] > 1:
-            # remove multiclass info
-            res = np.delete(res, slice(1, res.shape[1]), axis=1)
-        return res
-
-    def _split_criteria(self, data: np.array) -> np.array:
-        """Set the criteria to split arrays
-
-        :param data: [description]
-        :type data: np.array
-        :return: [description]
-        :rtype: np.array
-        """
-        return (
-            data > 0
-            if data.shape[0] >= self.min_samples_split
-            else np.ones((data.shape[0], 1), dtype=bool)
-        )
+        return {"requires_y": True}

    def fit(
        self, X: np.ndarray, y: np.ndarray, sample_weight: np.array = None
@@ -231,19 +376,95 @@ class Stree(BaseEstimator, ClassifierMixin):
                f"Maximum depth has to be greater than 1... got (max_depth=\
                    {self.max_depth})"
            )
+
        check_classification_targets(y)
        X, y = check_X_y(X, y)
        sample_weight = _check_sample_weight(sample_weight, X)
        check_classification_targets(y)
        # Initialize computed parameters
+        self.splitter_ = Splitter(
+            clf=self._build_clf(),
+            criterion=self.criterion,
+            splitter_type=self.splitter,
+            criteria=self.split_criteria,
+            random_state=self.random_state,
+            min_samples_split=self.min_samples_split,
+        )
+        if self.random_state is not None:
+            random.seed(self.random_state)
        self.classes_, y = np.unique(y, return_inverse=True)
+        self.n_classes_ = self.classes_.shape[0]
        self.n_iter_ = self.max_iter
        self.depth_ = 0
+        self.n_features_ = X.shape[1]
        self.n_features_in_ = X.shape[1]
+        self.max_features_ = self._initialize_max_features()
        self.tree_ = self.train(X, y, sample_weight, 1, "root")
        self._build_predictor()
        return self

+    def train(
+        self,
+        X: np.ndarray,
+        y: np.ndarray,
+        sample_weight: np.ndarray,
+        depth: int,
+        title: str,
+    ) -> Snode:
+        """Recursive function to split the original dataset into predictor
+        nodes (leaves)
+
+        :param X: samples dataset
+        :type X: np.ndarray
+        :param y: samples labels
+        :type y: np.ndarray
+        :param sample_weight: weight of samples. Rescale C per sample.
+        Hi weights force the classifier to put more emphasis on these points.
+        :type sample_weight: np.ndarray
+        :param depth: actual depth in the tree
+        :type depth: int
+        :param title: description of the node
+        :type title: str
+        :return: binary tree
+        :rtype: Snode
+        """
+        if depth > self.__max_depth:
+            return None
+        if np.unique(y).shape[0] == 1:
+            # only 1 class => pure dataset
+            return Snode(
+                clf=None,
+                X=X,
+                y=y,
+                features=X.shape[1],
+                impurity=0.0,
+                title=title + ", <pure>",
+            )
+        # Train the model
+        clf = self._build_clf()
+        Xs, features = self.splitter_.get_subspace(X, y, self.max_features_)
+        clf.fit(Xs, y, sample_weight=sample_weight)
+        impurity = self.splitter_.impurity(y)
+        node = Snode(clf, X, y, features, impurity, title)
+        self.depth_ = max(depth, self.depth_)
+        self.splitter_.partition(X, node)
+        X_U, X_D = self.splitter_.part(X)
+        y_u, y_d = self.splitter_.part(y)
+        sw_u, sw_d = self.splitter_.part(sample_weight)
+        if X_U is None or X_D is None:
+            # didn't part anything
+            return Snode(
+                clf,
+                X,
+                y,
+                features=X.shape[1],
+                impurity=impurity,
+                title=title + ", <cgaf>",
+            )
+        node.set_up(self.train(X_U, y_u, sw_u, depth + 1, title + " - Up"))
+        node.set_down(self.train(X_D, y_d, sw_d, depth + 1, title + " - Down"))
+        return node
+
    def _build_predictor(self):
        """Process the leaves to make them predictors
        """
@@ -278,53 +499,8 @@ class Stree(BaseEstimator, ClassifierMixin):
            )
        )

-    def train(
-        self,
-        X: np.ndarray,
-        y: np.ndarray,
-        sample_weight: np.ndarray,
-        depth: int,
-        title: str,
-    ) -> Snode:
-        """Recursive function to split the original dataset into predictor
-        nodes (leaves)
-
-        :param X: samples dataset
-        :type X: np.ndarray
-        :param y: samples labels
-        :type y: np.ndarray
-        :param sample_weight: weight of samples. Rescale C per sample.
-        Hi weights force the classifier to put more emphasis on these points.
-        :type sample_weight: np.ndarray
-        :param depth: actual depth in the tree
-        :type depth: int
-        :param title: description of the node
-        :type title: str
-        :return: binary tree
-        :rtype: Snode
-        """
-        if depth > self.__max_depth:
-            return None
-        if np.unique(y).shape[0] == 1:
-            # only 1 class => pure dataset
-            return Snode(None, X, y, title + ", <pure>")
-        # Train the model
-        clf = self._build_clf()
-        clf.fit(X, y, sample_weight=sample_weight)
-        tree = Snode(clf, X, y, title)
-        self.depth_ = max(depth, self.depth_)
-        down = self._split_criteria(self._distances(tree, X))
-        X_U, X_D = self._split_array(X, down)
-        y_u, y_d = self._split_array(y, down)
-        sw_u, sw_d = self._split_array(sample_weight, down)
-        if X_U is None or X_D is None:
-            # didn't part anything
-            return Snode(clf, X, y, title + ", <cgaf>")
-        tree.set_up(self.train(X_U, y_u, sw_u, depth + 1, title + " - Up"))
-        tree.set_down(self.train(X_D, y_d, sw_d, depth + 1, title + " - Down"))
-        return tree
-
-    def _reorder_results(self, y: np.array, indices: np.array) -> np.array:
+    @staticmethod
+    def _reorder_results(y: np.array, indices: np.array) -> np.array:
        """Reorder an array based on the array of indices passed

        :param y: data untidy
@@ -334,12 +510,8 @@ class Stree(BaseEstimator, ClassifierMixin):
        :return: array y ordered
        :rtype: np.array
        """
-        if y.ndim > 1 and y.shape[1] > 1:
-            # if predict_proba return np.array of floats
-            y_ordered = np.zeros(y.shape, dtype=float)
-        else:
-            # return array of same type given in y
-            y_ordered = y.copy()
+        # return array of same type given in y
+        y_ordered = y.copy()
        indices = indices.astype(int)
        for i, index in enumerate(indices):
            y_ordered[index] = y[i]
@@ -363,17 +535,22 @@ class Stree(BaseEstimator, ClassifierMixin):
                # set a class for every sample in dataset
                prediction = np.full((xp.shape[0], 1), node._class)
                return prediction, indices
-            down = self._split_criteria(self._distances(node, xp))
-            X_U, X_D = self._split_array(xp, down)
-            i_u, i_d = self._split_array(indices, down)
-            prx_u, prin_u = predict_class(X_U, i_u, node.get_up())
-            prx_d, prin_d = predict_class(X_D, i_d, node.get_down())
+            self.splitter_.partition(xp, node)
+            x_u, x_d = self.splitter_.part(xp)
+            i_u, i_d = self.splitter_.part(indices)
+            prx_u, prin_u = predict_class(x_u, i_u, node.get_up())
+            prx_d, prin_d = predict_class(x_d, i_d, node.get_down())
            return np.append(prx_u, prx_d), np.append(prin_u, prin_d)

        # sklearn check
        check_is_fitted(self, ["tree_"])
        # Input validation
        X = check_array(X)
+        if X.shape[1] != self.n_features_:
+            raise ValueError(
+                f"Expected {self.n_features_} features but got "
+                f"({X.shape[1]})"
+            )
        # setup prediction & make it happen
        indices = np.arange(X.shape[0])
        result = (
@@ -383,68 +560,6 @@ class Stree(BaseEstimator, ClassifierMixin):
        )
        return self.classes_[result]

-    def predict_proba(self, X: np.array) -> np.array:
-        """Computes an approximation of the probability of samples belonging to
-        class 0 and 1
-        :param X: dataset
-        :type X: np.array
-        :return: array array of shape (m, num_classes), probability of being
-        each class
-        :rtype: np.array
-        """
-
-        def predict_class(
-            xp: np.array, indices: np.array, dist: np.array, node: Snode
-        ) -> np.array:
-            """Run the tree to compute predictions
-
-            :param xp: subdataset of samples
-            :type xp: np.array
-            :param indices: indices of subdataset samples to rebuild original
-            order
-            :type indices: np.array
-            :param dist: distances of every sample to the hyperplane or the
-            father node
-            :type dist: np.array
-            :param node: node of the leaf with the class
-            :type node: Snode
-            :return: array of labels and distances, array of indices
-            :rtype: np.array
-            """
-            if xp is None:
-                return [], []
-            if node.is_leaf():
-                # set a class for every sample in dataset
-                prediction = np.full((xp.shape[0], 1), node._class)
-                prediction_proba = dist
-                return np.append(prediction, prediction_proba, axis=1), indices
-            distances = self._distances(node, xp)
-            down = self._split_criteria(distances)
-            X_U, X_D = self._split_array(xp, down)
-            i_u, i_d = self._split_array(indices, down)
-            di_u, di_d = self._split_array(distances, down)
-            prx_u, prin_u = predict_class(X_U, i_u, di_u, node.get_up())
-            prx_d, prin_d = predict_class(X_D, i_d, di_d, node.get_down())
-            return np.append(prx_u, prx_d), np.append(prin_u, prin_d)
-
-        # sklearn check
-        check_is_fitted(self, ["tree_"])
-        # Input validation
-        X = check_array(X)
-        # setup prediction & make it happen
-        indices = np.arange(X.shape[0])
-        empty_dist = np.empty((X.shape[0], 1), dtype=float)
-        result, indices = predict_class(X, indices, empty_dist, self.tree_)
-        result = result.reshape(X.shape[0], 2)
-        # Turn distances to hyperplane into probabilities based on fitting
-        # distances of samples to its hyperplane that classified them, to the
-        # sigmoid function
-        # Probability of being 1
-        result[:, 1] = 1 / (1 + np.exp(-result[:, 1]))
-        # Probability of being 0
-        result[:, 0] = 1 - result[:, 1]
-        return self._reorder_results(result, indices)
-
    def score(
        self, X: np.array, y: np.array, sample_weight: np.array = None
    ) -> float:
@@ -466,19 +581,14 @@ class Stree(BaseEstimator, ClassifierMixin):
        X, y = check_X_y(X, y)
        y_pred = self.predict(X).reshape(y.shape)
        # Compute accuracy for each possible representation
-        y_type, y_true, y_pred = _check_targets(y, y_pred)
+        _, y_true, y_pred = _check_targets(y, y_pred)
        check_consistent_length(y_true, y_pred, sample_weight)
-        if y_type.startswith("multilabel"):
-            differing_labels = count_nonzero(y_true - y_pred, axis=1)
-            score = differing_labels == 0
-        else:
-            score = y_true == y_pred
-
+        score = y_true == y_pred
        return _weighted_sum(score, sample_weight, normalize=True)

    def __iter__(self) -> Siterator:
-        """Create an iterator to be able to visit the nodes of the tree in preorder,
-        can make a list with all the nodes in preorder
+        """Create an iterator to be able to visit the nodes of the tree in
+        preorder, can make a list with all the nodes in preorder

        :return: an iterator, can for i in... and list(...)
        :rtype: Siterator
@@ -499,3 +609,34 @@ class Stree(BaseEstimator, ClassifierMixin):
        for i in self:
            output += str(i) + "\n"
        return output
+
+    def _initialize_max_features(self) -> int:
+        if isinstance(self.max_features, str):
+            if self.max_features == "auto":
+                max_features = max(1, int(np.sqrt(self.n_features_)))
+            elif self.max_features == "sqrt":
+                max_features = max(1, int(np.sqrt(self.n_features_)))
+            elif self.max_features == "log2":
+                max_features = max(1, int(np.log2(self.n_features_)))
+            else:
+                raise ValueError(
+                    "Invalid value for max_features. "
+                    "Allowed string values are 'auto', "
+                    "'sqrt' or 'log2'."
+                )
+        elif self.max_features is None:
+            max_features = self.n_features_
+        elif isinstance(self.max_features, numbers.Integral):
+            max_features = self.max_features
+        else:  # float
+            if self.max_features > 0.0:
+                max_features = max(
+                    1, int(self.max_features * self.n_features_)
+                )
+            else:
+                raise ValueError(
+                    "Invalid value for max_features."
+                    "Allowed float must be in range (0, 1] "
+                    f"got ({self.max_features})"
+                )
+        return max_features
--- a/stree/Strees_grapher.py
+++ b/stree/Strees_grapher.py
@@ -1,205 +0,0 @@
-"""
-__author__ = "Ricardo Montañana Gómez"
-__copyright__ = "Copyright 2020, Ricardo Montañana Gómez"
-__license__ = "MIT"
-__version__ = "0.9"
-Plot 3D views of nodes in Stree
-"""
-
-import os
-
-import matplotlib.pyplot as plt
-import numpy as np
-from sklearn.decomposition import PCA
-from mpl_toolkits.mplot3d import Axes3D
-
-from .Strees import Stree, Snode, Siterator
-
-
-class Snode_graph(Snode):
-    def __init__(self, node: Stree):
-        self._plot_size = (8, 8)
-        self._xlimits = (None, None)
-        self._ylimits = (None, None)
-        self._zlimits = (None, None)
-        n = Snode.copy(node)
-        super().__init__(n._clf, n._X, n._y, n._title)
-
-    def set_plot_size(self, size: tuple):
-        self._plot_size = size
-
-    def get_plot_size(self) -> tuple:
-        return self._plot_size
-
-    def _is_pure(self) -> bool:
-        """is considered pure a leaf node with one label
-        """
-        if self.is_leaf():
-            return self._belief == 1.0
-        return False
-
-    def set_axis_limits(self, limits: tuple):
-        self._xlimits, self._ylimits, self._zlimits = limits
-
-    def get_axis_limits(self) -> tuple:
-        return self._xlimits, self._ylimits, self._zlimits
-
-    def _set_graphics_axis(self, ax: Axes3D):
-        ax.set_xlim(self._xlimits)
-        ax.set_ylim(self._ylimits)
-        ax.set_zlim(self._zlimits)
-
-    def save_hyperplane(
-        self, save_folder: str = "./", save_prefix: str = "", save_seq: int = 1
-    ):
-        _, fig = self.plot_hyperplane()
-        name = os.path.join(save_folder, f"{save_prefix}STnode{save_seq}.png")
-        fig.savefig(name, bbox_inches="tight")
-        plt.close(fig)
-
-    def _get_cmap(self):
-        cmap = "jet"
-        if self._is_pure() and self._class == 1:
-            cmap = "jet_r"
-        return cmap
-
-    def _graph_title(self):
-        n_class, card = np.unique(self._y, return_counts=True)
-        return f"{self._title} {n_class} {card}"
-
-    def plot_hyperplane(self, plot_distribution: bool = True):
-        fig = plt.figure(figsize=self._plot_size)
-        ax = fig.add_subplot(1, 1, 1, projection="3d")
-        if not self._is_pure():
-            # Can't plot hyperplane of leaves with one label because it hasn't
-            # classiffier
-            # get the splitting hyperplane
-            def hyperplane(x, y):
-                return (
-                    -self._clf.intercept_
-                    - self._clf.coef_[0][0] * x
-                    - self._clf.coef_[0][1] * y
-                ) / self._clf.coef_[0][2]
-
-            tmpx = np.linspace(self._X[:, 0].min(), self._X[:, 0].max())
-            tmpy = np.linspace(self._X[:, 1].min(), self._X[:, 1].max())
-            xx, yy = np.meshgrid(tmpx, tmpy)
-            ax.plot_surface(
-                xx,
-                yy,
-                hyperplane(xx, yy),
-                alpha=0.5,
-                antialiased=True,
-                rstride=1,
-                cstride=1,
-                cmap="seismic",
-            )
-            self._set_graphics_axis(ax)
-        if plot_distribution:
-            self.plot_distribution(ax)
-        else:
-            plt.title(self._graph_title())
-            plt.show()
-        return ax, fig
-
-    def plot_distribution(self, ax: Axes3D = None):
-        if ax is None:
-            fig = plt.figure(figsize=self._plot_size)
-            ax = fig.add_subplot(1, 1, 1, projection="3d")
-        plt.title(self._graph_title())
-        cmap = self._get_cmap()
-        ax.scatter(
-            self._X[:, 0], self._X[:, 1], self._X[:, 2], c=self._y, cmap=cmap
-        )
-        ax.set_xlabel("X0")
-        ax.set_ylabel("X1")
-        ax.set_zlabel("X2")
-        plt.show()
-
-
-class Stree_grapher(Stree):
-    """Build 3d graphs of any dataset, if it's more than 3 features PCA shall
-    make its magic
-    """
-
-    def __init__(self, params: dict):
-        self._plot_size = (8, 8)
-        self._tree_gr = None
-        # make Snode store X's
-        os.environ["TESTING"] = "1"
-        self._fitted = False
-        self._pca = None
-        super().__init__(**params)
-
-    def __del__(self):
-        try:
-            os.environ.pop("TESTING")
-        except KeyError:
-            pass
-
-    def _copy_tree(self, node: Snode) -> Snode_graph:
-        mirror = Snode_graph(node)
-        # clone node
-        mirror._class = node._class
-        mirror._belief = node._belief
-        if node.get_down() is not None:
-            mirror.set_down(self._copy_tree(node.get_down()))
-        if node.get_up() is not None:
-            mirror.set_up(self._copy_tree(node.get_up()))
-        return mirror
-
-    def fit(
-        self, X: np.array, y: np.array, sample_weight: np.array = None
-    ) -> "Stree_grapher":
-        """Fit the Stree and copy the tree in a Snode_graph tree
-
-        :param X: Dataset
-        :type X: np.array
-        :param y: Labels
-        :type y: np.array
-        :return: Stree model
-        :rtype: Stree
-        """
-        if X.shape[1] != 3:
-            self._pca = PCA(n_components=3)
-            X = self._pca.fit_transform(X)
-        super().fit(X, y, sample_weight=sample_weight)
-        self._tree_gr = self._copy_tree(self.tree_)
-        self._fitted = True
-        return self
-
-    def score(self, X: np.array, y: np.array) -> float:
-        self._check_fitted()
-        if X.shape[1] != 3:
-            X = self._pca.transform(X)
-        return super().score(X, y)
-
-    def _check_fitted(self):
-        if not self._fitted:
-            raise Exception("Have to fit the grapher first!")
-
-    def save_all(self, save_folder: str = "./", save_prefix: str = ""):
-        """Save all the node plots in png format, each with a sequence number
-
-        :param save_folder: folder where the plots are saved, defaults to './'
-        :type save_folder: str, optional
-        """
-        self._check_fitted()
-        if not os.path.isdir(save_folder):
-            os.mkdir(save_folder)
-        seq = 1
-        for node in self:
-            node.save_hyperplane(
-                save_folder=save_folder, save_prefix=save_prefix, save_seq=seq
-            )
-            seq += 1
-
-    def plot_all(self):
-        """Plots all the nodes
-        """
-        self._check_fitted()
-        for node in self:
-            node.plot_hyperplane()
-
-    def __iter__(self):
-        return Siterator(self._tree_gr)
--- a/stree/init.py
+++ b/stree/init.py
@@ -1,4 +1,3 @@
-from .Strees import Stree, Snode, Siterator
-from .Strees_grapher import Stree_grapher, Snode_graph
+from .Strees import Stree, Snode, Siterator, Splitter

-__all__ = ["Stree", "Snode", "Siterator", "Stree_grapher", "Snode_graph"]
+__all__ = ["Stree", "Snode", "Siterator", "Splitter"]
--- a/stree/tests/Snode_test.py
+++ b/stree/tests/Snode_test.py
@@ -0,0 +1,91 @@
+import os
+import unittest
+
+import numpy as np
+
+from stree import Stree, Snode
+from .utils import load_dataset
+
+
+class Snode_test(unittest.TestCase):
+    def __init__(self, *args, **kwargs):
+        self._random_state = 1
+        self._clf = Stree(random_state=self._random_state)
+        self._clf.fit(*load_dataset(self._random_state))
+        super().__init__(*args, **kwargs)
+
+    @classmethod
+    def setUp(cls):
+        os.environ["TESTING"] = "1"
+
+    def test_attributes_in_leaves(self):
+        """Check if the attributes in leaves have correct values so they form a
+        predictor
+        """
+
+        def check_leave(node: Snode):
+            if not node.is_leaf():
+                check_leave(node.get_down())
+                check_leave(node.get_up())
+                return
+            # Check Belief in leave
+            classes, card = np.unique(node._y, return_counts=True)
+            max_card = max(card)
+            min_card = min(card)
+            if len(classes) > 1:
+                try:
+                    belief = max_card / (max_card + min_card)
+                except ZeroDivisionError:
+                    belief = 0.0
+            else:
+                belief = 1
+            self.assertEqual(belief, node._belief)
+            # Check Class
+            class_computed = classes[card == max_card]
+            self.assertEqual(class_computed, node._class)
+
+        check_leave(self._clf.tree_)
+
+    def test_nodes_coefs(self):
+        """Check if the nodes of the tree have the right attributes filled
+        """
+
+        def run_tree(node: Snode):
+            if node._belief < 1:
+                # only exclude pure leaves
+                self.assertIsNotNone(node._clf)
+                self.assertIsNotNone(node._clf.coef_)
+            if node.is_leaf():
+                return
+            run_tree(node.get_down())
+            run_tree(node.get_up())
+
+        run_tree(self._clf.tree_)
+
+    def test_make_predictor_on_leaf(self):
+        test = Snode(None, [1, 2, 3, 4], [1, 0, 1, 1], [], 0.0, "test")
+        test.make_predictor()
+        self.assertEqual(1, test._class)
+        self.assertEqual(0.75, test._belief)
+
+    def test_make_predictor_on_not_leaf(self):
+        test = Snode(None, [1, 2, 3, 4], [1, 0, 1, 1], [], 0.0, "test")
+        test.set_up(Snode(None, [1], [1], [], 0.0, "another_test"))
+        test.make_predictor()
+        self.assertIsNone(test._class)
+        self.assertEqual(0, test._belief)
+
+    def test_make_predictor_on_leaf_bogus_data(self):
+        test = Snode(None, [1, 2, 3, 4], [], [], 0.0, "test")
+        test.make_predictor()
+        self.assertIsNone(test._class)
+
+    def test_copy_node(self):
+        px = [1, 2, 3, 4]
+        py = [1]
+        test = Snode(Stree(), px, py, [], 0.0, "test")
+        computed = Snode.copy(test)
+        self.assertListEqual(computed._X, px)
+        self.assertListEqual(computed._y, py)
+        self.assertEqual("test", computed._title)
+        self.assertIsInstance(computed._clf, Stree)
--- a/stree/tests/Splitter_test.py
+++ b/stree/tests/Splitter_test.py
@@ -0,0 +1,142 @@
+import os
+import unittest
+
+import numpy as np
+from sklearn.svm import LinearSVC
+
+from stree import Splitter
+from .utils import load_dataset
+
+
+class Splitter_test(unittest.TestCase):
+    def __init__(self, *args, **kwargs):
+        self._random_state = 1
+        super().__init__(*args, **kwargs)
+
+    @staticmethod
+    def build(
+        clf=LinearSVC(),
+        min_samples_split=0,
+        splitter_type="random",
+        criterion="gini",
+        criteria="min_distance",
+        random_state=None,
+    ):
+        return Splitter(
+            clf=clf,
+            min_samples_split=min_samples_split,
+            splitter_type=splitter_type,
+            criterion=criterion,
+            criteria=criteria,
+            random_state=random_state,
+        )
+
+    @classmethod
+    def setUp(cls):
+        os.environ["TESTING"] = "1"
+
+    def test_init(self):
+        with self.assertRaises(ValueError):
+            self.build(criterion="duck")
+        with self.assertRaises(ValueError):
+            self.build(splitter_type="duck")
+        with self.assertRaises(ValueError):
+            self.build(criteria="duck")
+        with self.assertRaises(ValueError):
+            self.build(clf=None)
+        for splitter_type in ["best", "random"]:
+            for criterion in ["gini", "entropy"]:
+                for criteria in ["min_distance", "max_samples"]:
+                    tcl = self.build(
+                        splitter_type=splitter_type,
+                        criterion=criterion,
+                        criteria=criteria,
+                    )
+                    self.assertEqual(splitter_type, tcl._splitter_type)
+                    self.assertEqual(criterion, tcl._criterion)
+                    self.assertEqual(criteria, tcl._criteria)
+
+    def test_gini(self):
+        y = [0, 1, 1, 1, 1, 1, 0, 0, 0, 1]
+        expected = 0.48
+        self.assertEqual(expected, Splitter._gini(y))
+        tcl = self.build(criterion="gini")
+        self.assertEqual(expected, tcl.criterion_function(y))
+
+    def test_entropy(self):
+        y = [0, 1, 1, 1, 1, 1, 0, 0, 0, 1]
+        expected = 0.9709505944546686
+        self.assertAlmostEqual(expected, Splitter._entropy(y))
+        tcl = self.build(criterion="entropy")
+        self.assertEqual(expected, tcl.criterion_function(y))
+
+    def test_information_gain(self):
+        yu = np.array([0, 1, 1, 1, 1, 1])
+        yd = np.array([0, 0, 0, 1])
+        values_expected = [
+            ("gini", 0.31666666666666665),
+            ("entropy", 0.7145247027726656),
+        ]
+        for criterion, expected in values_expected:
+            tcl = self.build(criterion=criterion)
+            computed = tcl.information_gain(yu, yd)
+            self.assertAlmostEqual(expected, computed)
+
+    def test_max_samples(self):
+        tcl = self.build(criteria="max_samples")
+        data = np.array(
+            [
+                [-0.1, 0.2, -0.3],
+                [0.7, 0.01, -0.1],
+                [0.7, -0.9, 0.5],
+                [0.1, 0.2, 0.3],
+            ]
+        )
+        expected = np.array([0.2, 0.01, -0.9, 0.2])
+        y = [1, 2, 1, 0]
+        computed = tcl._max_samples(data, y)
+        self.assertEqual((4,), computed.shape)
+        self.assertListEqual(expected.tolist(), computed.tolist())
+
+    def test_min_distance(self):
+        tcl = self.build()
+        data = np.array(
+            [
+                [-0.1, 0.2, -0.3],
+                [0.7, 0.01, -0.1],
+                [0.7, -0.9, 0.5],
+                [0.1, 0.2, 0.3],
+            ]
+        )
+        expected = np.array([-0.1, 0.01, 0.5, 0.1])
+        computed = tcl._min_distance(data, None)
+        self.assertEqual((4,), computed.shape)
+        self.assertListEqual(expected.tolist(), computed.tolist())
+
+    def test_splitter_parameter(self):
+        expected_values = [
+            [1, 7, 9],
+            [1, 7, 9],
+            [1, 7, 9],
+            [1, 7, 9],
+            [0, 5, 6],
+            [0, 5, 6],
+            [0, 5, 6],
+            [0, 5, 6],
+        ]
+        X, y = load_dataset(self._random_state, n_features=12)
+        for splitter_type in ["best", "random"]:
+            for criterion in ["gini", "entropy"]:
+                for criteria in ["min_distance", "max_samples"]:
+                    tcl = self.build(
+                        splitter_type=splitter_type,
+                        criterion=criterion,
+                        criteria=criteria,
+                        random_state=self._random_state,
+                    )
+                    expected = expected_values.pop(0)
+                    dataset, computed = tcl.get_subspace(X, y, max_features=3)
+                    self.assertListEqual(expected, list(computed))
+                    self.assertListEqual(
+                        X[:, computed].tolist(), dataset.tolist()
+                    )
--- a/stree/tests/Strees_test.py
+++ b/stree/tests/Strees_test.py
@@ -2,26 +2,10 @@ import os
 import unittest

 import numpy as np
-from sklearn.datasets import make_classification
+from sklearn.datasets import load_iris

 from stree import Stree, Snode
-
-
-def get_dataset(random_state=0):
-    X, y = make_classification(
-        n_samples=1500,
-        n_features=3,
-        n_informative=3,
-        n_redundant=0,
-        n_repeated=0,
-        n_classes=2,
-        n_clusters_per_class=2,
-        class_sep=1.5,
-        flip_y=0,
-        weights=[0.5, 0.5],
-        random_state=random_state,
-    )
-    return X, y
+from .utils import load_dataset


 class Stree_test(unittest.TestCase):
@@ -80,12 +64,11 @@ class Stree_test(unittest.TestCase):
        warnings.filterwarnings("ignore")
        for kernel in self._kernels:
            clf = Stree(kernel=kernel, random_state=self._random_state)
-            clf.fit(*get_dataset(self._random_state))
+            clf.fit(*load_dataset(self._random_state))
            self._check_tree(clf.tree_)

-    def _find_out(
-        self, px: np.array, x_original: np.array, y_original
-    ) -> list:
+    @staticmethod
+    def _find_out(px: np.array, x_original: np.array, y_original) -> list:
        """Find the original values of y for a given array of samples

        Arguments:
@@ -104,9 +87,8 @@ class Stree_test(unittest.TestCase):
        return res

    def test_single_prediction(self):
-        probs = [0.29026400766, 0.73105613, 0.0307635]
-        X, y = get_dataset(self._random_state)
-        for kernel, prob in zip(self._kernels, probs):
+        X, y = load_dataset(self._random_state)
+        for kernel in self._kernels:
            clf = Stree(kernel=kernel, random_state=self._random_state)
            yp = clf.fit(X, y).predict((X[0, :].reshape(-1, X.shape[1])))
            self.assertEqual(yp[0], y[0])
@@ -114,18 +96,20 @@ class Stree_test(unittest.TestCase):
    def test_multiple_prediction(self):
        # First 27 elements the predictions are the same as the truth
        num = 27
-        X, y = get_dataset(self._random_state)
+        X, y = load_dataset(self._random_state)
        for kernel in self._kernels:
            clf = Stree(kernel=kernel, random_state=self._random_state)
            yp = clf.fit(X, y).predict(X[:num, :])
            self.assertListEqual(y[:num].tolist(), yp.tolist())

    def test_score(self):
-        X, y = get_dataset(self._random_state)
-        for kernel, accuracy_expected in zip(
-            self._kernels,
-            [0.9506666666666667, 0.9606666666666667, 0.9433333333333334],
-        ):
+        X, y = load_dataset(self._random_state)
+        accuracies = [
+            0.9506666666666667,
+            0.9606666666666667,
+            0.9433333333333334,
+        ]
+        for kernel, accuracy_expected in zip(self._kernels, accuracies):
            clf = Stree(random_state=self._random_state, kernel=kernel,)
            clf.fit(X, y)
            accuracy_score = clf.score(X, y)
@@ -134,43 +118,11 @@ class Stree_test(unittest.TestCase):
            self.assertEqual(accuracy_score, accuracy_computed)
            self.assertAlmostEqual(accuracy_expected, accuracy_score)

-    def test_single_predict_proba(self):
-        """Check the element 28 probability of being 1
-        """
-        decimals = 5
-        element = 28
-        probs = [0.29026400766, 0.73105613, 0.0307635]
-        X, y = get_dataset(self._random_state)
-        self.assertEqual(1, y[element])
-        for kernel, prob in zip(self._kernels, probs):
-            clf = Stree(kernel=kernel, random_state=self._random_state)
-            yp = clf.fit(X, y).predict_proba(
-                X[element, :].reshape(-1, X.shape[1])
-            )
-            self.assertAlmostEqual(
-                np.round(1 - prob, decimals), np.round(yp[0:, 0], decimals)
-            )
-            self.assertAlmostEqual(
-                round(prob, decimals), round(yp[0, 1], decimals), decimals
-            )
-
-    def test_multiple_predict_proba(self):
-        # First 27 elements the predictions are the same as the truth
-        num = 27
-        X, y = get_dataset(self._random_state)
-        for kernel in self._kernels:
-            clf = Stree(kernel=kernel, random_state=self._random_state)
-            clf.fit(X, y)
-            yp = clf.predict_proba(X[:num, :])
-            self.assertListEqual(
-                y[:num].tolist(), np.argmax(yp[:num], axis=1).tolist()
-            )
-
    def test_single_vs_multiple_prediction(self):
        """Check if predicting sample by sample gives the same result as
        predicting all samples at once
        """
-        X, y = get_dataset(self._random_state)
+        X, y = load_dataset(self._random_state)
        for kernel in self._kernels:
            clf = Stree(kernel=kernel, random_state=self._random_state)
            clf.fit(X, y)
@@ -188,29 +140,30 @@ class Stree_test(unittest.TestCase):
        """Check preorder iterator
        """
        expected = [
-            "root",
-            "root - Down",
-            "root - Down - Down, <cgaf> - Leaf class=1 belief= 0.975989 counts"
-            "=(array([0, 1]), array([ 17, 691]))",
-            "root - Down - Up",
+            "root feaures=(0, 1, 2) impurity=0.5000",
+            "root - Down feaures=(0, 1, 2) impurity=0.0671",
+            "root - Down - Down, <cgaf> - Leaf class=1 belief= 0.975989 "
+            "impurity=0.0469 counts=(array([0, 1]), array([ 17, 691]))",
+            "root - Down - Up feaures=(0, 1, 2) impurity=0.3967",
            "root - Down - Up - Down, <cgaf> - Leaf class=1 belief= 0.750000 "
-            "counts=(array([0, 1]), array([1, 3]))",
+            "impurity=0.3750 counts=(array([0, 1]), array([1, 3]))",
            "root - Down - Up - Up, <pure> - Leaf class=0 belief= 1.000000 "
-            "counts=(array([0]), array([7]))",
-            "root - Up, <cgaf> - Leaf class=0 belief= 0.928297 counts=(array("
-            "[0, 1]), array([725,  56]))",
+            "impurity=0.0000 counts=(array([0]), array([7]))",
+            "root - Up, <cgaf> - Leaf class=0 belief= 0.928297 impurity=0.1331"
+            " counts=(array([0, 1]), array([725,  56]))",
        ]
        computed = []
        expected_string = ""
        clf = Stree(kernel="linear", random_state=self._random_state)
-        clf.fit(*get_dataset(self._random_state))
+        clf.fit(*load_dataset(self._random_state))
        for node in clf:
            computed.append(str(node))
            expected_string += str(node) + "\n"
        self.assertListEqual(expected, computed)
        self.assertEqual(expected_string, str(clf))

-    def test_is_a_sklearn_classifier(self):
+    @staticmethod
+    def test_is_a_sklearn_classifier():
        import warnings
        from sklearn.exceptions import ConvergenceWarning

@@ -223,7 +176,12 @@ class Stree_test(unittest.TestCase):
    def test_exception_if_C_is_negative(self):
        tclf = Stree(C=-1)
        with self.assertRaises(ValueError):
-            tclf.fit(*get_dataset(self._random_state))
+            tclf.fit(*load_dataset(self._random_state))
+
+    def test_exception_if_bogus_split_criteria(self):
+        tclf = Stree(split_criteria="duck")
+        with self.assertRaises(ValueError):
+            tclf.fit(*load_dataset(self._random_state))

    def test_check_max_depth_is_positive_or_None(self):
        tcl = Stree()
@@ -232,13 +190,13 @@ class Stree_test(unittest.TestCase):
        self.assertGreaterEqual(1, tcl.max_depth)
        with self.assertRaises(ValueError):
            tcl = Stree(max_depth=-1)
-            tcl.fit(*get_dataset(self._random_state))
+            tcl.fit(*load_dataset(self._random_state))

    def test_check_max_depth(self):
        depths = (3, 4)
        for depth in depths:
            tcl = Stree(random_state=self._random_state, max_depth=depth)
-            tcl.fit(*get_dataset(self._random_state))
+            tcl.fit(*load_dataset(self._random_state))
            self.assertEqual(depth, tcl.depth_)

    def test_unfitted_tree_is_iterable(self):
@@ -246,95 +204,131 @@ class Stree_test(unittest.TestCase):
        self.assertEqual(0, len(list(tcl)))

    def test_min_samples_split(self):
-        tcl_split = Stree(min_samples_split=3)
-        tcl_nosplit = Stree(min_samples_split=4)
        dataset = [[1], [2], [3]], [1, 1, 0]
-        tcl_split.fit(*dataset)
+        tcl_split = Stree(min_samples_split=3).fit(*dataset)
        self.assertIsNotNone(tcl_split.tree_.get_down())
        self.assertIsNotNone(tcl_split.tree_.get_up())
-        tcl_nosplit.fit(*dataset)
+        tcl_nosplit = Stree(min_samples_split=4).fit(*dataset)
        self.assertIsNone(tcl_nosplit.tree_.get_down())
        self.assertIsNone(tcl_nosplit.tree_.get_up())

-    def test_muticlass_dataset(self):
+    def test_simple_muticlass_dataset(self):
        for kernel in self._kernels:
-            clf = Stree(kernel=kernel, random_state=self._random_state)
-            px = [[1, 2], [3, 4], [5, 6]]
-            py = [1, 2, 3]
+            clf = Stree(
+                kernel=kernel,
+                split_criteria="max_samples",
+                random_state=self._random_state,
+            )
+            px = [[1, 2], [5, 6], [9, 10]]
+            py = [0, 1, 2]
            clf.fit(px, py)
            self.assertEqual(1.0, clf.score(px, py))
-            self.assertListEqual([1, 2, 3], clf.predict(px).tolist())
+            self.assertListEqual(py, clf.predict(px).tolist())
+            self.assertListEqual(py, clf.classes_.tolist())

+    def test_muticlass_dataset(self):
+        datasets = {
+            "Synt": load_dataset(random_state=self._random_state, n_classes=3),
+            "Iris": load_iris(return_X_y=True),
+        }
+        outcomes = {
+            "Synt": {
+                "max_samples linear": 0.9533333333333334,
+                "max_samples rbf": 0.836,
+                "max_samples poly": 0.9473333333333334,
+                "min_distance linear": 0.9533333333333334,
+                "min_distance rbf": 0.836,
+                "min_distance poly": 0.9473333333333334,
+            },
+            "Iris": {
+                "max_samples linear": 0.98,
+                "max_samples rbf": 1.0,
+                "max_samples poly": 1.0,
+                "min_distance linear": 0.98,
+                "min_distance rbf": 1.0,
+                "min_distance poly": 1.0,
+            },
+        }
+        for name, dataset in datasets.items():
+            px, py = dataset
+            for criteria in ["max_samples", "min_distance"]:
+                for kernel in self._kernels:
+                    clf = Stree(
+                        C=1e4,
+                        max_iter=1e4,
+                        kernel=kernel,
+                        random_state=self._random_state,
+                    )
+                    clf.fit(px, py)
+                    outcome = outcomes[name][f"{criteria} {kernel}"]
+                    self.assertAlmostEqual(outcome, clf.score(px, py))

-class Snode_test(unittest.TestCase):
-    def __init__(self, *args, **kwargs):
-        self._random_state = 1
-        self._clf = Stree(random_state=self._random_state)
-        self._clf.fit(*get_dataset(self._random_state))
-        super().__init__(*args, **kwargs)
+    def test_max_features(self):
+        n_features = 16
+        expected_values = [
+            ("auto", 4),
+            ("log2", 4),
+            ("sqrt", 4),
+            (0.5, 8),
+            (3, 3),
+            (None, 16),
+        ]
+        clf = Stree()
+        clf.n_features_ = n_features
+        for max_features, expected in expected_values:
+            clf.set_params(**dict(max_features=max_features))
+            computed = clf._initialize_max_features()
+            self.assertEqual(expected, computed)
+        # Check bogus max_features
+        values = ["duck", -0.1, 0.0]
+        for max_features in values:
+            clf.set_params(**dict(max_features=max_features))
+            with self.assertRaises(ValueError):
+                _ = clf._initialize_max_features()

-    @classmethod
-    def setUp(cls):
-        os.environ["TESTING"] = "1"
+    def test_get_subspaces(self):
+        dataset = np.random.random((10, 16))
+        y = np.random.randint(0, 2, 10)
+        expected_values = [
+            ("auto", 4),
+            ("log2", 4),
+            ("sqrt", 4),
+            (0.5, 8),
+            (3, 3),
+            (None, 16),
+        ]
+        clf = Stree()
+        for max_features, expected in expected_values:
+            clf.set_params(**dict(max_features=max_features))
+            clf.fit(dataset, y)
+            computed, indices = clf.splitter_.get_subspace(
+                dataset, y, clf.max_features_
+            )
+            self.assertListEqual(
+                dataset[:, indices].tolist(), computed.tolist()
+            )
+            self.assertEqual(expected, len(indices))

-    def test_attributes_in_leaves(self):
-        """Check if the attributes in leaves have correct values so they form a
-        predictor
-        """
+    def test_bogus_criterion(self):
+        clf = Stree(criterion="duck")
+        with self.assertRaises(ValueError):
+            clf.fit(*load_dataset())

-        def check_leave(node: Snode):
-            if not node.is_leaf():
-                check_leave(node.get_down())
-                check_leave(node.get_up())
-                return
-            # Check Belief in leave
-            classes, card = np.unique(node._y, return_counts=True)
-            max_card = max(card)
-            min_card = min(card)
-            if len(classes) > 1:
-                try:
-                    belief = max_card / (max_card + min_card)
-                except ZeroDivisionError:
-                    belief = 0.0
-            else:
-                belief = 1
-            self.assertEqual(belief, node._belief)
-            # Check Class
-            class_computed = classes[card == max_card]
-            self.assertEqual(class_computed, node._class)
+    def test_predict_feature_dimensions(self):
+        X = np.random.rand(10, 5)
+        y = np.random.randint(0, 2, 10)
+        clf = Stree()
+        clf.fit(X, y)
+        with self.assertRaises(ValueError):
+            clf.predict(X[:, :3])

-        check_leave(self._clf.tree_)
+    def test_score_max_features(self):
+        X, y = load_dataset(self._random_state)
+        clf = Stree(random_state=self._random_state, max_features=2)
+        clf.fit(X, y)
+        self.assertAlmostEqual(0.9426666666666667, clf.score(X, y))

-    def test_nodes_coefs(self):
-        """Check if the nodes of the tree have the right attributes filled
-        """
-
-        def run_tree(node: Snode):
-            if node._belief < 1:
-                # only exclude pure leaves
-                self.assertIsNotNone(node._clf)
-                self.assertIsNotNone(node._clf.coef_)
-            if node.is_leaf():
-                return
-            run_tree(node.get_down())
-            run_tree(node.get_up())
-
-        run_tree(self._clf.tree_)
-
-    def test_make_predictor_on_leaf(self):
-        test = Snode(None, [1, 2, 3, 4], [1, 0, 1, 1], "test")
-        test.make_predictor()
-        self.assertEqual(1, test._class)
-        self.assertEqual(0.75, test._belief)
-
-    def test_make_predictor_on_not_leaf(self):
-        test = Snode(None, [1, 2, 3, 4], [1, 0, 1, 1], "test")
-        test.set_up(Snode(None, [1], [1], "another_test"))
-        test.make_predictor()
-        self.assertIsNone(test._class)
-        self.assertEqual(0, test._belief)
-
-    def test_make_predictor_on_leaf_bogus_data(self):
-        test = Snode(None, [1, 2, 3, 4], [], "test")
-        test.make_predictor()
-        self.assertIsNone(test._class)
+    def test_bogus_splitter_parameter(self):
+        clf = Stree(splitter="duck")
+        with self.assertRaises(ValueError):
+            clf.fit(*load_dataset())
--- a/stree/tests/Strees_grapher_test.py
+++ b/stree/tests/Strees_grapher_test.py
@@ -1,211 +0,0 @@
-import os
-import imghdr
-import unittest
-
-import numpy as np
-import matplotlib
-import matplotlib.pyplot as plt
-import warnings
-from sklearn.datasets import make_classification
-
-from stree import Stree_grapher, Snode_graph, Snode
-
-
-def get_dataset(random_state=0, n_features=3):
-    X, y = make_classification(
-        n_samples=1500,
-        n_features=n_features,
-        n_informative=3,
-        n_redundant=0,
-        n_repeated=0,
-        n_classes=2,
-        n_clusters_per_class=2,
-        class_sep=1.5,
-        flip_y=0,
-        weights=[0.5, 0.5],
-        random_state=random_state,
-    )
-    return X, y
-
-
-class Stree_grapher_test(unittest.TestCase):
-    def __init__(self, *args, **kwargs):
-        self._random_state = 1
-        self._clf = Stree_grapher(dict(random_state=self._random_state))
-        self._clf.fit(*get_dataset(self._random_state, n_features=4))
-        super().__init__(*args, **kwargs)
-
-    @classmethod
-    def setUp(cls):
-        os.environ["TESTING"] = "1"
-
-    def test_iterator(self):
-        """Check preorder iterator
-        """
-        expected = [
-            "root",
-            "root - Down",
-            "root - Down - Down, <cgaf> - Leaf class=1 belief= 0.976023 counts"
-            "=(array([0, 1]), array([ 17, 692]))",
-            "root - Down - Up",
-            "root - Down - Up - Down, <cgaf> - Leaf class=0 belief= 0.500000 "
-            "counts=(array([0, 1]), array([1, 1]))",
-            "root - Down - Up - Up, <cgaf> - Leaf class=0 belief= 0.888889 "
-            "counts=(array([0, 1]), array([8, 1]))",
-            "root - Up, <cgaf> - Leaf class=0 belief= 0.928205 counts=(array("
-            "[0, 1]), array([724,  56]))",
-        ]
-        computed = []
-        for node in self._clf:
-            computed.append(str(node))
-        self.assertListEqual(expected, computed)
-
-    def test_score(self):
-        X, y = get_dataset(self._random_state)
-        accuracy_score = self._clf.score(X, y)
-        yp = self._clf.predict(X)
-        accuracy_computed = np.mean(yp == y)
-        self.assertEqual(accuracy_score, accuracy_computed)
-        self.assertGreater(accuracy_score, 0.86)
-
-    def test_save_all(self):
-        folder_name = os.path.join(os.sep, "tmp", "stree")
-        if os.path.isdir(folder_name):
-            os.rmdir(folder_name)
-        file_names = [
-            os.path.join(folder_name, f"STnode{i}.png") for i in range(1, 8)
-        ]
-        with warnings.catch_warnings():
-            warnings.simplefilter("ignore")
-            matplotlib.use("Agg")
-            self._clf.save_all(save_folder=folder_name)
-        for file_name in file_names:
-            self.assertTrue(os.path.exists(file_name))
-            self.assertEqual("png", imghdr.what(file_name))
-            os.remove(file_name)
-        os.rmdir(folder_name)
-
-    def test_plot_all(self):
-        with warnings.catch_warnings():
-            warnings.simplefilter("ignore")
-            matplotlib.use("Agg")
-            num_figures_before = plt.gcf().number
-            self._clf.plot_all()
-            num_figures_after = plt.gcf().number
-        self.assertEqual(7, num_figures_after - num_figures_before)
-
-
-class Snode_graph_test(unittest.TestCase):
-    def __init__(self, *args, **kwargs):
-        self._random_state = 1
-        self._clf = Stree_grapher(dict(random_state=self._random_state))
-        self._clf.fit(*get_dataset(self._random_state))
-        super().__init__(*args, **kwargs)
-
-    @classmethod
-    def setUp(cls):
-        os.environ["TESTING"] = "1"
-
-    def test_plot_size(self):
-        default = self._clf._tree_gr.get_plot_size()
-        expected = (17, 3)
-        self._clf._tree_gr.set_plot_size(expected)
-        self.assertEqual(expected, self._clf._tree_gr.get_plot_size())
-        self._clf._tree_gr.set_plot_size(default)
-        self.assertEqual(default, self._clf._tree_gr.get_plot_size())
-
-    def test_attributes_in_leaves_graph(self):
-        """Check if the attributes in leaves have correct values so they form a
-        predictor
-        """
-
-        def check_leave(node: Snode_graph):
-            if not node.is_leaf():
-                check_leave(node.get_down())
-                check_leave(node.get_up())
-                return
-            # Check Belief in leave
-            classes, card = np.unique(node._y, return_counts=True)
-            max_card = max(card)
-            min_card = min(card)
-            if len(classes) > 1:
-                try:
-                    belief = max_card / (max_card + min_card)
-                except ZeroDivisionError:
-                    belief = 0.0
-            else:
-                belief = 1
-            self.assertEqual(belief, node._belief)
-            # Check Class
-            class_computed = classes[card == max_card]
-            self.assertEqual(class_computed, node._class)
-
-        check_leave(self._clf._tree_gr)
-
-    def test_nodes_graph_coefs(self):
-        """Check if the nodes of the tree have the right attributes filled
-        """
-
-        def run_tree(node: Snode_graph):
-            if node._belief < 1:
-                # only exclude pure leaves
-                self.assertIsNotNone(node._clf)
-                self.assertIsNotNone(node._clf.coef_)
-            if node.is_leaf():
-                return
-            run_tree(node.get_down())
-            run_tree(node.get_up())
-
-        run_tree(self._clf._tree_gr)
-
-    def test_save_hyperplane(self):
-        folder_name = "/tmp/"
-        file_name = os.path.join(folder_name, "STnode1.png")
-        with warnings.catch_warnings():
-            warnings.simplefilter("ignore")
-            matplotlib.use("Agg")
-            self._clf._tree_gr.save_hyperplane(folder_name)
-        self.assertTrue(os.path.exists(file_name))
-        self.assertEqual("png", imghdr.what(file_name))
-        os.remove(file_name)
-
-    def test_plot_hyperplane_with_distribution(self):
-        plt.close()
-        with warnings.catch_warnings():
-            warnings.simplefilter("ignore")
-            matplotlib.use("Agg")
-            num_figures_before = plt.gcf().number
-            self._clf._tree_gr.plot_hyperplane(plot_distribution=True)
-            num_figures_after = plt.gcf().number
-        self.assertEqual(1, num_figures_after - num_figures_before)
-
-    def test_plot_hyperplane_without_distribution(self):
-        plt.close()
-        with warnings.catch_warnings():
-            warnings.simplefilter("ignore")
-            matplotlib.use("Agg")
-            num_figures_before = plt.gcf().number
-            self._clf._tree_gr.plot_hyperplane(plot_distribution=False)
-            num_figures_after = plt.gcf().number
-        self.assertEqual(1, num_figures_after - num_figures_before)
-
-    def test_plot_distribution(self):
-        plt.close()
-        with warnings.catch_warnings():
-            warnings.simplefilter("ignore")
-            matplotlib.use("Agg")
-            num_figures_before = plt.gcf().number
-            self._clf._tree_gr.plot_distribution()
-            num_figures_after = plt.gcf().number
-        self.assertEqual(1, num_figures_after - num_figures_before)
-
-    def test_set_axis_limits(self):
-        node = Snode_graph(Snode(None, None, None, "test"))
-        limits = (-2, 2), (-3, 3), (-4, 4)
-        node.set_axis_limits(limits)
-        computed = node.get_axis_limits()
-        x, y, z = limits
-        xx, yy, zz = computed
-        self.assertEqual(x, xx)
-        self.assertEqual(y, yy)
-        self.assertEqual(z, zz)
--- a/stree/tests/init.py
+++ b/stree/tests/init.py
@@ -1,9 +1,5 @@
-from .Strees_test import Stree_test, Snode_test
-from .Strees_grapher_test import Stree_grapher_test, Snode_graph_test
+from .Stree_test import Stree_test
+from .Snode_test import Snode_test
+from .Splitter_test import Splitter_test

-__all__ = [
-    "Stree_test",
-    "Snode_test",
-    "Stree_grapher_test",
-    "Snode_graph_test",
-]
+__all__ = ["Stree_test", "Snode_test", "Splitter_test"]
--- a/stree/tests/utils.py
+++ b/stree/tests/utils.py
@@ -0,0 +1,17 @@
+from sklearn.datasets import make_classification
+
+
+def load_dataset(random_state=0, n_classes=2, n_features=3):
+    X, y = make_classification(
+        n_samples=1500,
+        n_features=n_features,
+        n_informative=3,
+        n_redundant=0,
+        n_repeated=0,
+        n_classes=n_classes,
+        n_clusters_per_class=2,
+        class_sep=1.5,
+        flip_y=0,
+        random_state=random_state,
+    )
+    return X, y
Author	SHA1	Message	Date
Ricardo Montañana	9334951d1b	#2 Cosmetic and style updates	2020-06-15 11:09:11 +02:00
Ricardo Montañana	736ab7ef20	#2 update benchmark notebook	2020-06-15 10:33:51 +02:00
Ricardo Montañana	c94bc068bd	#2 Refactor Stree & create Splitter Add and test splitter parameter	2020-06-15 00:22:57 +02:00
Ricardo Montañana	502ee72799	#2 Add predict and score support Add a test in features notebook Show max_features in main.py	2020-06-14 14:00:21 +02:00
Ricardo Montañana	f1ee4de37b	#2 - Add gini and entropy measures rename get_dataset to load_dataset add features and impurity to __str__ of node	2020-06-14 03:08:55 +02:00
Ricardo Montañana	ae1c199e21	# 2 - add max_features parameters	2020-06-13 17:58:45 +02:00
Ricardo Montañana	1bfe273a70	Fix problem in _min_distance Remove grapher (moved to another repo)	2020-06-12 00:50:25 +02:00
Ricardo Montañana Gómez	647d21bdb5	Merge pull request #9 from Doctorado-ML/add_multiclass #6 Add multiclass	2020-06-11 16:30:16 +02:00
Ricardo Montañana	1d392d534f	#6 - Update tests and codecov conf	2020-06-11 13:45:24 +02:00
Ricardo Montañana	f360a2640c	#6 - Add multiclass support Removed (by now) predict_proba. Created a notebook in jupyter Added split_criteria parameter with min_distance and max_samples values Refactor _distances Refactor _split_criteria Refactor _reorder_results	2020-06-11 13:10:52 +02:00
Ricardo Montañana Gómez	45510b43bc	Merge pull request #5 from Doctorado-ML/add_kernels #3 Add kernels to STree	2020-06-09 13:43:31 +02:00