Merge pull request #38 from arminwitte/print_rules

implementation to print rules; resolves #31

binarybeech/visualize.py

@@ -74,3 +74,30 @@ def _show(node, tree_view, parent=None, prefix=""):
     tree_view = treelib.Tree()
     _show(tree.root, tree_view)
     tree_view.show()
+
+def extract_rules(tree):
+    leafs = tree.leafs()
+    d = {}
+    for L in leafs:
+        rules = []
+        node = L
+        value = node.value
+        if value not in d:
+            d[value] = []
+        while node is not None:
+            rules.append((node.attribute,node.threshold))
+            node = node.parent
+        d[value].append(reversed(rules))
+    return d
+
+def print_rules(d: dict):
+    s = ""
+    for key, val in d.items():
+        s += key + "\n"
+        for i, rules in enumerate(val):
+            s += "    "
+            if i > 0:
+                s += "or "
+            s += str(rules) + "\n"
+    return s
+

tests/test_visualize.py

@@ -2,7 +2,7 @@
 import pandas as pd
 
 from binarybeech.binarybeech import CART
-from binarybeech.visualize import plot_areas, plot_pruning_quality
+from binarybeech.visualize import plot_areas, plot_pruning_quality, extract_rules, print_rules
 
 
 def test_plot_areas():
@@ -17,3 +17,16 @@ def test_plot_pruning_quality():
     c = CART(df=df_iris, y_name="species", method="classification", seed=42)
     c.train()
     assert isinstance(c.pruning_quality, dict)
+
+def test_print_rules():
+    df_iris = pd.read_csv("data/iris.csv")
+    c = CART(df=df_iris, y_name="species", method="classification", seed=42)
+    c.train()
+    rules = extract_rules(c.tree)
+    assert isinstance(rules,dict)
+
+    s = print_rules(rules)
+    print(s)
+
+    assert isinstance(s,str)
+    assert s.split()[0] == "setosa"